Repository: grpc/grpc-go
File: balancer/randomsubsetting/randomsubsetting.go (Go)

Code before:
	}

	if b.cfg == nil || b.cfg.ChildPolicy.Name != lbCfg.ChildPolicy.Name {
		err := b.child.SwitchTo(bb)
		if err != nil {
			return fmt.Errorf("subsetting: error switching to child of type %q: %v", lbCfg.ChildPolicy.Name, err)
		}
	}
	b.cfg = lbCfg

	err := b.child.UpdateClientConnState(balancer.ClientConnState{
		ResolverState:  b.prepareChildResolverState(s.ResolverState),
		BalancerConfig: b.cfg.ChildPolicy.Config,
	})

	return err
}

type AddressWithHash struct {
	hash uint64
	addr resolver.Address
}

// implements the subsetting algorithm, as described in A68: https://github.com/grpc/proposal/pull/423
func (b *subsettingBalancer) prepareChildResolverState(s resolver.State) resolver.State {
	addresses := s.Addresses
	backendCount := len(addresses)
	if backendCount <= int(b.cfg.SubsetSize) {
		return s
	}

	addressesSet := make([]AddressWithHash, backendCount)
	// calculate hash for each endpoint
	for i, endpoint := range addresses {

		b.hashf.Write([]byte(s.Addresses[0].String()))
		addressesSet[i] = AddressWithHash{
			hash: b.hashf.Sum64(),
			addr: endpoint,
		}
	}
	// sort addresses by hash
	sort.Slice(addressesSet, func(i, j int) bool {
		return addressesSet[i].hash < addressesSet[j].hash
	})

	b.logger.Infof("resulting subset: %v", addressesSet[:b.cfg.SubsetSize])

	// Convert back to resolver.addresses
	addressesSubset := make([]resolver.Address, b.cfg.SubsetSize)
	for _, eh := range addressesSet[:b.cfg.SubsetSize] {

Review comment (suggestion):
We should be using the `Endpoints` field in `resolver.State` for everything here (except for computing the hash where we use the first address inside of the endpoint).

Code after:
	cfg := []map[string]any{{lbCfg.ChildPolicy.Name: lbCfg.ChildPolicy.Config}}
	cfgJSON, _ := json.Marshal(cfg)
	parsedCfg, err := gracefulswitch.ParseConfig(cfgJSON)
	if err != nil {
		return fmt.Errorf("randomsubsetting: error switching to child of type %q: %v", lbCfg.ChildPolicy.Name, err)
	}
	b.cfg = lbCfg
	endpoints := resolver.State{
		Endpoints:     b.calculateSubset(s.ResolverState.Endpoints),
		ServiceConfig: s.ResolverState.ServiceConfig,
		Attributes:    s.ResolverState.Attributes,
	}

	return b.Balancer.UpdateClientConnState(balancer.ClientConnState{
		ResolverState:  endpoints,
		BalancerConfig: parsedCfg,
	})
}

// calculateSubset implements the subsetting algorithm, as described in A68:
// https://github.com/grpc/proposal/blob/master/A68-random-subsetting.md#subsetting-algorithm
func (b *subsettingBalancer) calculateSubset(endpoints []resolver.Endpoint) []resolver.Endpoint {
	// A helper struct to hold an endpoint and its hash.
	type endpointWithHash struct {
		hash uint64
		ep   resolver.Endpoint
	}

	subsetSize := b.cfg.SubsetSize
	if len(endpoints) <= int(subsetSize) {
		return endpoints
	}

	hashedEndpoints := make([]endpointWithHash, len(endpoints))
	for i, endpoint := range endpoints {
		// For every endpoint in the list, compute a hash with previously
		// generated seed - A68.
		//
		// The xxhash package's Sum64() function does not allow setting a seed.
		// This means that we need to reset the digest with the seed for every
		// endpoint. Without this, an endpoint will not retain the same hash
		// across resolver updates.
		//
		// Note that we only hash the first address of the endpoint, as per A68.
		b.hashDigest.ResetWithSeed(b.hashSeed)
		b.hashDigest.WriteString(endpoint.Addresses[0].String())
		hashedEndpoints[i] = endpointWithHash{
			hash: b.hashDigest.Sum64(),
			ep:   endpoint,
		}
	}