Merge pull request #56 from Jigsaw-Code/bemasc-metrics-bench

Add benchmarks for the metrics functions

Author: Benjamin M. Schwartz

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "metrics/test-data"]
+	path = third_party/maxmind
+	url = https://github.com/maxmind/MaxMind-DB.git

README.md

@@ -108,7 +108,16 @@ Run the iperf3 client tests listed above on port 10002.
 
 You can mix and match the libev and go servers and clients.
 
-## Benchmark
+## Tests and Benchmarks
+
+Before running tests, you should first run
+```
+git submodule update --init
+```
+to download test data used by the GeoIP metrics tests.  To run all tests, you can use
+```
+go test -v ./...
+```
 
 You can benchmark the cipher finding code with
 ```

metrics/metrics.go

@@ -63,8 +63,8 @@ type shadowsocksMetrics struct {
 	udpRemovedNatEntries prometheus.Counter
 }
 
-func NewShadowsocksMetrics(ipCountryDB *geoip2.Reader) ShadowsocksMetrics {
-	m := &shadowsocksMetrics{
+func newShadowsocksMetrics(ipCountryDB *geoip2.Reader) *shadowsocksMetrics {
+	return &shadowsocksMetrics{
 		ipCountryDB: ipCountryDB,
 		accessKeys: prometheus.NewGauge(prometheus.GaugeOpts{
 			Namespace: "shadowsocks",
@@ -130,8 +130,16 @@ func NewShadowsocksMetrics(ipCountryDB *geoip2.Reader) ShadowsocksMetrics {
 				Help:      "Entries removed from the UDP NAT table",
 			}),
 	}
+}
+
+// NewPrometheusShadowsocksMetrics constructs a metrics object that uses
+// `ipCountryDB` to convert IP addresses to countries, and reports all
+// metrics to Prometheus via `registerer`.  `ipCountryDB` may be nil, but
+// `registerer` must not be.
+func NewPrometheusShadowsocksMetrics(ipCountryDB *geoip2.Reader, registerer prometheus.Registerer) ShadowsocksMetrics {
+	m := newShadowsocksMetrics(ipCountryDB)
 	// TODO: Is it possible to pass where to register the collectors?
-	prometheus.MustRegister(m.accessKeys, m.ports, m.tcpOpenConnections, m.tcpProbes, m.tcpClosedConnections, m.tcpConnectionDurationMs,
+	registerer.MustRegister(m.accessKeys, m.ports, m.tcpOpenConnections, m.tcpProbes, m.tcpClosedConnections, m.tcpConnectionDurationMs,
 		m.dataBytes, m.timeToCipherMs, m.udpAddedNatEntries, m.udpRemovedNatEntries)
 	return m
 }

metrics/metrics_test.go

@@ -1,12 +1,16 @@
 package metrics
 
 import (
+	"net"
 	"testing"
 	"time"
+
+	geoip2 "github.com/oschwald/geoip2-golang"
+	"github.com/prometheus/client_golang/prometheus"
 )
 
 func TestMethodsDontPanic(t *testing.T) {
-	ssMetrics := NewShadowsocksMetrics(nil)
+	ssMetrics := NewPrometheusShadowsocksMetrics(nil, prometheus.NewPedanticRegistry())
 	proxyMetrics := ProxyMetrics{
 		ClientProxy: 1,
 		ProxyTarget: 2,
@@ -22,3 +26,94 @@ func TestMethodsDontPanic(t *testing.T) {
 	ssMetrics.AddUDPNatEntry()
 	ssMetrics.RemoveUDPNatEntry()
 }
+
+func BenchmarkGetLocation(b *testing.B) {
+	var ipCountryDB *geoip2.Reader
+	// The test data is in a git submodule that must be initialized before running the test.
+	dbPath := "../third_party/maxmind/test-data/GeoIP2-Country-Test.mmdb"
+	ipCountryDB, err := geoip2.Open(dbPath)
+	if err != nil {
+		b.Fatalf("Could not open geoip database at %v: %v", dbPath, err)
+	}
+	defer ipCountryDB.Close()
+
+	ssMetrics := NewPrometheusShadowsocksMetrics(ipCountryDB, prometheus.NewRegistry())
+	testIP := net.ParseIP("217.65.48.1")
+	testAddr := &net.TCPAddr{IP: testIP, Port: 12345}
+	b.ResetTimer()
+	// Repeatedly check the country for the same address.  This is realistic, because
+	// servers call this method for each new connection, but typically many connections
+	// come from a single user in succession.
+	for i := 0; i < b.N; i++ {
+		ssMetrics.GetLocation(testAddr)
+	}
+}
+
+func BenchmarkOpenTCP(b *testing.B) {
+	ssMetrics := NewPrometheusShadowsocksMetrics(nil, prometheus.NewRegistry())
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ssMetrics.AddOpenTCPConnection("ZZ")
+	}
+}
+
+func BenchmarkCloseTCP(b *testing.B) {
+	ssMetrics := NewPrometheusShadowsocksMetrics(nil, prometheus.NewRegistry())
+	clientLocation := "ZZ"
+	accessKey := "key 1"
+	status := "OK"
+	data := ProxyMetrics{}
+	timeToCipher := time.Microsecond
+	duration := time.Minute
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ssMetrics.AddClosedTCPConnection(clientLocation, accessKey, status, data, timeToCipher, duration)
+	}
+}
+
+func BenchmarkProbe(b *testing.B) {
+	ssMetrics := NewPrometheusShadowsocksMetrics(nil, prometheus.NewRegistry())
+	clientLocation := "ZZ"
+	status := "ERR_REPLAY"
+	drainResult := "other"
+	port := 12345
+	data := ProxyMetrics{}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ssMetrics.AddTCPProbe(clientLocation, status, drainResult, port, data)
+	}
+}
+
+func BenchmarkClientUDP(b *testing.B) {
+	ssMetrics := NewPrometheusShadowsocksMetrics(nil, prometheus.NewRegistry())
+	clientLocation := "ZZ"
+	accessKey := "key 1"
+	status := "OK"
+	size := 1000
+	timeToCipher := time.Microsecond
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ssMetrics.AddUDPPacketFromClient(clientLocation, accessKey, status, size, size, timeToCipher)
+	}
+}
+
+func BenchmarkTargetUDP(b *testing.B) {
+	ssMetrics := NewPrometheusShadowsocksMetrics(nil, prometheus.NewRegistry())
+	clientLocation := "ZZ"
+	accessKey := "key 1"
+	status := "OK"
+	size := 1000
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ssMetrics.AddUDPPacketFromTarget(clientLocation, accessKey, status, size, size)
+	}
+}
+
+func BenchmarkNAT(b *testing.B) {
+	ssMetrics := NewPrometheusShadowsocksMetrics(nil, prometheus.NewRegistry())
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ssMetrics.AddUDPNatEntry()
+		ssMetrics.RemoveUDPNatEntry()
+	}
+}

server.go

@@ -32,6 +32,7 @@ import (
 	"github.com/Jigsaw-Code/outline-ss-server/shadowsocks"
 	"github.com/op/go-logging"
 	"github.com/oschwald/geoip2-golang"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/shadowsocks/go-shadowsocks2/core"
 	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
@@ -246,7 +247,7 @@ func main() {
 		}
 		defer ipCountryDB.Close()
 	}
-	m := metrics.NewShadowsocksMetrics(ipCountryDB)
+	m := metrics.NewPrometheusShadowsocksMetrics(ipCountryDB, prometheus.DefaultRegisterer)
 	err = runSSServer(flags.ConfigFile, flags.natTimeout, m, flags.replayHistory)
 	if err != nil {
 		logger.Fatal(err)