upstream-grpc-go/internal/xds/balancer/wrrlocality/balancer.go

/*
 *
 * Copyright 2023 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 wrrlocality provides an implementation of the wrr locality LB policy,
// as defined in [A52 - xDS Custom LB Policies].
//
// [A52 - xDS Custom LB Policies]: https://github.com/grpc/proposal/blob/master/A52-xds-custom-lb-policies.md
package wrrlocality

import (
	"encoding/json"
	"errors"
	"fmt"

	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/weightedtarget"
	"google.golang.org/grpc/internal/grpclog"
	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
	xdsinternal "google.golang.org/grpc/internal/xds"
	"google.golang.org/grpc/resolver"
	"google.golang.org/grpc/serviceconfig"
)

// Name is the name of wrr_locality balancer.
const Name = "xds_wrr_locality_experimental"

func init() {
	balancer.Register(bb{})
}

type bb struct{}

func (bb) Name() string {
	return Name
}

// LBConfig is the config for the wrr locality balancer.
type LBConfig struct {
	serviceconfig.LoadBalancingConfig `json:"-"`
	// ChildPolicy is the config for the child policy.
	ChildPolicy *internalserviceconfig.BalancerConfig `json:"childPolicy,omitempty"`
}

// To plumb in a different child in tests.
var weightedTargetName = weightedtarget.Name

func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
	builder := balancer.Get(weightedTargetName)
	if builder == nil {
		// Shouldn't happen, registered through imported weighted target,
		// defensive programming.
		return nil
	}

	// Doesn't need to intercept any balancer.ClientConn operations; pass
	// through by just giving cc to child balancer.
	wtb := builder.Build(cc, bOpts)
	if wtb == nil {
		// shouldn't happen, defensive programming.
		return nil
	}
	wtbCfgParser, ok := builder.(balancer.ConfigParser)
	if !ok {
		// Shouldn't happen, imported weighted target builder has this method.
		return nil
	}
	wrrL := &wrrLocalityBalancer{
		child:       wtb,
		childParser: wtbCfgParser,
	}

	wrrL.logger = prefixLogger(wrrL)
	wrrL.logger.Infof("Created")
	return wrrL
}

func (bb) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
	var lbCfg *LBConfig
	if err := json.Unmarshal(s, &lbCfg); err != nil {
		return nil, fmt.Errorf("xds_wrr_locality: invalid LBConfig: %s, error: %v", string(s), err)
	}
	if lbCfg == nil || lbCfg.ChildPolicy == nil {
		return nil, errors.New("xds_wrr_locality: invalid LBConfig: child policy field must be set")
	}
	return lbCfg, nil
}

type attributeKey struct{}

// Equal allows the values to be compared by Attributes.Equal.
func (a AddrInfo) Equal(o any) bool {
	oa, ok := o.(AddrInfo)
	return ok && oa.LocalityWeight == a.LocalityWeight
}

// AddrInfo is the locality weight of the locality an address is a part of.
type AddrInfo struct {
	LocalityWeight uint32
}

// SetAddrInfo returns a copy of endpoint in which the Attributes field is
// updated with AddrInfo.
func SetAddrInfo(endpoint resolver.Endpoint, addrInfo AddrInfo) resolver.Endpoint {
	endpoint.Attributes = endpoint.Attributes.WithValue(attributeKey{}, addrInfo)
	return endpoint
}

func (a AddrInfo) String() string {
	return fmt.Sprintf("Locality Weight: %d", a.LocalityWeight)
}

// getAddrInfo returns the AddrInfo stored in the Attributes field of
// ep. Returns false if no AddrInfo found.
func getAddrInfo(ep resolver.Endpoint) (AddrInfo, bool) {
	v := ep.Attributes.Value(attributeKey{})
	ai, ok := v.(AddrInfo)
	return ai, ok
}

// wrrLocalityBalancer wraps a weighted target balancer, and builds
// configuration for the weighted target once it receives configuration
// specifying the weighted target child balancer and locality weight
// information.
type wrrLocalityBalancer struct {
	// child will be a weighted target balancer, and will be built it at
	// wrrLocalityBalancer build time. Other than preparing configuration, other
	// balancer operations are simply pass through.
	child balancer.Balancer

	childParser balancer.ConfigParser

	logger *grpclog.PrefixLogger
}

func (b *wrrLocalityBalancer) ExitIdle() {
	b.child.ExitIdle()
}

func (b *wrrLocalityBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
	lbCfg, ok := s.BalancerConfig.(*LBConfig)
	if !ok {
		b.logger.Errorf("Received config with unexpected type %T: %v", s.BalancerConfig, s.BalancerConfig)
		return balancer.ErrBadResolverState
	}

	weightedTargets := make(map[string]weightedtarget.Target)
	for _, ep := range s.ResolverState.Endpoints {
		// This get of LocalityID could potentially return a zero value. This
		// shouldn't happen though (this attribute that is set actually gets
		// used to build localities in the first place), and thus don't error
		// out, and just build a weighted target with undefined behavior.
		locality := xdsinternal.LocalityString(xdsinternal.LocalityIDFromEndpoint(ep))
		ai, ok := getAddrInfo(ep)
		if !ok {
			return fmt.Errorf("xds_wrr_locality: missing locality weight information in endpoint %q", ep)
		}
		weightedTargets[locality] = weightedtarget.Target{Weight: ai.LocalityWeight, ChildPolicy: lbCfg.ChildPolicy}
	}
	wtCfg := &weightedtarget.LBConfig{Targets: weightedTargets}
	wtCfgJSON, err := json.Marshal(wtCfg)
	if err != nil {
		// Shouldn't happen.
		return fmt.Errorf("xds_wrr_locality: error marshalling prepared config: %v", wtCfg)
	}
	var sc serviceconfig.LoadBalancingConfig
	if sc, err = b.childParser.ParseConfig(wtCfgJSON); err != nil {
		return fmt.Errorf("xds_wrr_locality: config generated %v is invalid: %v", wtCfgJSON, err)
	}

	return b.child.UpdateClientConnState(balancer.ClientConnState{
		ResolverState:  s.ResolverState,
		BalancerConfig: sc,
	})
}

func (b *wrrLocalityBalancer) ResolverError(err error) {
	b.child.ResolverError(err)
}

func (b *wrrLocalityBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
	b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
}

func (b *wrrLocalityBalancer) Close() {
	b.child.Close()
}