Export snowflake broker metrics for prometheus

This change adds a prometheus exporter for our existing snowflake broker
metrics. Current values for the metrics can be fetched by sending a GET
request to /prometheus.

broker/broker.go

@@ -24,6 +24,8 @@ import (
 
 	"git.torproject.org/pluggable-transports/snowflake.git/common/messages"
 	"git.torproject.org/pluggable-transports/snowflake.git/common/safelog"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"golang.org/x/crypto/acme/autocert"
 )
 
@@ -212,6 +214,7 @@ func proxyPolls(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 	if nil == offer {
 		ctx.metrics.lock.Lock()
 		ctx.metrics.proxyIdleCount++
+		promMetrics.ProxyPollTotal.With(prometheus.Labels{"nat": natType, "status": "idle"}).Inc()
 		ctx.metrics.lock.Unlock()
 
 		b, err = messages.EncodePollResponse("", false, "")
@@ -223,6 +226,7 @@ func proxyPolls(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 		w.Write(b)
 		return
 	}
+	promMetrics.ProxyPollTotal.With(prometheus.Labels{"nat": natType, "status": "matched"}).Inc()
 	b, err = messages.EncodePollResponse(string(offer.sdp), true, offer.natType)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
@@ -276,6 +280,7 @@ func clientOffers(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 	if numSnowflakes <= 0 {
 		ctx.metrics.lock.Lock()
 		ctx.metrics.clientDeniedCount++
+		promMetrics.ClientPollTotal.With(prometheus.Labels{"nat": offer.natType, "status": "denied"}).Inc()
 		if offer.natType == NATUnrestricted {
 			ctx.metrics.clientUnrestrictedDeniedCount++
 		} else {
@@ -297,6 +302,7 @@ func clientOffers(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 	case answer := <-snowflake.answerChannel:
 		ctx.metrics.lock.Lock()
 		ctx.metrics.clientProxyMatchCount++
+		promMetrics.ClientPollTotal.With(prometheus.Labels{"nat": offer.natType, "status": "matched"}).Inc()
 		ctx.metrics.lock.Unlock()
 		if _, err := w.Write(answer); err != nil {
 			log.Printf("unable to write answer with error: %v", err)
@@ -497,6 +503,9 @@ func main() {
 	http.Handle("/answer", SnowflakeHandler{ctx, proxyAnswers})
 	http.Handle("/debug", SnowflakeHandler{ctx, debugHandler})
 	http.Handle("/metrics", MetricsHandler{metricsFilename, metricsHandler})
+	http.Handle("/prometheus", promhttp.Handler())
+
+	InitPrometheus()
 
 	server := http.Server{
 		Addr: addr,

broker/metrics.go

@@ -13,13 +13,20 @@ import (
 	"sort"
 	"sync"
 	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
 var (
-	once sync.Once
+	once        sync.Once
+	promMetrics *PromMetrics
 )
 
-const metricsResolution = 60 * 60 * 24 * time.Second //86400 seconds
+const (
+	PrometheusNamespace = "snowflake"
+	metricsResolution   = 60 * 60 * 24 * time.Second //86400 seconds
+)
 
 type CountryStats struct {
 	standalone map[string]bool
@@ -140,6 +147,11 @@ 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,
+		"cc":   country,
+	}).Inc()
 
 	switch natType {
 	case NATRestricted:
@@ -246,3 +258,43 @@ func (m *Metrics) zeroMetrics() {
 func binCount(count uint) uint {
 	return uint((math.Ceil(float64(count) / 8)) * 8)
 }
+
+type PromMetrics struct {
+	ProxyTotal      *prometheus.CounterVec
+	ProxyPollTotal  *prometheus.CounterVec
+	ClientPollTotal *prometheus.CounterVec
+}
+
+//Initialize metrics for prometheus exporter
+func InitPrometheus() {
+
+	promMetrics = &PromMetrics{}
+
+	promMetrics.ProxyTotal = promauto.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace: PrometheusNamespace,
+			Name:      "proxy_total",
+			Help:      "The number of unique snowflake IPs",
+		},
+		[]string{"type", "nat", "cc"},
+	)
+
+	promMetrics.ProxyPollTotal = promauto.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace: PrometheusNamespace,
+			Name:      "proxy_poll_total",
+			Help:      "The number of snowflake proxy polls",
+		},
+		[]string{"nat", "status"},
+	)
+
+	promMetrics.ClientPollTotal = promauto.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace: PrometheusNamespace,
+			Name:      "client_poll_total",
+			Help:      "The number of snowflake client polls",
+		},
+		[]string{"nat", "status"},
+	)
+
+}

broker/snowflake-broker_test.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"sync"
 	"testing"
 	"time"
 
@@ -21,8 +22,12 @@ func NullLogger() *log.Logger {
 	return logger
 }
 
+var promOnce sync.Once
+
 func TestBroker(t *testing.T) {
 
+	promOnce.Do(InitPrometheus)
+
 	Convey("Context", t, func() {
 		ctx := NewBrokerContext(NullLogger())
 
@@ -298,6 +303,8 @@ func TestBroker(t *testing.T) {
 }
 
 func TestSnowflakeHeap(t *testing.T) {
+	promOnce.Do(InitPrometheus)
+
 	Convey("SnowflakeHeap", t, func() {
 		h := new(SnowflakeHeap)
 		heap.Init(h)
@@ -341,6 +348,8 @@ 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")
@@ -445,6 +454,7 @@ func TestGeoip(t *testing.T) {
 }
 
 func TestMetrics(t *testing.T) {
+	promOnce.Do(InitPrometheus)
 
 	Convey("Test metrics...", t, func() {
 		done := make(chan bool)