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

internal/xds/balancer/wrrlocality/balancer.go

@@ -0,0 +1,201 @@
+/*
+ *
+ * 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()
+}

internal/xds/balancer/wrrlocality/balancer_test.go

@@ -0,0 +1,250 @@
+/*
+ *
+ * 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
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/roundrobin"
+	"google.golang.org/grpc/balancer/weightedtarget"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/grpctest"
+	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/internal/testutils"
+	xdsinternal "google.golang.org/grpc/internal/xds"
+	"google.golang.org/grpc/internal/xds/clients"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+const (
+	defaultTestTimeout = 5 * time.Second
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestParseConfig(t *testing.T) {
+	const errParseConfigName = "errParseConfigBalancer"
+	stub.Register(errParseConfigName, stub.BalancerFuncs{
+		ParseConfig: func(json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+			return nil, errors.New("some error")
+		},
+	})
+
+	parser := bb{}
+	tests := []struct {
+		name    string
+		input   string
+		wantCfg serviceconfig.LoadBalancingConfig
+		wantErr string
+	}{
+		{
+			name:  "happy-case-round robin-child",
+			input: `{"childPolicy": [{"round_robin": {}}]}`,
+			wantCfg: &LBConfig{
+				ChildPolicy: &internalserviceconfig.BalancerConfig{
+					Name: roundrobin.Name,
+				},
+			},
+		},
+		{
+			name:    "invalid-json",
+			input:   "{{invalidjson{{",
+			wantErr: "invalid character",
+		},
+
+		{
+			name:    "child-policy-field-isn't-set",
+			input:   `{}`,
+			wantErr: "child policy field must be set",
+		},
+		{
+			name:    "child-policy-type-is-empty",
+			input:   `{"childPolicy": []}`,
+			wantErr: "invalid loadBalancingConfig: no supported policies found in []",
+		},
+		{
+			name:    "child-policy-empty-config",
+			input:   `{"childPolicy": [{"": {}}]}`,
+			wantErr: "invalid loadBalancingConfig: no supported policies found in []",
+		},
+		{
+			name:    "child-policy-type-isn't-registered",
+			input:   `{"childPolicy": [{"doesNotExistBalancer": {"cluster": "test_cluster"}}]}`,
+			wantErr: "invalid loadBalancingConfig: no supported policies found in [doesNotExistBalancer]",
+		},
+		{
+			name:    "child-policy-config-is-invalid",
+			input:   `{"childPolicy": [{"errParseConfigBalancer": {"cluster": "test_cluster"}}]}`,
+			wantErr: "error parsing loadBalancingConfig for policy \"errParseConfigBalancer\"",
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			gotCfg, gotErr := parser.ParseConfig(json.RawMessage(test.input))
+			// Substring match makes this very tightly coupled to the
+			// internalserviceconfig.BalancerConfig error strings. However, it
+			// is important to distinguish the different types of error messages
+			// possible as the parser has a few defined buckets of ways it can
+			// error out.
+			if (gotErr != nil) != (test.wantErr != "") {
+				t.Fatalf("ParseConfig(%v) = %v, wantErr %v", test.input, gotErr, test.wantErr)
+			}
+			if gotErr != nil && !strings.Contains(gotErr.Error(), test.wantErr) {
+				t.Fatalf("ParseConfig(%v) = %v, wantErr %v", test.input, gotErr, test.wantErr)
+			}
+			if test.wantErr != "" {
+				return
+			}
+			if diff := cmp.Diff(gotCfg, test.wantCfg); diff != "" {
+				t.Fatalf("ParseConfig(%v) got unexpected output, diff (-got +want): %v", test.input, diff)
+			}
+		})
+	}
+}
+
+// TestUpdateClientConnState tests the UpdateClientConnState method of the
+// wrr_locality_experimental balancer. This UpdateClientConn operation should
+// take the localities and their weights in the addresses passed in, alongside
+// the endpoint picking policy defined in the Balancer Config and construct a
+// weighted target configuration corresponding to these inputs.
+func (s) TestUpdateClientConnState(t *testing.T) {
+	// Configure the stub balancer defined below as the child policy of
+	// wrrLocalityBalancer.
+	cfgCh := testutils.NewChannel()
+	oldWeightedTargetName := weightedTargetName
+	defer func() {
+		weightedTargetName = oldWeightedTargetName
+	}()
+	weightedTargetName = "fake_weighted_target"
+	stub.Register("fake_weighted_target", stub.BalancerFuncs{
+		ParseConfig: func(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+			var cfg weightedtarget.LBConfig
+			if err := json.Unmarshal(c, &cfg); err != nil {
+				return nil, err
+			}
+			return &cfg, nil
+		},
+		UpdateClientConnState: func(_ *stub.BalancerData, ccs balancer.ClientConnState) error {
+			wtCfg, ok := ccs.BalancerConfig.(*weightedtarget.LBConfig)
+			if !ok {
+				return errors.New("child received config that was not a weighted target config")
+			}
+			defer cfgCh.Send(wtCfg)
+			return nil
+		},
+	})
+
+	builder := balancer.Get(Name)
+	if builder == nil {
+		t.Fatalf("balancer.Get(%q) returned nil", Name)
+	}
+	tcc := testutils.NewBalancerClientConn(t)
+	bal := builder.Build(tcc, balancer.BuildOptions{})
+	defer bal.Close()
+	wrrL := bal.(*wrrLocalityBalancer)
+
+	// Create the addresses with two localities with certain locality weights.
+	// This represents what addresses the wrr_locality balancer will receive in
+	// UpdateClientConnState.
+	ep1 := resolver.Endpoint{Addresses: []resolver.Address{{Addr: "locality-1"}}}
+	ep1 = xdsinternal.SetLocalityIDInEndpoint(ep1, clients.Locality{
+		Region:  "region-1",
+		Zone:    "zone-1",
+		SubZone: "subzone-1",
+	})
+	ep1 = SetAddrInfo(ep1, AddrInfo{LocalityWeight: 2})
+
+	ep2 := resolver.Endpoint{Addresses: []resolver.Address{{Addr: "locality-2"}}}
+	ep2 = xdsinternal.SetLocalityIDInEndpoint(ep2, clients.Locality{
+		Region:  "region-2",
+		Zone:    "zone-2",
+		SubZone: "subzone-2",
+	})
+	ep2 = SetAddrInfo(ep2, AddrInfo{LocalityWeight: 1})
+
+	eps := []resolver.Endpoint{ep1, ep2}
+
+	err := wrrL.UpdateClientConnState(balancer.ClientConnState{
+		BalancerConfig: &LBConfig{
+			ChildPolicy: &internalserviceconfig.BalancerConfig{
+				Name: "round_robin",
+			},
+		},
+		ResolverState: resolver.State{
+			Endpoints: eps,
+		},
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error from UpdateClientConnState: %v", err)
+	}
+
+	// Note that these inline strings declared as the key in Targets built from
+	// Locality ID are not exactly what is shown in the example in the gRFC.
+	// However, this is an implementation detail that does not affect
+	// correctness (confirmed with Java team). The important thing is to get
+	// those three pieces of information region, zone, and subzone down to the
+	// child layer.
+	wantWtCfg := &weightedtarget.LBConfig{
+		Targets: map[string]weightedtarget.Target{
+			"{region=\"region-1\", zone=\"zone-1\", sub_zone=\"subzone-1\"}": {
+				Weight: 2,
+				ChildPolicy: &internalserviceconfig.BalancerConfig{
+					Name: "round_robin",
+				},
+			},
+			"{region=\"region-2\", zone=\"zone-2\", sub_zone=\"subzone-2\"}": {
+				Weight: 1,
+				ChildPolicy: &internalserviceconfig.BalancerConfig{
+					Name: "round_robin",
+				},
+			},
+		},
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	cfg, err := cfgCh.Receive(ctx)
+	if err != nil {
+		t.Fatalf("No signal received from UpdateClientConnState() on the child: %v", err)
+	}
+
+	gotWtCfg, ok := cfg.(*weightedtarget.LBConfig)
+	if !ok {
+		// Shouldn't happen - only sends a config on this channel.
+		t.Fatalf("Unexpected config type: %T", gotWtCfg)
+	}
+
+	if diff := cmp.Diff(gotWtCfg, wantWtCfg); diff != "" {
+		t.Fatalf("Child received unexpected config, diff (-got, +want): %v", diff)
+	}
+}

internal/xds/balancer/wrrlocality/logging.go

@@ -0,0 +1,34 @@
+/*
+ *
+ * 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
+
+import (
+	"fmt"
+
+	"google.golang.org/grpc/grpclog"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+)
+
+const prefix = "[wrrlocality-lb %p] "
+
+var logger = grpclog.Component("xds")
+
+func prefixLogger(p *wrrLocalityBalancer) *internalgrpclog.PrefixLogger {
+	return internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(prefix, p))
+}