From 648609dbea31d5ba604d00a7f765a1a47d673896 Mon Sep 17 00:00:00 2001
From: Cecylia Bocovich <cohosh@torproject.org>
Date: Tue, 31 Oct 2023 13:15:52 -0400
Subject: [PATCH] Refactor disabling the stats logger

Have Snowflake proxy periodically collect throughput stats even if the
stats logger is disabled so that it can be handled by the prometheus
metrics.
---
 proxy/lib/pt_event_logger.go | 16 ++++++++++++----
 proxy/lib/snowflake.go       | 10 +++-------
 proxy/main.go                |  5 ++---
 3 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/proxy/lib/pt_event_logger.go b/proxy/lib/pt_event_logger.go
index 4ce4b68..810259a 100644
--- a/proxy/lib/pt_event_logger.go
+++ b/proxy/lib/pt_event_logger.go
@@ -9,17 +9,25 @@ import (
 	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/task"
 )
 
-func NewProxyEventLogger(output io.Writer) event.SnowflakeEventReceiver {
+func NewProxyEventLogger(output io.Writer, disableStats bool) event.SnowflakeEventReceiver {
 	logger := log.New(output, "", log.LstdFlags|log.LUTC)
-	return &proxyEventLogger{logger: logger}
+	return &proxyEventLogger{logger: logger, disableStats: disableStats}
 }
 
 type proxyEventLogger struct {
-	logger *log.Logger
+	logger       *log.Logger
+	disableStats bool
 }
 
 func (p *proxyEventLogger) OnNewSnowflakeEvent(e event.SnowflakeEvent) {
-	p.logger.Println(e.String())
+	switch e.(type) {
+	case event.EventOnProxyStats:
+		if !p.disableStats {
+			p.logger.Println(e.String())
+		}
+	default:
+		p.logger.Println(e.String())
+	}
 }
 
 type periodicProxyStats struct {
diff --git a/proxy/lib/snowflake.go b/proxy/lib/snowflake.go
index 0c85754..82a1780 100644
--- a/proxy/lib/snowflake.go
+++ b/proxy/lib/snowflake.go
@@ -140,8 +140,6 @@ type SnowflakeProxy struct {
 	EventDispatcher event.SnowflakeEventDispatcher
 	shutdown        chan struct{}
 
-	// DisableStatsLogger indicates whether proxy stats will be logged periodically
-	DisableStatsLogger bool
 	// SummaryInterval is the time interval at which proxy stats will be logged
 	SummaryInterval time.Duration
 
@@ -658,11 +656,9 @@ func (sf *SnowflakeProxy) Start() error {
 		sf.EventDispatcher = event.NewSnowflakeEventDispatcher()
 	}
 
-	if !sf.DisableStatsLogger {
-		sf.bytesLogger = newBytesSyncLogger()
-		sf.periodicProxyStats = newPeriodicProxyStats(sf.SummaryInterval, sf.EventDispatcher, sf.bytesLogger)
-		sf.EventDispatcher.AddSnowflakeEventListener(sf.periodicProxyStats)
-	}
+	sf.bytesLogger = newBytesSyncLogger()
+	sf.periodicProxyStats = newPeriodicProxyStats(sf.SummaryInterval, sf.EventDispatcher, sf.bytesLogger)
+	sf.EventDispatcher.AddSnowflakeEventListener(sf.periodicProxyStats)
 
 	broker, err = newSignalingServer(sf.BrokerURL, sf.KeepLocalAddresses)
 	if err != nil {
diff --git a/proxy/main.go b/proxy/main.go
index 3713515..8653b80 100644
--- a/proxy/main.go
+++ b/proxy/main.go
@@ -97,8 +97,7 @@ func main() {
 		RelayDomainNamePattern: *allowedRelayHostNamePattern,
 		AllowNonTLSRelay:       *allowNonTLSRelay,
 
-		DisableStatsLogger: *disableStatsLogger,
-		SummaryInterval:    *summaryInterval,
+		SummaryInterval: *summaryInterval,
 	}
 
 	var logOutput = ioutil.Discard
@@ -127,7 +126,7 @@ func main() {
 		log.SetOutput(&safelog.LogScrubber{Output: logOutput})
 	}
 
-	proxyEventLogger := sf.NewProxyEventLogger(eventlogOutput)
+	proxyEventLogger := sf.NewProxyEventLogger(eventlogOutput, *disableStatsLogger)
 	eventLogger.AddSnowflakeEventListener(proxyEventLogger)
 
 	if *enableMetrics {