From d3d3289d631ab7837434dc14ff840e945be30892 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal.matczuk@saucelabs.com>
Date: Fri, 8 Nov 2024 12:28:40 +0100
Subject: [PATCH] net: MultiListener add name label to prometheus metrics

---
 net.go                                       |  9 ++-
 net_metrics.go                               | 31 +++++++++
 net_test.go                                  | 67 ++++++++++++++++++++
 testdata/TestMultiListenerMetrics.golden.txt | 12 ++++
 4 files changed, 114 insertions(+), 5 deletions(-)
 create mode 100644 testdata/TestMultiListenerMetrics.golden.txt

diff --git a/net.go b/net.go
index 21895be5..04fbcc17 100644
--- a/net.go
+++ b/net.go
@@ -197,6 +197,7 @@ type NamedListenerConfig struct {
 }
 
 // MultiListener is a builder for multiple listeners sharing the same prometheus configuration.
+// The listener name is added as a label to the metrics.
 type MultiListener struct {
 	ListenerConfigs []NamedListenerConfig
 	TLSConfig       func(NamedListenerConfig) *tls.Config
@@ -204,10 +205,6 @@ type MultiListener struct {
 }
 
 func (ml MultiListener) Listen() (_ []net.Listener, ferr error) {
-	prototype := Listener{
-		metrics: newListenerMetrics(ml.PromRegistry, ml.PromNamespace),
-	}
-
 	listeners := make([]net.Listener, 0, len(ml.ListenerConfigs))
 	defer func() {
 		if ferr != nil {
@@ -217,13 +214,15 @@ func (ml MultiListener) Listen() (_ []net.Listener, ferr error) {
 		}
 	}()
 
+	mf := newListenerMetricsWithNameFunc(ml.PromRegistry, ml.PromNamespace)
+
 	for _, lc := range ml.ListenerConfigs {
 		l := new(Listener)
-		*l = prototype
 		l.ListenerConfig = lc.ListenerConfig
 		if ml.TLSConfig != nil {
 			l.TLSConfig = ml.TLSConfig(lc)
 		}
+		l.metrics = mf(lc.Name)
 		if err := l.Listen(); err != nil {
 			return nil, err
 		}
diff --git a/net_metrics.go b/net_metrics.go
index 574fab4c..a01c0993 100644
--- a/net_metrics.go
+++ b/net_metrics.go
@@ -123,3 +123,34 @@ func (m *listenerMetrics) error() {
 func (m *listenerMetrics) close() {
 	m.closed.Inc()
 }
+
+func newListenerMetricsWithNameFunc(r prometheus.Registerer, namespace string) func(name string) *listenerMetrics {
+	if r == nil {
+		r = prometheus.NewRegistry() // This registry will be discarded.
+	}
+	f := promauto.With(r)
+
+	accepted := f.NewCounterVec(prometheus.CounterOpts{
+		Name:      "listener_accepted_total",
+		Namespace: namespace,
+		Help:      "Number of accepted connections",
+	}, []string{"name"})
+	errors := f.NewCounterVec(prometheus.CounterOpts{
+		Name:      "listener_errors_total",
+		Namespace: namespace,
+		Help:      "Number of listener errors when accepting connections",
+	}, []string{"name"})
+	closed := f.NewCounterVec(prometheus.CounterOpts{
+		Name:      "listener_closed_total",
+		Namespace: namespace,
+		Help:      "Number of closed connections",
+	}, []string{"name"})
+
+	return func(name string) *listenerMetrics {
+		return &listenerMetrics{
+			accepted: accepted.WithLabelValues(name),
+			errors:   errors.WithLabelValues(name),
+			closed:   closed.WithLabelValues(name),
+		}
+	}
+}
diff --git a/net_test.go b/net_test.go
index c2519bb4..29eef53d 100644
--- a/net_test.go
+++ b/net_test.go
@@ -394,3 +394,70 @@ func selfSingedCert() *tls.Config {
 		Certificates: []tls.Certificate{cert},
 	}
 }
+
+func (ml *MultiListener) listenAndWait(t *testing.T) []net.Listener {
+	t.Helper()
+
+	listeners, err := ml.Listen()
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, l := range listeners {
+		for {
+			if l.Addr() != nil {
+				break
+			}
+		}
+	}
+	return listeners
+}
+
+func TestMultiListenerMetrics(t *testing.T) {
+	r := prometheus.NewRegistry()
+	ml := MultiListener{
+		ListenerConfigs: []NamedListenerConfig{
+			{
+				Name: "a",
+				ListenerConfig: ListenerConfig{
+					Address: "localhost:0",
+				},
+			},
+			{
+				Name: "b",
+				ListenerConfig: ListenerConfig{
+					Address: "localhost:0",
+				},
+			},
+		},
+		PromConfig: PromConfig{
+			PromNamespace: "test",
+			PromRegistry:  r,
+		},
+	}
+	listeners := ml.listenAndWait(t)
+	defer func() {
+		for _, l := range listeners {
+			l.Close()
+		}
+	}()
+
+	for _, l := range listeners {
+		go l.(*Listener).acceptAndCopy() //nolint:forcetypeassert // trust the test
+	}
+
+	for _, l := range listeners {
+		for range 10 {
+			conn, err := net.Dial("tcp", l.Addr().String())
+			if err != nil {
+				t.Fatalf("net.Dial(): got %v, want no error", err)
+			}
+			fmt.Fprintf(conn, "Hello, World!\n")
+			if _, err := conn.Read(make([]byte, 1)); err != nil {
+				t.Fatal(err)
+			}
+			conn.Close()
+		}
+	}
+
+	golden.DiffPrometheusMetrics(t, r)
+}
diff --git a/testdata/TestMultiListenerMetrics.golden.txt b/testdata/TestMultiListenerMetrics.golden.txt
new file mode 100644
index 00000000..428e55b2
--- /dev/null
+++ b/testdata/TestMultiListenerMetrics.golden.txt
@@ -0,0 +1,12 @@
+# HELP test_listener_accepted_total Number of accepted connections
+# TYPE test_listener_accepted_total counter
+test_listener_accepted_total{name="a"} 10
+test_listener_accepted_total{name="b"} 10
+# HELP test_listener_closed_total Number of closed connections
+# TYPE test_listener_closed_total counter
+test_listener_closed_total{name="a"} 10
+test_listener_closed_total{name="b"} 10
+# HELP test_listener_errors_total Number of listener errors when accepting connections
+# TYPE test_listener_errors_total counter
+test_listener_errors_total{name="a"} 0
+test_listener_errors_total{name="b"} 0