Implement binned counts for polling metrics

4 files changed, 108 insertions, 23 deletions

diff --git a/broker/metrics.go b/broker/metrics.go
index 6939742..24ff9b0 100644
--- a/broker/metrics.go
+++ b/broker/metrics.go
@@ -20,11 +20,11 @@ import (
 
 var (
 	once        sync.Once
-	promMetrics *PromMetrics
+	promMetrics = initPrometheus()
 )
 
 const (
-	PrometheusNamespace = "snowflake"
+	prometheusNamespace = "snowflake"
 	metricsResolution   = 60 * 60 * 24 * time.Second //86400 seconds
 )
 
@@ -147,6 +147,7 @@ func (m *Metrics) UpdateCountryStats(addr string, proxyType string, natType stri
 	} else {
 		m.countryStats.unknown[addr] = true
 	}
+
 	promMetrics.ProxyTotal.With(prometheus.Labels{
 		"nat":  natType,
 		"type": proxyType,
@@ -261,40 +262,47 @@ func binCount(count uint) uint {
 
 type PromMetrics struct {
 	ProxyTotal      *prometheus.CounterVec
-	ProxyPollTotal  *prometheus.CounterVec
-	ClientPollTotal *prometheus.CounterVec
+	ProxyPollTotal  *RoundedCounterVec
+	ClientPollTotal *RoundedCounterVec
 }
 
 //Initialize metrics for prometheus exporter
-func InitPrometheus() {
+func initPrometheus() *PromMetrics {
 
-	promMetrics = &PromMetrics{}
+	promMetrics := &PromMetrics{}
 
 	promMetrics.ProxyTotal = promauto.NewCounterVec(
 		prometheus.CounterOpts{
-			Namespace: PrometheusNamespace,
+			Namespace: prometheusNamespace,
 			Name:      "proxy_total",
 			Help:      "The number of unique snowflake IPs",
 		},
 		[]string{"type", "nat", "cc"},
 	)
 
-	promMetrics.ProxyPollTotal = promauto.NewCounterVec(
+	promMetrics.ProxyPollTotal = NewRoundedCounterVec(
 		prometheus.CounterOpts{
-			Namespace: PrometheusNamespace,
-			Name:      "proxy_poll_total",
-			Help:      "The number of snowflake proxy polls",
+			Namespace: prometheusNamespace,
+			Name:      "rounded_proxy_poll_total",
+			Help:      "The number of snowflake proxy polls, rounded up to a multiple of 8",
 		},
 		[]string{"nat", "status"},
 	)
 
-	promMetrics.ClientPollTotal = promauto.NewCounterVec(
+	promMetrics.ClientPollTotal = NewRoundedCounterVec(
 		prometheus.CounterOpts{
-			Namespace: PrometheusNamespace,
-			Name:      "client_poll_total",
-			Help:      "The number of snowflake client polls",
+			Namespace: prometheusNamespace,
+			Name:      "rounded_client_poll_total",
+			Help:      "The number of snowflake client polls, rounded up to a multiple of 8",
 		},
 		[]string{"nat", "status"},
 	)
 
+	// We need to register this new metric type because there is no constructor
+	// for it in promauto.
+	prometheus.DefaultRegisterer.MustRegister(promMetrics.ClientPollTotal)
+	prometheus.DefaultRegisterer.MustRegister(promMetrics.ProxyPollTotal)
+
+	return promMetrics
+
 }
diff --git a/broker/prometheus.go b/broker/prometheus.go
new file mode 100644
index 0000000..d7592ec
--- /dev/null
+++ b/broker/prometheus.go
@@ -0,0 +1,83 @@
+/*
+Implements some additional prometheus metrics that we need for privacy preserving
+counts of users and proxies
+*/
+
+package main
+
+import (
+	"sync/atomic"
+
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
+	"google.golang.org/protobuf/proto"
+)
+
+// New Prometheus counter type that produces rounded counts of metrics
+// for privacy preserving reasons
+type RoundedCounter interface {
+	prometheus.Metric
+
+	Inc()
+}
+
+type roundedCounter struct {
+	total uint64 //reflects the true count
+	value uint64 //reflects the rounded count
+
+	desc       *prometheus.Desc
+	labelPairs []*dto.LabelPair
+}
+
+// Implements the RoundedCounter interface
+func (c *roundedCounter) Inc() {
+	atomic.AddUint64(&c.total, 1)
+	if c.total > c.value {
+		atomic.AddUint64(&c.value, 8)
+	}
+}
+
+// Implements the prometheus.Metric interface
+func (c *roundedCounter) Desc() *prometheus.Desc {
+	return c.desc
+}
+
+// Implements the prometheus.Metric interface
+func (c *roundedCounter) Write(m *dto.Metric) error {
+	m.Label = c.labelPairs
+
+	m.Counter = &dto.Counter{Value: proto.Float64(float64(c.value))}
+	return nil
+}
+
+// New prometheus vector type that will track RoundedCounter metrics
+// accross multiple labels
+type RoundedCounterVec struct {
+	*prometheus.MetricVec
+}
+
+func NewRoundedCounterVec(opts prometheus.CounterOpts, labelNames []string) *RoundedCounterVec {
+	desc := prometheus.NewDesc(
+		prometheus.BuildFQName(opts.Namespace, opts.Subsystem, opts.Name),
+		opts.Help,
+		labelNames,
+		opts.ConstLabels,
+	)
+	return &RoundedCounterVec{
+		MetricVec: prometheus.NewMetricVec(desc, func(lvs ...string) prometheus.Metric {
+			if len(lvs) != len(labelNames) {
+				panic("inconsistent cardinality")
+			}
+			return &roundedCounter{desc: desc, labelPairs: prometheus.MakeLabelPairs(desc, lvs)}
+		}),
+	}
+}
+
+// Helper function to return the underlying RoundedCounter metric from MetricVec
+func (v *RoundedCounterVec) With(labels prometheus.Labels) RoundedCounter {
+	metric, err := v.GetMetricWith(labels)
+	if err != nil {
+		panic(err)
+	}
+	return metric.(RoundedCounter)
+}
diff --git a/broker/snowflake-broker_test.go b/broker/snowflake-broker_test.go
index 987aae8..b676b04 100644
--- a/broker/snowflake-broker_test.go
+++ b/broker/snowflake-broker_test.go
@@ -26,8 +26,6 @@ var promOnce sync.Once
 
 func TestBroker(t *testing.T) {
 
-	promOnce.Do(InitPrometheus)
-
 	Convey("Context", t, func() {
 		ctx := NewBrokerContext(NullLogger())
 
@@ -303,8 +301,6 @@ func TestBroker(t *testing.T) {
 }
 
 func TestSnowflakeHeap(t *testing.T) {
-	promOnce.Do(InitPrometheus)
-
 	Convey("SnowflakeHeap", t, func() {
 		h := new(SnowflakeHeap)
 		heap.Init(h)
@@ -348,8 +344,6 @@ func TestSnowflakeHeap(t *testing.T) {
 }
 
 func TestGeoip(t *testing.T) {
-	promOnce.Do(InitPrometheus)
-
 	Convey("Geoip", t, func() {
 		tv4 := new(GeoIPv4Table)
 		err := GeoIPLoadFile(tv4, "test_geoip")
@@ -454,8 +448,6 @@ func TestGeoip(t *testing.T) {
 }
 
 func TestMetrics(t *testing.T) {
-	promOnce.Do(InitPrometheus)
-
 	Convey("Test metrics...", t, func() {
 		done := make(chan bool)
 		buf := new(bytes.Buffer)
diff --git a/go.mod b/go.mod
index ab3dc96..ed07394 100644
--- a/go.mod
+++ b/go.mod
@@ -12,10 +12,12 @@ require (
 	github.com/pion/transport v0.12.3 // indirect
 	github.com/pion/webrtc/v3 v3.0.15
 	github.com/prometheus/client_golang v1.10.0
+	github.com/prometheus/client_model v0.2.0
 	github.com/smartystreets/goconvey v1.6.4
 	github.com/xtaci/kcp-go/v5 v5.5.12
 	github.com/xtaci/smux v1.5.12
 	golang.org/x/crypto v0.0.0-20210317152858-513c2a44f670
 	golang.org/x/net v0.0.0-20210316092652-d523dce5a7f4
 	golang.org/x/sys v0.0.0-20210317225723-c4fcb01b228e // indirect
+	google.golang.org/protobuf v1.23.0
 )