From 2f64d9797a883bd91fa60fa148e3bd02165d6fad Mon Sep 17 00:00:00 2001
From: Joseph Anttila Hall <jkh@google.com>
Date: Wed, 25 Jan 2023 15:07:48 -0800
Subject: [PATCH] Fix konnectivity-client metric registration.

Kubernetes-commit: a9edc2c2191ae9b9fb527f9d65fa7d73406261ed
---
 pkg/server/egressselector/egress_selector.go  |  4 +-
 .../egressselector/egress_selector_test.go    | 69 +++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/pkg/server/egressselector/egress_selector.go b/pkg/server/egressselector/egress_selector.go
index 6d2763397..0936d6ef4 100644
--- a/pkg/server/egressselector/egress_selector.go
+++ b/pkg/server/egressselector/egress_selector.go
@@ -36,7 +36,7 @@ import (
 	utilnet "k8s.io/apimachinery/pkg/util/net"
 	"k8s.io/apiserver/pkg/apis/apiserver"
 	egressmetrics "k8s.io/apiserver/pkg/server/egressselector/metrics"
-	compbasemetrics "k8s.io/component-base/metrics"
+	"k8s.io/component-base/metrics/legacyregistry"
 	"k8s.io/component-base/tracing"
 	"k8s.io/klog/v2"
 	client "sigs.k8s.io/apiserver-network-proxy/konnectivity-client/pkg/client"
@@ -45,7 +45,7 @@ import (
 var directDialer utilnet.DialFunc = http.DefaultTransport.(*http.Transport).DialContext
 
 func init() {
-	client.Metrics.RegisterMetrics(compbasemetrics.NewKubeRegistry().Registerer())
+	client.Metrics.RegisterMetrics(legacyregistry.Registerer())
 }
 
 // EgressSelector is the map of network context type to context dialer, for network egress.
diff --git a/pkg/server/egressselector/egress_selector_test.go b/pkg/server/egressselector/egress_selector_test.go
index 1ac555d2e..896b24461 100644
--- a/pkg/server/egressselector/egress_selector_test.go
+++ b/pkg/server/egressselector/egress_selector_test.go
@@ -18,6 +18,7 @@ package egressselector
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
 	"strings"
@@ -31,6 +32,9 @@ import (
 	"k8s.io/component-base/metrics/legacyregistry"
 	"k8s.io/component-base/metrics/testutil"
 	testingclock "k8s.io/utils/clock/testing"
+	clientmetrics "sigs.k8s.io/apiserver-network-proxy/konnectivity-client/pkg/client/metrics"
+	ccmetrics "sigs.k8s.io/apiserver-network-proxy/konnectivity-client/pkg/common/metrics"
+	"sigs.k8s.io/apiserver-network-proxy/konnectivity-client/proto/client"
 )
 
 type fakeEgressSelection struct {
@@ -272,5 +276,70 @@ func TestMetrics(t *testing.T) {
 			}
 		})
 	}
+}
 
+func TestKonnectivityClientMetrics(t *testing.T) {
+	testcases := []struct {
+		name    string
+		metrics []string
+		trigger func()
+		want    string
+	}{
+		{
+			name:    "stream packets",
+			metrics: []string{"konnectivity_network_proxy_client_stream_packets_total"},
+			trigger: func() {
+				clientmetrics.Metrics.ObservePacket(ccmetrics.SegmentFromClient, client.PacketType_DIAL_REQ)
+			},
+			want: `
+# HELP konnectivity_network_proxy_client_stream_packets_total Count of packets processed, by segment and packet type (example: from_client, DIAL_REQ)
+# TYPE konnectivity_network_proxy_client_stream_packets_total counter
+konnectivity_network_proxy_client_stream_packets_total{packet_type="DIAL_REQ",segment="from_client"} 1
+`,
+		},
+		{
+			name:    "stream errors",
+			metrics: []string{"konnectivity_network_proxy_client_stream_errors_total"},
+			trigger: func() {
+				clientmetrics.Metrics.ObserveStreamError(ccmetrics.SegmentToClient, errors.New("example"), client.PacketType_DIAL_RSP)
+			},
+			want: `
+# HELP konnectivity_network_proxy_client_stream_errors_total Count of gRPC stream errors, by segment, grpc Code, packet type. (example: from_agent, Code.Unavailable, DIAL_RSP)
+# TYPE konnectivity_network_proxy_client_stream_errors_total counter
+konnectivity_network_proxy_client_stream_errors_total{code="Unknown",packet_type="DIAL_RSP",segment="to_client"} 1
+`,
+		},
+		{
+			name:    "dial failure",
+			metrics: []string{"konnectivity_network_proxy_client_dial_failure_total"},
+			trigger: func() {
+				clientmetrics.Metrics.ObserveDialFailure(clientmetrics.DialFailureTimeout)
+			},
+			want: `
+# HELP konnectivity_network_proxy_client_dial_failure_total Number of dial failures observed, by reason (example: remote endpoint error)
+# TYPE konnectivity_network_proxy_client_dial_failure_total counter
+konnectivity_network_proxy_client_dial_failure_total{reason="timeout"} 1
+`,
+		},
+		{
+			name:    "client connections",
+			metrics: []string{"konnectivity_network_proxy_client_client_connections"},
+			trigger: func() {
+				clientmetrics.Metrics.GetClientConnectionsMetric().WithLabelValues("dialing").Inc()
+			},
+			want: `
+# HELP konnectivity_network_proxy_client_client_connections Number of open client connections, by status (Example: dialing)
+# TYPE konnectivity_network_proxy_client_client_connections gauge
+konnectivity_network_proxy_client_client_connections{status="dialing"} 1
+`,
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			tc.trigger()
+			if err := testutil.GatherAndCompare(legacyregistry.DefaultGatherer, strings.NewReader(tc.want), tc.metrics...); err != nil {
+				t.Errorf("GatherAndCompare error: %v", err)
+			}
+		})
+	}
 }