snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

balancer/endpointsharding/endpointsharding.go

@@ -0,0 +1,385 @@
+/*
+ *
+ * Copyright 2024 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package endpointsharding implements a load balancing policy that manages
+// homogeneous child policies each owning a single endpoint.
+//
+// # Experimental
+//
+// Notice: This package is EXPERIMENTAL and may be changed or removed in a
+// later release.
+package endpointsharding
+
+import (
+	"errors"
+	rand "math/rand/v2"
+	"sync"
+	"sync/atomic"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/resolver"
+)
+
+var randIntN = rand.IntN
+
+// ChildState is the state of a child balancer.
+type ChildState struct {
+	Endpoint resolver.Endpoint // Endpoint of the child balancer.
+	State    balancer.State    // State of the child balancer.
+	ExitIdle func()            // Function to exit the child balancer from IDLE state.
+}
+
+// Options configure the behaviour of the endpointsharding balancer.
+type Options struct {
+	// DisableAutoReconnect allows the balancer to keep child balancer in the
+	// IDLE state until they are explicitly triggered to exit using the
+	// ChildState obtained from the endpointsharding picker. When set to false,
+	// the endpointsharding balancer will automatically call ExitIdle on child
+	// connections that report IDLE.
+	DisableAutoReconnect bool
+}
+
+// ChildBuilderFunc creates a new balancer with the ClientConn. It has the same
+// type as the balancer.Builder.Build method.
+type ChildBuilderFunc func(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer
+
+// NewBalancer returns a load balancing policy that manages homogeneous child
+// policies each owning a single endpoint. The endpointsharding balancer
+// forwards the LoadBalancingConfig in ClientConn state updates to its children.
+func NewBalancer(cc balancer.ClientConn, opts balancer.BuildOptions, childBuilder ChildBuilderFunc, esOpts Options) balancer.Balancer {
+	return &endpointSharding{
+		cc:           cc,
+		bOpts:        opts,
+		esOpts:       esOpts,
+		childBuilder: childBuilder,
+		endpoints:    resolver.NewEndpointMap[*endpointState](),
+	}
+}
+
+// endpointSharding is a balancer that wraps child balancers. It creates a child
+// balancer with child config for every unique Endpoint received. It updates the
+// child states on any update from parent or child.
+type endpointSharding struct {
+	cc           balancer.ClientConn
+	bOpts        balancer.BuildOptions
+	esOpts       Options
+	childBuilder ChildBuilderFunc
+
+	// mu guards access to the below fields and guarantees mutual exclusion
+	// between top-down methods (like UpdateClientConnState, ResolverError etc,
+	// which are already serialized) and bottom-up methods (like UpdateState)
+	// that can be called concurrently.
+	//
+	// Lock ordering & deadlock prevention:
+	// Child callbacks (like UpdateState) acquire this mutex to push state updates
+	// to the parent. This establishes a strict lock ordering:
+	//   [endpointState.childMu] -> [endpointSharding.mu]
+	//
+	// To prevent deadlocks, we must never invert this order. Therefore, we must
+	// never call any methods on a child balancer while holding this mutex. If we
+	// did, and that child balancer invoked UpdateState synchronously, it would
+	// attempt to re-acquire this mutex, causing a deadlock.
+	//
+	// Concurrency:
+	// Top-down operations need to iterate over all children. To ensure a
+	// single, clean aggregated update at the end of such operations, we inhibit
+	// intermediate updates from children. We grab this mutex briefly to set
+	// `inhibitChildUpdates = true` and immediately release it. This allows us
+	// to perform the potentially slow, top-down child updates without holding
+	// any parent locks. Once finished, we grab the mutex again to unset the
+	// flag and push the final aggregated state. An update from a child during
+	// this time will *only* update the child state, and will not access the
+	// endpoints map or push an aggregated state to the parent.
+	mu                  sync.Mutex
+	endpoints           *resolver.EndpointMap[*endpointState]
+	inhibitChildUpdates bool
+}
+
+// rotateEndpoints returns a slice of all the input endpoints rotated a random
+// amount.
+func rotateEndpoints(es []resolver.Endpoint) []resolver.Endpoint {
+	n := len(es)
+	if n == 0 {
+		return es
+	}
+	r := randIntN(n)
+
+	// Make a copy to avoid mutating data beyond the end of es.
+	ret := make([]resolver.Endpoint, n)
+	copy(ret, es[r:])
+	copy(ret[n-r:], es[:r])
+	return ret
+}
+
+// UpdateClientConnState creates a child for new endpoints and deletes children
+// for endpoints that are no longer present. It also updates all the children,
+// and sends a single synchronous update of the childrens' aggregated state at
+// the end of the UpdateClientConnState operation.
+//
+// Returns the first error found from a child, but fully processes the update.
+func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState) error {
+	es.inhibitUpdatesFromChildren()
+
+	// Update/create child balancers for each endpoint in the update. Note that we
+	// don't hold the mutex here, but this is fine because inhibitChildUpdates is
+	// true, and therefore UpdateState will not access es.endpoints.
+	var retErr error
+	newEndpoints := resolver.NewEndpointMap[*endpointState]()
+	for _, endpoint := range rotateEndpoints(state.ResolverState.Endpoints) {
+		if _, ok := newEndpoints.Get(endpoint); ok {
+			// Skip duplicate endpoints.
+			continue
+		}
+		epState, ok := es.endpoints.Get(endpoint)
+		if ok {
+			// Endpoint child already exists, update the stored endpoint.
+			epState.endpoint = endpoint
+		} else {
+			// Endpoint child does not exist, create a new one.
+			epState = &endpointState{
+				ClientConn:           es.cc,
+				parent:               es,
+				endpoint:             endpoint,
+				disableAutoReconnect: es.esOpts.DisableAutoReconnect,
+			}
+			epState.childLB = es.childBuilder(epState, es.bOpts)
+		}
+		// Update the endpoint state for the endpoint.
+		newEndpoints.Set(endpoint, epState)
+
+		if err := epState.updateClientConnState(balancer.ClientConnState{
+			BalancerConfig: state.BalancerConfig,
+			ResolverState: resolver.State{
+				Endpoints:  []resolver.Endpoint{endpoint},
+				Attributes: state.ResolverState.Attributes,
+			},
+		}); err != nil && retErr == nil {
+			// Keep the first error found from any child.
+			retErr = err
+		}
+	}
+
+	// Delete old children that are no longer present.
+	for e, child := range es.endpoints.All() {
+		if _, ok := newEndpoints.Get(e); !ok {
+			child.close()
+		}
+	}
+
+	if newEndpoints.Len() == 0 {
+		retErr = balancer.ErrBadResolverState
+	}
+
+	es.mu.Lock()
+	es.endpoints = newEndpoints
+	es.inhibitChildUpdates = false
+	es.updateStateLocked()
+	es.mu.Unlock()
+
+	return retErr
+}
+
+// ResolverError forwards the resolver error to all of the endpointSharding's
+// children and sends a single synchronous update of the childStates at the end
+// of the ResolverError operation.
+func (es *endpointSharding) ResolverError(err error) {
+	es.inhibitUpdatesFromChildren()
+	for _, child := range es.endpoints.All() {
+		child.resolverError(err)
+	}
+	es.allowUpdatesFromChildren()
+}
+
+func (es *endpointSharding) UpdateSubConnState(balancer.SubConn, balancer.SubConnState) {
+	// UpdateSubConnState is deprecated.
+}
+
+func (es *endpointSharding) Close() {
+	es.inhibitUpdatesFromChildren()
+	for _, child := range es.endpoints.All() {
+		child.close()
+	}
+}
+
+func (es *endpointSharding) ExitIdle() {
+	es.inhibitUpdatesFromChildren()
+	for _, child := range es.endpoints.All() {
+		child.exitIdle()
+	}
+	es.allowUpdatesFromChildren()
+}
+
+func (es *endpointSharding) inhibitUpdatesFromChildren() {
+	es.mu.Lock()
+	es.inhibitChildUpdates = true
+	es.mu.Unlock()
+}
+
+func (es *endpointSharding) allowUpdatesFromChildren() {
+	es.mu.Lock()
+	es.inhibitChildUpdates = false
+	es.updateStateLocked()
+	es.mu.Unlock()
+}
+
+// updateStateLocked updates this component's state. It sends the aggregated
+// state, and a picker with round robin behavior with all the child states
+// present if needed. This method must only be called when inhibitChildUpdates
+// is false.
+//
+// Caller must hold es.mu.
+func (es *endpointSharding) updateStateLocked() {
+	var readyPickers, connectingPickers, idlePickers, transientFailurePickers []balancer.Picker
+
+	childStates := make([]ChildState, 0, es.endpoints.Len())
+	for _, epState := range es.endpoints.All() {
+		childState := ChildState{
+			Endpoint: epState.endpoint,
+			State:    epState.state,
+			ExitIdle: func() { go epState.exitIdle() },
+		}
+		childStates = append(childStates, childState)
+		childPicker := childState.State.Picker
+		switch childState.State.ConnectivityState {
+		case connectivity.Ready:
+			readyPickers = append(readyPickers, childPicker)
+		case connectivity.Connecting:
+			connectingPickers = append(connectingPickers, childPicker)
+		case connectivity.Idle:
+			idlePickers = append(idlePickers, childPicker)
+		case connectivity.TransientFailure:
+			transientFailurePickers = append(transientFailurePickers, childPicker)
+			// connectivity.Shutdown shouldn't appear.
+		}
+	}
+
+	// Construct the round robin picker based off the aggregated state. Whatever
+	// the aggregated state, use the pickers present that are currently in that
+	// state only.
+	var aggState connectivity.State
+	var pickers []balancer.Picker
+	if len(readyPickers) >= 1 {
+		aggState = connectivity.Ready
+		pickers = readyPickers
+	} else if len(connectingPickers) >= 1 {
+		aggState = connectivity.Connecting
+		pickers = connectingPickers
+	} else if len(idlePickers) >= 1 {
+		aggState = connectivity.Idle
+		pickers = idlePickers
+	} else if len(transientFailurePickers) >= 1 {
+		aggState = connectivity.TransientFailure
+		pickers = transientFailurePickers
+	} else {
+		aggState = connectivity.TransientFailure
+		pickers = []balancer.Picker{base.NewErrPicker(errors.New("no children to pick from"))}
+	} // No children (resolver error before valid update).
+
+	es.cc.UpdateState(balancer.State{
+		ConnectivityState: aggState,
+		Picker: &pickerWithChildStates{
+			pickers:     pickers,
+			childStates: childStates,
+			next:        uint32(randIntN(len(pickers))),
+		},
+	})
+}
+
+// pickerWithChildStates delegates to the pickers it holds in a round robin
+// fashion. It also contains the childStates of all the endpointSharding's
+// children.
+type pickerWithChildStates struct {
+	pickers     []balancer.Picker
+	childStates []ChildState
+	next        uint32
+}
+
+func (p *pickerWithChildStates) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	nextIndex := atomic.AddUint32(&p.next, 1)
+	picker := p.pickers[nextIndex%uint32(len(p.pickers))]
+	return picker.Pick(info)
+}
+
+// ChildStatesFromPicker returns the state of all the children managed by the
+// endpoint sharding balancer that created this picker.
+func ChildStatesFromPicker(picker balancer.Picker) []ChildState {
+	p, ok := picker.(*pickerWithChildStates)
+	if !ok {
+		return nil
+	}
+	return p.childStates
+}
+
+// endpointState is the internal state maintained for each endpoint.
+type endpointState struct {
+	balancer.ClientConn // Embedded to intercept UpdateState
+
+	parent               *endpointSharding // Parent endpointsharding balancer.
+	endpoint             resolver.Endpoint // Endpoint of the child balancer.
+	state                balancer.State    // State of the child balancer.
+	disableAutoReconnect bool              // Whether to disable auto reconnect for this child.
+
+	childMu sync.Mutex        // Guarantees mutual exclusion for Balancer API calls to the child balancer.
+	childLB balancer.Balancer // Child balancer.
+	closed  bool              // Tracks closure of the child balancer to ensure ExitIdle is not called after Close().
+}
+
+func (es *endpointState) UpdateState(state balancer.State) {
+	es.parent.mu.Lock()
+	es.state = state
+	if !es.parent.inhibitChildUpdates {
+		es.parent.updateStateLocked()
+	}
+	es.parent.mu.Unlock()
+
+	if state.ConnectivityState == connectivity.Idle && !es.disableAutoReconnect {
+		go es.exitIdle()
+	}
+}
+
+func (es *endpointState) updateClientConnState(state balancer.ClientConnState) error {
+	es.childMu.Lock()
+	err := es.childLB.UpdateClientConnState(state)
+	es.childMu.Unlock()
+	return err
+}
+
+func (es *endpointState) resolverError(err error) {
+	es.childMu.Lock()
+	es.childLB.ResolverError(err)
+	es.childMu.Unlock()
+}
+
+func (es *endpointState) close() {
+	es.childMu.Lock()
+	if !es.closed {
+		es.closed = true
+		es.childLB.Close()
+	}
+	es.childMu.Unlock()
+}
+
+func (es *endpointState) exitIdle() {
+	es.childMu.Lock()
+	if !es.closed {
+		es.childLB.ExitIdle()
+	}
+	es.childMu.Unlock()
+}