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 +}