Skip to content

Commit

Permalink
balancer: Enforce embedding requirement for balancer.ClientConn (grpc…
Browse files Browse the repository at this point in the history
  • Loading branch information
arjan-bal authored Feb 6, 2025
1 parent b963f4b commit 3e27c17
Show file tree
Hide file tree
Showing 7 changed files with 32 additions and 21 deletions.
12 changes: 12 additions & 0 deletions balancer/balancer.go
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,13 @@ type State struct {
// brand new implementation of this interface. For the situations like
// testing, the new implementation should embed this interface. This allows
// gRPC to add new methods to this interface.
//
// NOTICE: This interface is intended to be implemented by gRPC, or intercepted
// by custom load balancing polices. Users should not need their own complete
// implementation of this interface -- they should always delegate to a
// ClientConn passed to Builder.Build() by embedding it in their
// implementations. An embedded ClientConn must never be nil, or runtime panics
// will occur.
type ClientConn interface {
// NewSubConn is called by balancer to create a new SubConn.
// It doesn't block and wait for the connections to be established.
Expand Down Expand Up @@ -167,6 +174,11 @@ type ClientConn interface {
//
// Deprecated: Use the Target field in the BuildOptions instead.
Target() string

// EnforceClientConnEmbedding is included to force implementers to embed
// another implementation of this interface, allowing gRPC to add methods
// without breaking users.
internal.EnforceClientConnEmbedding
}

// BuildOptions contains additional information for Build.
Expand Down
2 changes: 1 addition & 1 deletion balancer/subconn.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ import (
// should only use a single address.
//
// NOTICE: This interface is intended to be implemented by gRPC, or intercepted
// by custom load balancing poilices. Users should not need their own complete
// by custom load balancing polices. Users should not need their own complete
// implementation of this interface -- they should always delegate to a SubConn
// returned by ClientConn.NewSubConn() by embedding it in their implementations.
// An embedded SubConn must never be nil, or runtime panics will occur.
Expand Down
1 change: 1 addition & 0 deletions balancer_wrapper.go
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ var (
// It uses the gracefulswitch.Balancer internally to ensure that balancer
// switches happen in a graceful manner.
type ccBalancerWrapper struct {
internal.EnforceClientConnEmbedding
// The following fields are initialized when the wrapper is created and are
// read-only afterwards, and therefore can be accessed without a mutex.
cc *ClientConn
Expand Down
10 changes: 4 additions & 6 deletions internal/balancer/gracefulswitch/gracefulswitch.go
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,9 @@ func (gsb *Balancer) switchTo(builder balancer.Builder) (*balancerWrapper, error
return nil, errBalancerClosed
}
bw := &balancerWrapper{
builder: builder,
gsb: gsb,
ClientConn: gsb.cc,
builder: builder,
gsb: gsb,
lastState: balancer.State{
ConnectivityState: connectivity.Connecting,
Picker: base.NewErrPicker(balancer.ErrNoSubConnAvailable),
Expand Down Expand Up @@ -293,6 +294,7 @@ func (gsb *Balancer) Close() {
// State updates from the wrapped balancer can result in invocation of the
// graceful switch logic.
type balancerWrapper struct {
balancer.ClientConn
balancer.Balancer
gsb *Balancer
builder balancer.Builder
Expand Down Expand Up @@ -413,7 +415,3 @@ func (bw *balancerWrapper) UpdateAddresses(sc balancer.SubConn, addrs []resolver
bw.gsb.mu.Unlock()
bw.gsb.cc.UpdateAddresses(sc, addrs)
}

func (bw *balancerWrapper) Target() string {
return bw.gsb.cc.Target()
}
6 changes: 6 additions & 0 deletions internal/internal.go
Original file line number Diff line number Diff line change
Expand Up @@ -277,3 +277,9 @@ const RLSLoadBalancingPolicyName = "rls_experimental"
type EnforceSubConnEmbedding interface {
enforceSubConnEmbedding()
}

// EnforceClientConnEmbedding is used to enforce proper ClientConn implementation
// embedding.
type EnforceClientConnEmbedding interface {
enforceClientConnEmbedding()
}
2 changes: 2 additions & 0 deletions internal/testutils/balancer.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (

"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/internal/grpcsync"
"google.golang.org/grpc/resolver"
)
Expand Down Expand Up @@ -98,6 +99,7 @@ func (*TestSubConn) RegisterHealthListener(func(balancer.SubConnState)) {}

// BalancerClientConn is a mock balancer.ClientConn used in tests.
type BalancerClientConn struct {
internal.EnforceClientConnEmbedding
logger Logger

NewSubConnAddrsCh chan []resolver.Address // the last 10 []Address to create subconn.
Expand Down
20 changes: 6 additions & 14 deletions xds/internal/balancer/outlierdetection/balancer.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ type bb struct{}

func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
b := &outlierDetectionBalancer{
cc: cc,
ClientConn: cc,
closed: grpcsync.NewEvent(),
done: grpcsync.NewEvent(),
addrs: make(map[string]*endpointInfo),
Expand Down Expand Up @@ -158,6 +158,7 @@ type scHealthUpdate struct {
}

type outlierDetectionBalancer struct {
balancer.ClientConn
// These fields are safe to be accessed without holding any mutex because
// they are synchronized in run(), which makes these field accesses happen
// serially.
Expand All @@ -171,7 +172,6 @@ type outlierDetectionBalancer struct {

closed *grpcsync.Event
done *grpcsync.Event
cc balancer.ClientConn
logger *grpclog.PrefixLogger
channelzParent channelz.Identifier

Expand Down Expand Up @@ -480,7 +480,7 @@ func (b *outlierDetectionBalancer) NewSubConn(addrs []resolver.Address, opts bal
opts.StateListener = func(state balancer.SubConnState) { b.updateSubConnState(scw, state) }
b.mu.Lock()
defer b.mu.Unlock()
sc, err := b.cc.NewSubConn(addrs, opts)
sc, err := b.ClientConn.NewSubConn(addrs, opts)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -544,7 +544,7 @@ func (b *outlierDetectionBalancer) UpdateAddresses(sc balancer.SubConn, addrs []
return
}

b.cc.UpdateAddresses(scw.SubConn, addrs)
b.ClientConn.UpdateAddresses(scw.SubConn, addrs)
b.mu.Lock()
defer b.mu.Unlock()

Expand Down Expand Up @@ -586,14 +586,6 @@ func (b *outlierDetectionBalancer) UpdateAddresses(sc balancer.SubConn, addrs []
scw.addresses = addrs
}

func (b *outlierDetectionBalancer) ResolveNow(opts resolver.ResolveNowOptions) {
b.cc.ResolveNow(opts)
}

func (b *outlierDetectionBalancer) Target() string {
return b.cc.Target()
}

// handleSubConnUpdate stores the recent state and forward the update
// if the SubConn is not ejected.
func (b *outlierDetectionBalancer) handleSubConnUpdate(u *scUpdate) {
Expand Down Expand Up @@ -628,7 +620,7 @@ func (b *outlierDetectionBalancer) handleChildStateUpdate(u balancer.State) {
noopCfg := b.noopConfig()
b.mu.Unlock()
b.recentPickerNoop = noopCfg
b.cc.UpdateState(balancer.State{
b.ClientConn.UpdateState(balancer.State{
ConnectivityState: b.childState.ConnectivityState,
Picker: &wrappedPicker{
childPicker: b.childState.Picker,
Expand All @@ -652,7 +644,7 @@ func (b *outlierDetectionBalancer) handleLBConfigUpdate(u lbCfgUpdate) {
// the bit.
if b.childState.Picker != nil && noopCfg != b.recentPickerNoop || b.updateUnconditionally {
b.recentPickerNoop = noopCfg
b.cc.UpdateState(balancer.State{
b.ClientConn.UpdateState(balancer.State{
ConnectivityState: b.childState.ConnectivityState,
Picker: &wrappedPicker{
childPicker: b.childState.Picker,
Expand Down

0 comments on commit 3e27c17

Please sign in to comment.