From 9fd8d46a836f7297601caaa86891a1cfc87ba053 Mon Sep 17 00:00:00 2001
From: Tolga Ozen <mtolgaozen@gmail.com>
Date: Tue, 11 Mar 2025 14:37:41 +0300
Subject: [PATCH] refactor(balancer): update ClientConn to use
 ClientConnWrapper

---
 pkg/balancer/balancer.go      |  2 +-
 pkg/balancer/balancer_test.go |  6 ++++++
 pkg/balancer/wrapper.go       | 25 +++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 pkg/balancer/wrapper.go

diff --git a/pkg/balancer/balancer.go b/pkg/balancer/balancer.go
index 1a720e792..e44583cdc 100644
--- a/pkg/balancer/balancer.go
+++ b/pkg/balancer/balancer.go
@@ -18,7 +18,7 @@ type Balancer struct {
 	state connectivity.State
 
 	// The ClientConn to communicate with the gRPC client.
-	clientConn balancer.ClientConn
+	clientConn ClientConnWrapper
 
 	// Current picker used to select SubConns for requests.
 	picker balancer.Picker
diff --git a/pkg/balancer/balancer_test.go b/pkg/balancer/balancer_test.go
index 6af9914da..da8801311 100644
--- a/pkg/balancer/balancer_test.go
+++ b/pkg/balancer/balancer_test.go
@@ -4,6 +4,8 @@ import (
 	"errors"
 	"testing"
 
+	estats "google.golang.org/grpc/experimental/stats"
+
 	"github.com/cespare/xxhash/v2"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
@@ -27,6 +29,10 @@ type mockClientConn struct {
 	newSubConn  balancer.SubConn
 }
 
+func (m *mockClientConn) MetricsRecorder() estats.MetricsRecorder {
+	return nil
+}
+
 func (m *mockClientConn) UpdateState(state balancer.State) {
 	m.updateState = state
 }
diff --git a/pkg/balancer/wrapper.go b/pkg/balancer/wrapper.go
new file mode 100644
index 000000000..621f45ac9
--- /dev/null
+++ b/pkg/balancer/wrapper.go
@@ -0,0 +1,25 @@
+package balancer
+
+import (
+	"google.golang.org/grpc/balancer"
+	estats "google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/resolver"
+)
+
+// ClientConnWrapper is an interface that wraps the gRPC ClientConn interface.
+type ClientConnWrapper interface {
+	// NewSubConn creates a new SubConn with the specified addresses and options.
+	NewSubConn([]resolver.Address, balancer.NewSubConnOptions) (balancer.SubConn, error)
+	// RemoveSubConn removes the specified SubConn.
+	RemoveSubConn(balancer.SubConn)
+	// UpdateAddresses updates the addresses for the specified SubConn.
+	UpdateAddresses(balancer.SubConn, []resolver.Address)
+	// UpdateState updates the overall connectivity state of the balancer.
+	UpdateState(balancer.State)
+	// ResolveNow triggers an immediate resolution of the target.
+	ResolveNow(resolver.ResolveNowOptions)
+	// Target returns the target URI of the connection.
+	Target() string
+	// MetricsRecorder returns the metrics recorder for the connection.
+	MetricsRecorder() estats.MetricsRecorder
+}