interop/orcalb: Delegate SubConn management to pick_first #8914
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -24,135 +24,204 @@ import ( | |
| "sync" | ||
| "time" | ||
|
|
||
| "google.golang.org/grpc/balancer" | ||
| "google.golang.org/grpc/balancer/endpointsharding" | ||
| "google.golang.org/grpc/balancer/pickfirst" | ||
| "google.golang.org/grpc/connectivity" | ||
| "google.golang.org/grpc/grpclog" | ||
| internalgrpclog "google.golang.org/grpc/internal/grpclog" | ||
| "google.golang.org/grpc/orca" | ||
| "google.golang.org/grpc/resolver" | ||
|
|
||
| v3orcapb "github.com/cncf/xds/go/xds/data/orca/v3" | ||
| ) | ||
|
|
||
| var orcaLogger = grpclog.Component("orca") | ||
|
|
||
| func init() { | ||
| balancer.Register(orcabb{}) | ||
| } | ||
|
|
||
| type orcabb struct{} | ||
|
|
||
| // Build creates a new orcab balancer. | ||
| func (orcabb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer { | ||
| b := &orcab{ | ||
| ClientConn: cc, | ||
| stopOOBListeners: make(map[balancer.SubConn]func()), | ||
| oobState: &oobState{}, | ||
| } | ||
| b.logger = internalgrpclog.NewPrefixLogger(orcaLogger, fmt.Sprintf("[%p] ", b)) | ||
| b.child = endpointsharding.NewBalancer(b, bOpts, balancer.Get(pickfirst.Name).Build, endpointsharding.Options{}) | ||
| b.logger.Infof("Created") | ||
| return b | ||
| } | ||
|
|
||
| func (orcabb) Name() string { | ||
| return "test_backend_metrics_load_balancer" | ||
| } | ||
|
|
||
| // orcab is the balancer for the test_backend_metrics_load_balancer policy. | ||
| // It delegates SubConn management to endpointsharding + pick_first and | ||
| // intercepts NewSubConn calls to register OOB listeners on READY SubConns. | ||
| type orcab struct { | ||
| // Embeds balancer.ClientConn to intercept NewSubConn and UpdateState | ||
| // calls from child balancers. | ||
| balancer.ClientConn | ||
| child balancer.Balancer | ||
| // mu guards stopOOBListeners. | ||
| mu sync.Mutex | ||
| stopOOBListeners map[balancer.SubConn]func() | ||
|
|
||
| oobState *oobState | ||
| logger *internalgrpclog.PrefixLogger | ||
| } | ||
|
|
||
| func (b *orcab) UpdateClientConnState(s balancer.ClientConnState) error { | ||
| // Delegate to the child endpoint sharding balancer, which distributes | ||
| // state updates to its pick_first children. | ||
| return b.child.UpdateClientConnState(s) | ||
| } | ||
|
|
||
| func (b *orcab) ResolverError(err error) { | ||
| // Will cause an inline picker update from endpoint sharding. | ||
| b.child.ResolverError(err) | ||
| } | ||
|
|
||
| func (b *orcab) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) { | ||
| orcaLogger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state) | ||
| } | ||
|
|
||
| func (b *orcab) ExitIdle() { | ||
| // Propagate to the child endpoint sharding balancer. | ||
| b.child.ExitIdle() | ||
| } | ||
|
|
||
| func (b *orcab) Close() { | ||
| b.mu.Lock() | ||
| listeners := b.stopOOBListeners | ||
| b.stopOOBListeners = nil | ||
| b.mu.Unlock() | ||
| for _, stop := range listeners { | ||
| stop() | ||
| } | ||
| b.child.Close() | ||
| } | ||
|
|
||
| // NewSubConn intercepts SubConn creation from pick_first children to wrap the | ||
| // StateListener for OOB listener management on READY transitions. | ||
| func (b *orcab) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) { | ||
| // The variable sc is declared before the closure so the closure captures | ||
| // the variable, not its (zero) value. After ClientConn.NewSubConn assigns | ||
| // to sc, the closure sees the real SubConn when invoked. | ||
| var sc balancer.SubConn | ||
| oldListener := opts.StateListener | ||
| opts.StateListener = func(state balancer.SubConnState) { | ||
| b.updateSubConnState(sc, state) | ||
| if oldListener != nil { | ||
| oldListener(state) | ||
| } | ||
| } | ||
| sc, err := b.ClientConn.NewSubConn(addrs, opts) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return sc, nil | ||
| } | ||
|
|
||
| func (b *orcab) updateSubConnState(sc balancer.SubConn, state balancer.SubConnState) { | ||
| b.mu.Lock() | ||
| if b.stopOOBListeners == nil { | ||
| b.mu.Unlock() | ||
| return | ||
| } | ||
|
|
||
| if state.ConnectivityState == connectivity.Ready { | ||
zarinn3pal marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| oldStop, exists := b.stopOOBListeners[sc] | ||
| stop := orca.RegisterOOBListener(sc, &orcaOOBListener{subConn: sc, balancer: b}, orca.OOBListenerOptions{ReportInterval: time.Second}) | ||
| b.stopOOBListeners[sc] = stop | ||
| b.mu.Unlock() | ||
|
|
||
| if exists { | ||
| oldStop() | ||
| } | ||
|
Comment on lines
+145
to
+147
Contributor
There was a problem hiding this comment. Here and below, using |
||
| return | ||
| } | ||
|
|
||
| stop, ok := b.stopOOBListeners[sc] | ||
| if ok { | ||
| delete(b.stopOOBListeners, sc) | ||
|
Contributor
There was a problem hiding this comment. Can't we call Also, following up from the previous comment, we could just check for |
||
| } | ||
| b.mu.Unlock() | ||
|
|
||
| if ok { | ||
| stop() | ||
| } | ||
|
|
||
| if state.ConnectivityState == connectivity.Shutdown { | ||
| b.oobState.reports.Delete(sc) | ||
| } | ||
| } | ||
|
|
||
| // UpdateState intercepts state updates from endpointsharding to wrap the | ||
| // picker with ORCA load report handling. | ||
| func (b *orcab) UpdateState(state balancer.State) { | ||
| // If READY, wrap the picker to inject the ORCA Done callback; otherwise, | ||
| // pass through the child picker as-is. | ||
| if state.ConnectivityState == connectivity.Ready { | ||
| state = balancer.State{ | ||
| ConnectivityState: state.ConnectivityState, | ||
| Picker: &orcaPicker{childPicker: state.Picker, oobState: b.oobState}, | ||
| } | ||
| } | ||
| b.ClientConn.UpdateState(state) | ||
| } | ||
|
|
||
| type orcaOOBListener struct { | ||
| subConn balancer.SubConn | ||
| balancer *orcab | ||
| } | ||
|
|
||
| // OnLoadReport implements orca.OOBListener. | ||
| func (l *orcaOOBListener) OnLoadReport(r *v3orcapb.OrcaLoadReport) { | ||
| if r == nil { | ||
| return | ||
| } | ||
| l.balancer.oobState.reports.Store(l.subConn, r) | ||
| } | ||
|
|
||
| type oobState struct { | ||
| reports sync.Map // map[balancer.SubConn]*v3orcapb.OrcaLoadReport | ||
|
Contributor
There was a problem hiding this comment. The The Map type is optimized for two common use cases: (1) when the entry for a given key is only ever written once but read many times, as in caches that only grow, or (2) when multiple goroutines read, write, and overwrite entries for disjoint sets of keys. In these two cases, use of a Map may significantly reduce lock contention compared to a Go map paired with a separate Mutex or RWMutex. I don't think the usage here falls into either of these two cases.
I think we should revert back to using a vanilla map with appropriate synchronization using a |
||
| } | ||
| type orcaPicker struct { | ||
| childPicker balancer.Picker | ||
| oobState *oobState | ||
| } | ||
|
|
||
| func (p *orcaPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) { | ||
| res, err := p.childPicker.Pick(info) | ||
| if err != nil { | ||
| return res, err | ||
| } | ||
|
|
||
| var lr *v3orcapb.OrcaLoadReport | ||
| if val, ok := p.oobState.reports.Load(res.SubConn); ok { | ||
| lr = val.(*v3orcapb.OrcaLoadReport) | ||
| } | ||
|
|
||
| origDone := res.Done | ||
| res.Done = func(di balancer.DoneInfo) { | ||
| if perRPCLR, _ := di.ServerLoad.(*v3orcapb.OrcaLoadReport); perRPCLR != nil && | ||
| (perRPCLR.CpuUtilization != 0 || perRPCLR.MemUtilization != 0 || len(perRPCLR.Utilization) > 0 || len(perRPCLR.RequestCost) > 0) { | ||
|
Comment on lines
+214
to
+215
Contributor
There was a problem hiding this comment. I understand that some of this was existing code. But this is quite hard to read. Can we instead write this as: perRPCLR, _ := di.ServerLoad.(*v3orcapb.OrcaLoadReport)
if perRPCLR != nil && (perRPCLR.CpuUtilization != 0 || perRPCLR.MemUtilization != 0 || len(perRPCLR.Utilization) > 0 || len(perRPCLR.RequestCost) > 0) {
setContextCMR(info.Ctx, perRPCLR)
} else if lr != nil {
setContextCMR(info.Ctx, lr)
}See: https://google.github.io/styleguide/go/guide#line-length |
||
| setContextCMR(info.Ctx, perRPCLR) | ||
| } else if lr != nil { | ||
| setContextCMR(info.Ctx, lr) | ||
| } | ||
| if origDone != nil { | ||
| origDone(di) | ||
| } | ||
| } | ||
| return res, nil | ||
| } | ||
|
|
||
| func setContextCMR(ctx context.Context, lr *v3orcapb.OrcaLoadReport) { | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is a good example of how we format our structs in grpc-go:
grpc-go/internal/xds/xdsdepmgr/xds_dependency_manager.go
Line 98 in 12e91dd
Please try to follow that pattern here.