resolver: convert EndpointMap to use generics

balancer/endpointsharding/endpointsharding.go

@@ -73,7 +73,7 @@ func NewBalancer(cc balancer.ClientConn, opts balancer.BuildOptions, childBuilde
 		esOpts:       esOpts,
 		childBuilder: childBuilder,
 	}
-	es.children.Store(resolver.NewEndpointMap())
+	es.children.Store(resolver.NewEndpointMap[*balancerWrapper]())
 	return es
 }
 
@@ -90,7 +90,7 @@ type endpointSharding struct {
 	// calls into a child. To avoid deadlocks, do not acquire childMu while
 	// holding mu.
 	childMu  sync.Mutex
-	children atomic.Pointer[resolver.EndpointMap] // endpoint -> *balancerWrapper
+	children atomic.Pointer[resolver.EndpointMap[*balancerWrapper]]
 
 	// inhibitChildUpdates is set during UpdateClientConnState/ResolverError
 	// calls (calls to children will each produce an update, only want one
@@ -122,7 +122,7 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 	var ret error
 
 	children := es.children.Load()
-	newChildren := resolver.NewEndpointMap()
+	newChildren := resolver.NewEndpointMap[*balancerWrapper]()
 
 	// Update/Create new children.
 	for _, endpoint := range state.ResolverState.Endpoints {
@@ -131,9 +131,8 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 			// update.
 			continue
 		}
-		var childBalancer *balancerWrapper
-		if val, ok := children.Get(endpoint); ok {
-			childBalancer = val.(*balancerWrapper)
+		childBalancer, ok := children.Get(endpoint)
+		if ok {
 			// Endpoint attributes may have changed, update the stored endpoint.
 			es.mu.Lock()
 			childBalancer.childState.Endpoint = endpoint
@@ -166,7 +165,7 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 	for _, e := range children.Keys() {
 		child, _ := children.Get(e)
 		if _, ok := newChildren.Get(e); !ok {
-			child.(*balancerWrapper).closeLocked()
+			child.closeLocked()
 		}
 	}
 	es.children.Store(newChildren)
@@ -189,7 +188,7 @@ func (es *endpointSharding) ResolverError(err error) {
 	}()
 	children := es.children.Load()
 	for _, child := range children.Values() {
-		child.(*balancerWrapper).resolverErrorLocked(err)
+		child.resolverErrorLocked(err)
 	}
 }
 
@@ -202,7 +201,7 @@ func (es *endpointSharding) Close() {
 	defer es.childMu.Unlock()
 	children := es.children.Load()
 	for _, child := range children.Values() {
-		child.(*balancerWrapper).closeLocked()
+		child.closeLocked()
 	}
 }
 
@@ -222,8 +221,7 @@ func (es *endpointSharding) updateState() {
 	childStates := make([]ChildState, 0, children.Len())
 
 	for _, child := range children.Values() {
-		bw := child.(*balancerWrapper)
-		childState := bw.childState
+		childState := child.childState
 		childStates = append(childStates, childState)
 		childPicker := childState.State.Picker
 		switch childState.State.ConnectivityState {

balancer/leastrequest/leastrequest.go

@@ -88,7 +88,7 @@ func (bb) Name() string {
 func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
 	b := &leastRequestBalancer{
 		ClientConn:        cc,
-		endpointRPCCounts: resolver.NewEndpointMap(),
+		endpointRPCCounts: resolver.NewEndpointMap[*atomic.Int32](),
 	}
 	b.child = endpointsharding.NewBalancer(b, bOpts, balancer.Get(pickfirstleaf.Name).Build, endpointsharding.Options{})
 	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[%p] ", b))
@@ -110,7 +110,7 @@ type leastRequestBalancer struct {
 	choiceCount uint32
 	// endpointRPCCounts holds RPC counts to keep track for subsequent picker
 	// updates.
-	endpointRPCCounts *resolver.EndpointMap // endpoint -> *atomic.Int32
+	endpointRPCCounts *resolver.EndpointMap[*atomic.Int32]
 }
 
 func (lrb *leastRequestBalancer) Close() {
@@ -164,7 +164,7 @@ func (lrb *leastRequestBalancer) UpdateState(state balancer.State) {
 	}
 
 	// Reconcile endpoints.
-	newEndpoints := resolver.NewEndpointMap() // endpoint -> nil
+	newEndpoints := resolver.NewEndpointMap[any]()
 	for _, child := range readyEndpoints {
 		newEndpoints.Set(child.Endpoint, nil)
 	}
@@ -179,13 +179,11 @@ func (lrb *leastRequestBalancer) UpdateState(state balancer.State) {
 	// Copy refs to counters into picker.
 	endpointStates := make([]endpointState, 0, len(readyEndpoints))
 	for _, child := range readyEndpoints {
-		var counter *atomic.Int32
-		if val, ok := lrb.endpointRPCCounts.Get(child.Endpoint); !ok {
+		counter, ok := lrb.endpointRPCCounts.Get(child.Endpoint)
+		if !ok {
 			// Create new counts if needed.
 			counter = new(atomic.Int32)
 			lrb.endpointRPCCounts.Set(child.Endpoint, counter)
-		} else {
-			counter = val.(*atomic.Int32)
 		}
 		endpointStates = append(endpointStates, endpointState{
 			picker:  child.State.Picker,

balancer/weightedroundrobin/balancer.go

@@ -105,7 +105,7 @@ func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Ba
 		target:           bOpts.Target.String(),
 		metricsRecorder:  cc.MetricsRecorder(),
 		addressWeights:   resolver.NewAddressMapV2[*endpointWeight](),
-		endpointToWeight: resolver.NewEndpointMap(),
+		endpointToWeight: resolver.NewEndpointMap[*endpointWeight](),
 		scToWeight:       make(map[balancer.SubConn]*endpointWeight),
 	}
 
@@ -155,17 +155,15 @@ func (bb) Name() string {
 //
 // Caller must hold b.mu.
 func (b *wrrBalancer) updateEndpointsLocked(endpoints []resolver.Endpoint) {
-	endpointSet := resolver.NewEndpointMap()
+	endpointSet := resolver.NewEndpointMap[*endpointWeight]()
 	addressSet := resolver.NewAddressMapV2[*endpointWeight]()
 	for _, endpoint := range endpoints {
 		endpointSet.Set(endpoint, nil)
 		for _, addr := range endpoint.Addresses {
 			addressSet.Set(addr, nil)
 		}
-		var ew *endpointWeight
-		if ewi, ok := b.endpointToWeight.Get(endpoint); ok {
-			ew = ewi.(*endpointWeight)
-		} else {
+		ew, ok := b.endpointToWeight.Get(endpoint)
+		if !ok {
 			ew = &endpointWeight{
 				logger:            b.logger,
 				connectivityState: connectivity.Connecting,
@@ -215,7 +213,7 @@ type wrrBalancer struct {
 	locality         string
 	stopPicker       *grpcsync.Event
 	addressWeights   *resolver.AddressMapV2[*endpointWeight]
-	endpointToWeight *resolver.EndpointMap // endpoint -> endpointWeight
+	endpointToWeight *resolver.EndpointMap[*endpointWeight]
 	scToWeight       map[balancer.SubConn]*endpointWeight
 }
 
@@ -260,13 +258,12 @@ func (b *wrrBalancer) UpdateState(state balancer.State) {
 
 	for _, childState := range childStates {
 		if childState.State.ConnectivityState == connectivity.Ready {
-			ewv, ok := b.endpointToWeight.Get(childState.Endpoint)
+			ew, ok := b.endpointToWeight.Get(childState.Endpoint)
 			if !ok {
 				// Should never happen, simply continue and ignore this endpoint
 				// for READY pickers.
 				continue
 			}
-			ew := ewv.(*endpointWeight)
 			readyPickersWeight = append(readyPickersWeight, pickerWeightedEndpoint{
 				picker:           childState.State.Picker,
 				weightedEndpoint: ew,
@@ -398,8 +395,7 @@ func (b *wrrBalancer) Close() {
 	b.mu.Unlock()
 
 	// Ensure any lingering OOB watchers are stopped.
-	for _, ewv := range b.endpointToWeight.Values() {
-		ew := ewv.(*endpointWeight)
+	for _, ew := range b.endpointToWeight.Values() {
 		if ew.stopORCAListener != nil {
 			ew.stopORCAListener()
 		}

resolver/map.go

@@ -162,21 +162,21 @@ type endpointMapKey string
 // unordered set of address strings within an endpoint. This map is not thread
 // safe, thus it is unsafe to access concurrently. Must be created via
 // NewEndpointMap; do not construct directly.
-type EndpointMap struct {
-	endpoints map[endpointMapKey]endpointData
+type EndpointMap[T any] struct {
+	endpoints map[endpointMapKey]endpointData[T]
 }
 
-type endpointData struct {
+type endpointData[T any] struct {
 	// decodedKey stores the original key to avoid decoding when iterating on
 	// EndpointMap keys.
 	decodedKey Endpoint
-	value      any
+	value      T
 }
 
 // NewEndpointMap creates a new EndpointMap.
-func NewEndpointMap() *EndpointMap {
-	return &EndpointMap{
-		endpoints: make(map[endpointMapKey]endpointData),
+func NewEndpointMap[T any]() *EndpointMap[T] {
+	return &EndpointMap[T]{
+		endpoints: make(map[endpointMapKey]endpointData[T]),
 	}
 }
 
@@ -196,25 +196,25 @@ func encodeEndpoint(e Endpoint) endpointMapKey {
 }
 
 // Get returns the value for the address in the map, if present.
-func (em *EndpointMap) Get(e Endpoint) (value any, ok bool) {
+func (em *EndpointMap[T]) Get(e Endpoint) (value T, ok bool) {
 	val, found := em.endpoints[encodeEndpoint(e)]
 	if found {
 		return val.value, true
 	}
-	return nil, false
+	return value, false
 }
 
 // Set updates or adds the value to the address in the map.
-func (em *EndpointMap) Set(e Endpoint, value any) {
+func (em *EndpointMap[T]) Set(e Endpoint, value T) {
 	en := encodeEndpoint(e)
-	em.endpoints[en] = endpointData{
+	em.endpoints[en] = endpointData[T]{
 		decodedKey: Endpoint{Addresses: e.Addresses},
 		value:      value,
 	}
 }
 
 // Len returns the number of entries in the map.
-func (em *EndpointMap) Len() int {
+func (em *EndpointMap[T]) Len() int {
 	return len(em.endpoints)
 }
 
@@ -223,7 +223,7 @@ func (em *EndpointMap) Len() int {
 // the unordered set of addresses. Thus, endpoint information returned is not
 // the full endpoint data (drops duplicated addresses and attributes) but can be
 // used for EndpointMap accesses.
-func (em *EndpointMap) Keys() []Endpoint {
+func (em *EndpointMap[T]) Keys() []Endpoint {
 	ret := make([]Endpoint, 0, len(em.endpoints))
 	for _, en := range em.endpoints {
 		ret = append(ret, en.decodedKey)
@@ -232,16 +232,16 @@ func (em *EndpointMap) Keys() []Endpoint {
 }
 
 // Values returns a slice of all current map values.
-func (em *EndpointMap) Values() []any {
-	ret := make([]any, 0, len(em.endpoints))
+func (em *EndpointMap[T]) Values() []T {
+	ret := make([]T, 0, len(em.endpoints))
 	for _, val := range em.endpoints {
 		ret = append(ret, val.value)
 	}
 	return ret
 }
 
 // Delete removes the specified endpoint from the map.
-func (em *EndpointMap) Delete(e Endpoint) {
+func (em *EndpointMap[T]) Delete(e Endpoint) {
 	en := encodeEndpoint(e)
 	delete(em.endpoints, en)
 }