diff --git a/broker/metrics.go b/broker/metrics.go index 639d505..cd1ca37 100644 --- a/broker/metrics.go +++ b/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 +}