Add distinct IP counter to metrics

broker/metrics.go

@@ -14,6 +14,7 @@ import (
 	"sync"
 	"time"
 
+	"git.torproject.org/pluggable-transports/snowflake.git/v2/common/ipsetsink/sinkcluster"
 	"git.torproject.org/pluggable-transports/snowflake.git/v2/common/messages"
 	"github.com/prometheus/client_golang/prometheus"
 	"gitlab.torproject.org/tpo/anti-censorship/geoip"
@@ -41,6 +42,8 @@ type Metrics struct {
 	logger  *log.Logger
 	geoipdb *geoip.Geoip
 
+	distinctIPWriter *sinkcluster.ClusterWriter
+
 	countryStats                  CountryStats
 	clientRoundtripEstimate       time.Duration
 	proxyIdleCount                uint
@@ -324,3 +327,13 @@ func initPrometheus() *PromMetrics {
 
 	return promMetrics
 }
+
+func (m *Metrics) RecordIPAddress(ip string) {
+	if m.distinctIPWriter != nil {
+		m.distinctIPWriter.AddIPToSet(ip)
+	}
+}
+
+func (m *Metrics) SetIPAddressRecorder(recorder *sinkcluster.ClusterWriter) {
+	m.distinctIPWriter = recorder
+}