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

balancer/autosharding/autosharding.go

@@ -0,0 +1,112 @@
+/*
+ *
+ * Copyright 2026 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 autosharding implements the autosharding load balancing policy.
+package autosharding
+
+import (
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/endpointsharding"
+	iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// Name is the name of the autosharding balancer.
+const Name = "autosharding_experimental"
+
+func init() {
+	balancer.Register(bb{})
+}
+
+// lbConfig is the balancer config for the autosharding balancer.
+type lbConfig struct {
+	serviceconfig.LoadBalancingConfig `json:"-"`
+
+	ChannelFactoryKey        string                  `json:"channelFactoryKey,omitempty"`
+	AutoShardingTarget       string                  `json:"autoshardingTarget,omitempty"`
+	KeyHeaderName            string                  `json:"keyHeaderName,omitempty"`
+	EnableFallback           bool                    `json:"enableFallback,omitempty"`
+	InitialAssignmentTimeout iserviceconfig.Duration `json:"initialAssignmentTimeout,omitempty"`
+}
+
+type bb struct{}
+
+func (bb) Name() string {
+	return Name
+}
+
+func (bb) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	lbConfig := &lbConfig{InitialAssignmentTimeout: iserviceconfig.Duration(60 * time.Second)}
+	if err := json.Unmarshal(s, lbConfig); err != nil {
+		return nil, fmt.Errorf("autosharding: unable to unmarshal LBConfig: %v", err)
+	}
+	if lbConfig.ChannelFactoryKey == "" {
+		return nil, fmt.Errorf("autosharding: channelFactoryKey field is required")
+	}
+	if lbConfig.AutoShardingTarget == "" {
+		return nil, fmt.Errorf("autosharding: autoshardingTarget field is required")
+	}
+	if lbConfig.KeyHeaderName == "" {
+		return nil, fmt.Errorf("autosharding: keyHeaderName field is required")
+	}
+	return lbConfig, nil
+}
+
+func (bb) Build(balancer.ClientConn, balancer.BuildOptions) balancer.Balancer {
+	return &autoshardingBalancer{}
+}
+
+// slice represents a key range and its assigned endpoints.
+//
+//lint:ignore U1000 Struct fields planned for future implementation
+type slice struct {
+	startKey  []byte // Inclusive start key of the key-range
+	endKey    []byte // Exclusive, nil for sentinel/infinity
+	endpoints []int  // Indices into assignment.endpointNames
+}
+
+// assignment represents a complete snapshot of sharding assignments, and is
+// expected to cover the entire key range.
+type assignment struct {
+	slices        []slice  // Sorted by startKey
+	endpointNames []string // Complete list of endpoint names
+	generation    int64
+}
+
+// endpointState represents the state associated with an endpoint in the LB policy.
+//
+//lint:ignore U1000 Struct fields planned for future implementation
+type endpointState struct {
+	index      int                         // Index of the endpoint within the NR update
+	endpoint   resolver.Endpoint           // The actual endpoint returned by the NR
+	childState endpointsharding.ChildState // State as reported by the child policy
+}
+
+// endpointMap maps from endpoint hostname to endpoint state.
+type endpointMap struct {
+	m map[string]*endpointState
+}
+
+type autoshardingBalancer struct {
+	balancer.Balancer
+}

balancer/autosharding/autosharding_test.go

@@ -0,0 +1,144 @@
+/*
+ *
+ * Copyright 2026 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 autosharding
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/internal/grpctest"
+	iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestParseConfig_Success(t *testing.T) {
+	parser := bb{}
+	tests := []struct {
+		name    string
+		input   string
+		wantCfg serviceconfig.LoadBalancingConfig
+	}{
+		{
+			name: "all-fields",
+			input: `{
+				"channelFactoryKey": "factory-key",
+				"autoshardingTarget": "target",
+				"keyHeaderName": "header",
+				"enableFallback": true,
+				"initialAssignmentTimeout": "30s"
+			}`,
+			wantCfg: &lbConfig{
+				ChannelFactoryKey:        "factory-key",
+				AutoShardingTarget:       "target",
+				KeyHeaderName:            "header",
+				EnableFallback:           true,
+				InitialAssignmentTimeout: iserviceconfig.Duration(30 * time.Second),
+			},
+		},
+		{
+			name: "default-timeout",
+			input: `{
+				"channelFactoryKey": "factory-key",
+				"autoshardingTarget": "target",
+				"keyHeaderName": "header"
+			}`,
+			wantCfg: &lbConfig{
+				ChannelFactoryKey:        "factory-key",
+				AutoShardingTarget:       "target",
+				KeyHeaderName:            "header",
+				InitialAssignmentTimeout: iserviceconfig.Duration(60 * time.Second),
+			},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			gotCfg, err := parser.ParseConfig(json.RawMessage(test.input))
+			if err != nil {
+				t.Fatalf("ParseConfig() error = %v, want nil", err)
+			}
+			if diff := cmp.Diff(test.wantCfg, gotCfg); diff != "" {
+				t.Errorf("ParseConfig() config diff (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func (s) TestParseConfig_Failure(t *testing.T) {
+	parser := bb{}
+	tests := []struct {
+		name  string
+		input string
+	}{
+		{
+			name:  "invalid-json",
+			input: "{{invalidjson{{",
+		},
+		{
+			name: "invalid-duration",
+			input: `{
+				"channelFactoryKey": "factory-key",
+				"autoshardingTarget": "target",
+				"keyHeaderName": "header",
+				"initialAssignmentTimeout": "invalid"
+			}`,
+		},
+		{
+			name: "missing-channel-factory-key",
+			input: `{
+				"autoshardingTarget": "target",
+				"keyHeaderName": "header"
+			}`,
+		},
+		{
+			name: "missing-autosharding-target",
+			input: `{
+				"channelFactoryKey": "factory-key",
+				"keyHeaderName": "header"
+			}`,
+		},
+		{
+			name: "missing-key-header-name",
+			input: `{
+				"channelFactoryKey": "factory-key",
+				"autoshardingTarget": "target"
+			}`,
+		},
+		{
+			name:  "empty-config",
+			input: `{}`,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if _, err := parser.ParseConfig(json.RawMessage(test.input)); err == nil {
+				t.Fatalf("ParseConfig() succeeded, want error")
+			}
+		})
+	}
+}

balancer/autosharding/slice_map.go

@@ -0,0 +1,120 @@
+/*
+ *
+ * Copyright 2026 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 autosharding
+
+import (
+	"bytes"
+	"slices"
+)
+
+// sliceMapEntry represents an entry for a key-range in the sliceMap.
+type sliceMapEntry struct {
+	startKey  []byte // Inclusive start key of the key-range
+	endpoints []int  // Indices into list[PickerEndpoint] in the Picker
+}
+
+// sliceMap is an immutable data structure used by the Picker to map request
+// keys to their assigned key-range. Because it is immutable, it allows the
+// Picker to access it without any explicit synchronization with the LB policy.
+//
+// The sliceMap is meant to be used by the Picker in conjunction with a list of
+// pickerEndpoints such that the list can be swapped out, as long as the number
+// and order of endpoints don't change.
+type sliceMap struct {
+	slices       []sliceMapEntry // Sorted by startKey
+	fallbackPool []int           // Indices into list[PickerEndpoint] for all endpoints
+	generation   int64           // Snapshot generation number
+}
+
+// lookup returns the index of the slice covering the given key.
+//
+// Because assignments are pre-validated to have no gaps and cover the full key
+// range, and since slices is sorted by startKey, lookup boils down to a binary
+// search, looking for the insertion point.
+//
+// If the key matches the startKey of a slice, that slice index is returned.
+// Otherwise, the index of the slice immediately preceding the insertion point
+// is returned (i.e., the slice covering the range [startKey, nextStartKey)).
+//
+// Returns -1 when the sliceMap is empty, which is the case before the first
+// assignment is received.
+func (sm *sliceMap) lookup(key []byte) int {
+	if len(sm.slices) == 0 {
+		return -1
+	}
+
+	idx, found := slices.BinarySearchFunc(sm.slices, key, func(e sliceMapEntry, k []byte) int {
+		return bytes.Compare(e.startKey, k)
+	})
+
+	if found {
+		return idx
+	}
+
+	// Key falls in range [slices[idx - 1].startKey, slices[idx].startKey).
+	return idx - 1
+}
+
+// buildSliceMap is used to generate a new sliceMap from the EndpointMap and
+// Assignment when either of them change.
+func buildSliceMap(endpointMap *endpointMap, assignment *assignment) *sliceMap {
+	sm := &sliceMap{}
+
+	// Populate fallbackPool with values [0, 1, 2, ... N-1] where N is the
+	// number of endpoints returned by the name resolver. The only reason to
+	// have this slice is to allow the Picker to share code that picks a random
+	// endpoint from a slice of indices, either from the fallbackPool or from a
+	// sliceMapEntry.
+	sm.fallbackPool = make([]int, len(endpointMap.m))
+	for i := range sm.fallbackPool {
+		sm.fallbackPool[i] = i
+	}
+
+	// If no assignment has been received yet (startup case), return early with
+	// empty slices.
+	if assignment == nil {
+		return sm
+	}
+
+	sm.generation = assignment.generation
+
+	// Build sliceMapEntry for each Slice in the assignment.
+	sm.slices = make([]sliceMapEntry, 0, len(assignment.slices))
+	for _, s := range assignment.slices {
+		entry := sliceMapEntry{
+			startKey:  s.startKey,
+			endpoints: make([]int, 0, len(s.endpoints)),
+		}
+
+		// Populate the endpoints for this slice by looking up the endpoint
+		// names in the EndpointMap. If an endpoint name is not found in the
+		// EndpointMap, it is skipped. This ensures that the sliceMap only
+		// contains valid endpoints that are currently available.
+		for _, idx := range s.endpoints {
+			hostname := assignment.endpointNames[idx]
+			if es, ok := endpointMap.m[hostname]; ok {
+				entry.endpoints = append(entry.endpoints, es.index)
+			}
+		}
+
+		sm.slices = append(sm.slices, entry)
+	}
+
+	return sm
+}

balancer/autosharding/slice_map_test.go

@@ -0,0 +1,216 @@
+/*
+ *
+ * Copyright 2026 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 autosharding
+
+import (
+	"bytes"
+	"fmt"
+	"slices"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func (s) TestSliceMap_Lookup(t *testing.T) {
+	sm := &sliceMap{
+		slices: []sliceMapEntry{
+			{startKey: []byte("a"), endpoints: []int{1}},
+			{startKey: []byte("d"), endpoints: []int{2}},
+			{startKey: []byte("f"), endpoints: []int{3}},
+		},
+	}
+
+	tests := []struct {
+		name string
+		key  []byte
+		want int
+	}{
+		{
+			name: "exact-match-first",
+			key:  []byte("a"),
+			want: 0,
+		},
+		{
+			name: "exact-match-middle",
+			key:  []byte("d"),
+			want: 1,
+		},
+		{
+			name: "exact-match-last",
+			key:  []byte("f"),
+			want: 2,
+		},
+		{
+			name: "between-first-and-middle",
+			key:  []byte("c"),
+			want: 0,
+		},
+		{
+			name: "between-middle-and-last",
+			key:  []byte("e"),
+			want: 1,
+		},
+		{
+			name: "after-last",
+			key:  []byte("g"),
+			want: 2,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := sm.lookup(tc.key); got != tc.want {
+				t.Errorf("lookup(%q) = %d, want %d", tc.key, got, tc.want)
+			}
+		})
+	}
+}
+
+func (s) TestSliceMap_Lookup_Empty(t *testing.T) {
+	sm := &sliceMap{}
+	if got := sm.lookup([]byte("a")); got != -1 {
+		t.Errorf("lookup() on empty map = %d, want -1", got)
+	}
+}
+
+func (s) TestBuildSliceMap(t *testing.T) {
+	// Setup endpointMap with unsorted order in map to verify deterministic
+	// sorting of fallbackPool.
+	epMap := &endpointMap{
+		m: map[string]*endpointState{
+			"hostC": {index: 2},
+			"hostA": {index: 0},
+			"hostB": {index: 1},
+		},
+	}
+
+	tests := []struct {
+		name       string
+		assignment *assignment
+		want       *sliceMap
+	}{
+		{
+			name:       "nil-assignment-startup",
+			assignment: nil,
+			want:       &sliceMap{fallbackPool: []int{0, 1, 2}},
+		},
+		{
+			name: "valid-assignment",
+			assignment: &assignment{
+				endpointNames: []string{"hostA", "hostB", "hostC"},
+				slices: []slice{
+					{startKey: []byte("a"), endpoints: []int{0, 1}}, // hostA, hostB -> indices 0, 1
+					{startKey: []byte("m"), endpoints: []int{1, 2}}, // hostB, hostC -> indices 1, 2
+				},
+				generation: 42,
+			},
+			want: &sliceMap{
+				slices: []sliceMapEntry{
+					{startKey: []byte("a"), endpoints: []int{0, 1}},
+					{startKey: []byte("m"), endpoints: []int{1, 2}},
+				},
+				fallbackPool: []int{0, 1, 2},
+				generation:   42,
+			},
+		},
+		{
+			name: "assignment-with-unknown-host",
+			assignment: &assignment{
+				endpointNames: []string{"hostA", "hostUnknown", "hostC"},
+				slices: []slice{
+					{startKey: []byte("a"), endpoints: []int{0, 1}}, // hostUnknown is skipped
+				},
+				generation: 43,
+			},
+			want: &sliceMap{
+				slices: []sliceMapEntry{
+					{startKey: []byte("a"), endpoints: []int{0}}, // Only hostA (index 0)
+				},
+				fallbackPool: []int{0, 1, 2},
+				generation:   43,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := buildSliceMap(epMap, tc.assignment)
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("buildSliceMap() diff (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func (e sliceMapEntry) Equal(b sliceMapEntry) bool {
+	return bytes.Equal(e.startKey, b.startKey) && slices.Equal(e.endpoints, b.endpoints)
+}
+
+func (sm *sliceMap) Equal(b *sliceMap) bool {
+	if sm == nil || b == nil {
+		return sm == b
+	}
+	fallbackPoolEqual := slices.Equal(sm.fallbackPool, b.fallbackPool)
+	slicesEqual := slices.EqualFunc(sm.slices, b.slices, func(x, y sliceMapEntry) bool { return x.Equal(y) })
+	return sm.generation == b.generation && fallbackPoolEqual && slicesEqual
+}
+
+func BenchmarkSliceMap_Lookup(b *testing.B) {
+	for _, numSlices := range []int{1, 10, 100, 1000, 10000} {
+		for _, keySize := range []int{16, 32, 64, 128, 256, 512} {
+			b.Run(fmt.Sprintf("slices_%d_keySize_%d", numSlices, keySize), func(b *testing.B) {
+				sm := &sliceMap{
+					slices: make([]sliceMapEntry, numSlices),
+				}
+				for i := 0; i < numSlices; i++ {
+					// Generate lexicographically sorted keys to serve as slice
+					// boundaries.  We multiply by 1000 to create ranges (gaps) between
+					// successive slices (e.g., Slice 0 covers [0, 1000), Slice 1 covers
+					// [1000, 2000)). This allows us to test lookups that fall
+					// *mid-slice*, rather than just exact matches.
+					//
+					// We use zero-padding on the left via "%0*d" (where * is the keySize
+					// width).  This fills the key to its full length (e.g., 512 bytes)
+					// with leading zeros.  Scanning long identical prefixes forces
+					// bytes.Compare to traverse the full length, simulating a
+					// conservative, worst-case latency scenario.
+					key := fmt.Appendf(nil, "%0*d", keySize, i*1000)
+					sm.slices[i] = sliceMapEntry{startKey: key, endpoints: []int{i}}
+				}
+
+				// Pre-generate a pool of lookup keys. This helps avoid
+				// measuring string formatting overhead inside the timer loop.
+				lookupKeys := make([][]byte, 10000)
+				for i := 0; i < 10000; i++ {
+					// Distribute the 10,000 lookup keys proportionally across the entire
+					// synthetic keyspace [0, numSlices * 1000). This ensures a good mix
+					// of boundary hits and interior hits.
+					val := i * (numSlices * 1000) / 10000
+					lookupKeys[i] = fmt.Appendf(nil, "%0*d", keySize, val)
+				}
+
+				var i int
+				for b.Loop() {
+					sm.lookup(lookupKeys[i%10000])
+					i++
+				}
+			})
+		}
+	}
+}

balancer/balancer.go

@@ -0,0 +1,394 @@
+/*
+ *
+ * Copyright 2017 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 balancer defines APIs for load balancing in gRPC.
+// All APIs in this package are experimental.
+package balancer
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"net"
+	"strings"
+
+	"google.golang.org/grpc/channelz"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials"
+	estats "google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/envconfig"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+var (
+	// m is a map from name to balancer builder.
+	m = make(map[string]Builder)
+
+	logger = grpclog.Component("balancer")
+)
+
+// Register registers the balancer builder to the balancer map. b.Name
+// will be used as the name registered with this builder. If the Builder
+// implements ConfigParser, ParseConfig will be called when new service
+// configs are received by the resolver, and the result will be provided to the
+// Balancer in UpdateClientConnState.
+//
+// NOTE: this function must only be called during initialization time (i.e. in
+// an init() function), and is not thread-safe. If multiple Balancers are
+// registered with the same name, the one registered last will take effect.
+func Register(b Builder) {
+	name := b.Name()
+	if !envconfig.CaseSensitiveBalancerRegistries {
+		name = strings.ToLower(name)
+		if name != b.Name() {
+			logger.Warningf("Balancer registered with name %q. grpc-go has switched to case sensitive balancer registries. GRPC_GO_EXPERIMENTAL_CASE_SENSITIVE_BALANCER_REGISTRIES env variable will be removed in release v1.82.0", b.Name())
+		}
+	}
+	m[name] = b
+}
+
+// unregisterForTesting deletes the balancer with the given name from the
+// balancer map.
+//
+// This function is not thread-safe.
+func unregisterForTesting(name string) {
+	delete(m, name)
+}
+
+func init() {
+	internal.BalancerUnregister = unregisterForTesting
+}
+
+// Get returns the resolver builder registered with the given name.
+// Note that the compare is done in a case-sensitive fashion.
+// If no builder is register with the name, nil will be returned.
+func Get(name string) Builder {
+	if !envconfig.CaseSensitiveBalancerRegistries {
+		lowerName := strings.ToLower(name)
+		if lowerName != name {
+			logger.Warningf("Balancer retrieved for name %q. grpc-go has switched to case sensitive balancer registries. GRPC_GO_EXPERIMENTAL_CASE_SENSITIVE_BALANCER_REGISTRIES env variable will be removed in release v1.82.0", name)
+		}
+		name = lowerName
+	}
+	if b, ok := m[name]; ok {
+		return b
+	}
+	return nil
+}
+
+// NewSubConnOptions contains options to create new SubConn.
+type NewSubConnOptions struct {
+	// CredsBundle is the credentials bundle that will be used in the created
+	// SubConn. If it's nil, the original creds from grpc DialOptions will be
+	// used.
+	//
+	// Deprecated: Use the Attributes field in resolver.Address to pass
+	// arbitrary data to the credential handshaker.
+	CredsBundle credentials.Bundle
+	// HealthCheckEnabled indicates whether health check service should be
+	// enabled on this SubConn
+	HealthCheckEnabled bool
+	// StateListener is called when the state of the subconn changes.  If nil,
+	// Balancer.UpdateSubConnState will be called instead.  Will never be
+	// invoked until after Connect() is called on the SubConn created with
+	// these options.
+	StateListener func(SubConnState)
+}
+
+// State contains the balancer's state relevant to the gRPC ClientConn.
+type State struct {
+	// State contains the connectivity state of the balancer, which is used to
+	// determine the state of the ClientConn.
+	ConnectivityState connectivity.State
+	// Picker is used to choose connections (SubConns) for RPCs.
+	Picker Picker
+}
+
+// ClientConn represents a gRPC ClientConn.
+//
+// This interface is to be implemented by gRPC. Users should not need a
+// brand new implementation of this interface. For the situations like
+// testing, the new implementation should embed this interface. This allows
+// gRPC to add new methods to this interface.
+//
+// NOTICE: This interface is intended to be implemented by gRPC, or intercepted
+// by custom load balancing polices.  Users should not need their own complete
+// implementation of this interface -- they should always delegate to a
+// ClientConn passed to Builder.Build() by embedding it in their
+// implementations. An embedded ClientConn must never be nil, or runtime panics
+// will occur.
+type ClientConn interface {
+	// NewSubConn is called by balancer to create a new SubConn.
+	// It doesn't block and wait for the connections to be established.
+	// Behaviors of the SubConn can be controlled by options.
+	//
+	// Deprecated: please be aware that in a future version, SubConns will only
+	// support one address per SubConn.
+	NewSubConn([]resolver.Address, NewSubConnOptions) (SubConn, error)
+	// RemoveSubConn removes the SubConn from ClientConn.
+	// The SubConn will be shutdown.
+	//
+	// Deprecated: use SubConn.Shutdown instead.
+	RemoveSubConn(SubConn)
+	// UpdateAddresses updates the addresses used in the passed in SubConn.
+	// gRPC checks if the currently connected address is still in the new list.
+	// If so, the connection will be kept. Else, the connection will be
+	// gracefully closed, and a new connection will be created.
+	//
+	// This may trigger a state transition for the SubConn.
+	//
+	// Deprecated: this method will be removed.  Create new SubConns for new
+	// addresses instead.
+	UpdateAddresses(SubConn, []resolver.Address)
+
+	// UpdateState notifies gRPC that the balancer's internal state has
+	// changed.
+	//
+	// gRPC will update the connectivity state of the ClientConn, and will call
+	// Pick on the new Picker to pick new SubConns.
+	UpdateState(State)
+
+	// ResolveNow is called by balancer to notify gRPC to do a name resolving.
+	ResolveNow(resolver.ResolveNowOptions)
+
+	// Target returns the dial target for this ClientConn.
+	//
+	// Deprecated: Use the Target field in the BuildOptions instead.
+	Target() string
+
+	// MetricsRecorder provides the metrics recorder that balancers can use to
+	// record metrics. Balancer implementations which do not register metrics on
+	// metrics registry and record on them can ignore this method. The returned
+	// MetricsRecorder is guaranteed to never be nil.
+	MetricsRecorder() estats.MetricsRecorder
+
+	// EnforceClientConnEmbedding is included to force implementers to embed
+	// another implementation of this interface, allowing gRPC to add methods
+	// without breaking users.
+	internal.EnforceClientConnEmbedding
+}
+
+// BuildOptions contains additional information for Build.
+type BuildOptions struct {
+	// DialCreds is the transport credentials to use when communicating with a
+	// remote load balancer server. Balancer implementations which do not
+	// communicate with a remote load balancer server can ignore this field.
+	DialCreds credentials.TransportCredentials
+	// CredsBundle is the credentials bundle to use when communicating with a
+	// remote load balancer server. Balancer implementations which do not
+	// communicate with a remote load balancer server can ignore this field.
+	CredsBundle credentials.Bundle
+	// Dialer is the custom dialer to use when communicating with a remote load
+	// balancer server. Balancer implementations which do not communicate with a
+	// remote load balancer server can ignore this field.
+	Dialer func(context.Context, string) (net.Conn, error)
+	// Authority is the server name to use as part of the authentication
+	// handshake when communicating with a remote load balancer server. Balancer
+	// implementations which do not communicate with a remote load balancer
+	// server can ignore this field.
+	Authority string
+	// ChannelzParent is the parent ClientConn's channelz channel.
+	ChannelzParent channelz.Identifier
+	// CustomUserAgent is the custom user agent set on the parent ClientConn.
+	// The balancer should set the same custom user agent if it creates a
+	// ClientConn.
+	CustomUserAgent string
+	// Target contains the parsed address info of the dial target. It is the
+	// same resolver.Target as passed to the resolver. See the documentation for
+	// the resolver.Target type for details about what it contains.
+	Target resolver.Target
+}
+
+// Builder creates a balancer.
+type Builder interface {
+	// Build creates a new balancer with the ClientConn.
+	Build(cc ClientConn, opts BuildOptions) Balancer
+	// Name returns the name of balancers built by this builder.
+	// It will be used to pick balancers (for example in service config).
+	Name() string
+}
+
+// ConfigParser parses load balancer configs.
+type ConfigParser interface {
+	// ParseConfig parses the JSON load balancer config provided into an
+	// internal form or returns an error if the config is invalid.  For future
+	// compatibility reasons, unknown fields in the config should be ignored.
+	ParseConfig(LoadBalancingConfigJSON json.RawMessage) (serviceconfig.LoadBalancingConfig, error)
+}
+
+// PickInfo contains additional information for the Pick operation.
+type PickInfo struct {
+	// FullMethodName is the method name that NewClientStream() is called
+	// with. The canonical format is /service/Method.
+	FullMethodName string
+	// Ctx is the RPC's context, and may contain relevant RPC-level information
+	// like the outgoing header metadata.
+	Ctx context.Context
+}
+
+// DoneInfo contains additional information for done.
+type DoneInfo struct {
+	// Err is the rpc error the RPC finished with. It could be nil.
+	Err error
+	// Trailer contains the metadata from the RPC's trailer, if present.
+	Trailer metadata.MD
+	// BytesSent indicates if any bytes have been sent to the server.
+	BytesSent bool
+	// BytesReceived indicates if any byte has been received from the server.
+	BytesReceived bool
+	// ServerLoad is the load received from server. It's usually sent as part of
+	// trailing metadata.
+	//
+	// The only supported type now is *orca_v3.LoadReport.
+	ServerLoad any
+}
+
+var (
+	// ErrNoSubConnAvailable indicates no SubConn is available for pick().
+	// gRPC will block the RPC until a new picker is available via UpdateState().
+	ErrNoSubConnAvailable = errors.New("no SubConn is available")
+	// ErrTransientFailure indicates all SubConns are in TransientFailure.
+	// WaitForReady RPCs will block, non-WaitForReady RPCs will fail.
+	//
+	// Deprecated: return an appropriate error based on the last resolution or
+	// connection attempt instead.  The behavior is the same for any non-gRPC
+	// status error.
+	ErrTransientFailure = errors.New("all SubConns are in TransientFailure")
+)
+
+// PickResult contains information related to a connection chosen for an RPC.
+type PickResult struct {
+	// SubConn is the connection to use for this pick, if its state is Ready.
+	// If the state is not Ready, gRPC will block the RPC until a new Picker is
+	// provided by the balancer (using ClientConn.UpdateState).  The SubConn
+	// must be one returned by ClientConn.NewSubConn.
+	SubConn SubConn
+
+	// Done is called when the RPC is completed.  If the SubConn is not ready,
+	// this will be called with a nil parameter.  If the SubConn is not a valid
+	// type, Done may not be called.  May be nil if the balancer does not wish
+	// to be notified when the RPC completes.
+	Done func(DoneInfo)
+
+	// Metadata provides a way for LB policies to inject arbitrary per-call
+	// metadata. Any metadata returned here will be merged with existing
+	// metadata added by the client application.
+	//
+	// LB policies with child policies are responsible for propagating metadata
+	// injected by their children to the ClientConn, as part of Pick().
+	Metadata metadata.MD
+}
+
+// TransientFailureError returns e.  It exists for backward compatibility and
+// will be deleted soon.
+//
+// Deprecated: no longer necessary, picker errors are treated this way by
+// default.
+func TransientFailureError(e error) error { return e }
+
+// Picker is used by gRPC to pick a SubConn to send an RPC.
+// Balancer is expected to generate a new picker from its snapshot every time its
+// internal state has changed.
+//
+// The pickers used by gRPC can be updated by ClientConn.UpdateState().
+type Picker interface {
+	// Pick returns the connection to use for this RPC and related information.
+	//
+	// Pick should not block.  If the balancer needs to do I/O or any blocking
+	// or time-consuming work to service this call, it should return
+	// ErrNoSubConnAvailable, and the Pick call will be repeated by gRPC when
+	// the Picker is updated (using ClientConn.UpdateState).
+	//
+	// If an error is returned:
+	//
+	// - If the error is ErrNoSubConnAvailable, gRPC will block until a new
+	//   Picker is provided by the balancer (using ClientConn.UpdateState).
+	//
+	// - If the error is a status error (implemented by the grpc/status
+	//   package), gRPC will terminate the RPC with the code and message
+	//   provided.
+	//
+	// - For all other errors, wait for ready RPCs will wait, but non-wait for
+	//   ready RPCs will be terminated with this error's Error() string and
+	//   status code Unavailable.
+	Pick(info PickInfo) (PickResult, error)
+}
+
+// Balancer takes input from gRPC, manages SubConns, and collects and aggregates
+// the connectivity states.
+//
+// It also generates and updates the Picker used by gRPC to pick SubConns for RPCs.
+//
+// UpdateClientConnState, ResolverError, UpdateSubConnState, and Close are
+// guaranteed to be called synchronously from the same goroutine.  There's no
+// guarantee on picker.Pick, it may be called anytime.
+type Balancer interface {
+	// UpdateClientConnState is called by gRPC when the state of the ClientConn
+	// changes.  If the error returned is ErrBadResolverState, the ClientConn
+	// will begin calling ResolveNow on the active name resolver with
+	// exponential backoff until a subsequent call to UpdateClientConnState
+	// returns a nil error.  Any other errors are currently ignored.
+	UpdateClientConnState(ClientConnState) error
+	// ResolverError is called by gRPC when the name resolver reports an error.
+	ResolverError(error)
+	// UpdateSubConnState is called by gRPC when the state of a SubConn
+	// changes.
+	//
+	// Deprecated: Use NewSubConnOptions.StateListener when creating the
+	// SubConn instead.
+	UpdateSubConnState(SubConn, SubConnState)
+	// Close closes the balancer. The balancer is not currently required to
+	// call SubConn.Shutdown for its existing SubConns; however, this will be
+	// required in a future release, so it is recommended.
+	Close()
+	// ExitIdle instructs the LB policy to reconnect to backends / exit the
+	// IDLE state, if appropriate and possible.  Note that SubConns that enter
+	// the IDLE state will not reconnect until SubConn.Connect is called.
+	ExitIdle()
+}
+
+// ExitIdler is an optional interface for balancers to implement.  If
+// implemented, ExitIdle will be called when ClientConn.Connect is called, if
+// the ClientConn is idle.  If unimplemented, ClientConn.Connect will cause
+// all SubConns to connect.
+//
+// Deprecated: All balancers must implement this interface. This interface will
+// be removed in a future release.
+type ExitIdler interface {
+	// ExitIdle instructs the LB policy to reconnect to backends / exit the
+	// IDLE state, if appropriate and possible.  Note that SubConns that enter
+	// the IDLE state will not reconnect until SubConn.Connect is called.
+	ExitIdle()
+}
+
+// ClientConnState describes the state of a ClientConn relevant to the
+// balancer.
+type ClientConnState struct {
+	ResolverState resolver.State
+	// The parsed load balancing configuration returned by the builder's
+	// ParseConfig method, if implemented.
+	BalancerConfig serviceconfig.LoadBalancingConfig
+}
+
+// ErrBadResolverState may be returned by UpdateClientConnState to indicate a
+// problem with the provided name resolver data.
+var ErrBadResolverState = errors.New("bad resolver state")

balancer/base/balancer.go

@@ -0,0 +1,260 @@
+/*
+ *
+ * Copyright 2017 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 base
+
+import (
+	"errors"
+	"fmt"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/resolver"
+)
+
+var logger = grpclog.Component("balancer")
+
+type baseBuilder struct {
+	name          string
+	pickerBuilder PickerBuilder
+	config        Config
+}
+
+func (bb *baseBuilder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer {
+	bal := &baseBalancer{
+		cc:            cc,
+		pickerBuilder: bb.pickerBuilder,
+
+		subConns: resolver.NewAddressMapV2[balancer.SubConn](),
+		scStates: make(map[balancer.SubConn]connectivity.State),
+		csEvltr:  &balancer.ConnectivityStateEvaluator{},
+		config:   bb.config,
+		state:    connectivity.Connecting,
+	}
+	// Initialize picker to a picker that always returns
+	// ErrNoSubConnAvailable, because when state of a SubConn changes, we
+	// may call UpdateState with this picker.
+	bal.picker = NewErrPicker(balancer.ErrNoSubConnAvailable)
+	return bal
+}
+
+func (bb *baseBuilder) Name() string {
+	return bb.name
+}
+
+type baseBalancer struct {
+	cc            balancer.ClientConn
+	pickerBuilder PickerBuilder
+
+	csEvltr *balancer.ConnectivityStateEvaluator
+	state   connectivity.State
+
+	subConns *resolver.AddressMapV2[balancer.SubConn]
+	scStates map[balancer.SubConn]connectivity.State
+	picker   balancer.Picker
+	config   Config
+
+	resolverErr error // the last error reported by the resolver; cleared on successful resolution
+	connErr     error // the last connection error; cleared upon leaving TransientFailure
+}
+
+func (b *baseBalancer) ResolverError(err error) {
+	b.resolverErr = err
+	if b.subConns.Len() == 0 {
+		b.state = connectivity.TransientFailure
+	}
+
+	if b.state != connectivity.TransientFailure {
+		// The picker will not change since the balancer does not currently
+		// report an error.
+		return
+	}
+	b.regeneratePicker()
+	b.cc.UpdateState(balancer.State{
+		ConnectivityState: b.state,
+		Picker:            b.picker,
+	})
+}
+
+func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
+	// TODO: handle s.ResolverState.ServiceConfig?
+	if logger.V(2) {
+		logger.Info("base.baseBalancer: got new ClientConn state: ", s)
+	}
+	// Successful resolution; clear resolver error and ensure we return nil.
+	b.resolverErr = nil
+	// addrsSet is the set converted from addrs, it's used for quick lookup of an address.
+	addrsSet := resolver.NewAddressMapV2[any]()
+	for _, a := range s.ResolverState.Addresses {
+		addrsSet.Set(a, nil)
+		if _, ok := b.subConns.Get(a); !ok {
+			// a is a new address (not existing in b.subConns).
+			var sc balancer.SubConn
+			opts := balancer.NewSubConnOptions{
+				HealthCheckEnabled: b.config.HealthCheck,
+				StateListener:      func(scs balancer.SubConnState) { b.updateSubConnState(sc, scs) },
+			}
+			sc, err := b.cc.NewSubConn([]resolver.Address{a}, opts)
+			if err != nil {
+				logger.Warningf("base.baseBalancer: failed to create new SubConn: %v", err)
+				continue
+			}
+			b.subConns.Set(a, sc)
+			b.scStates[sc] = connectivity.Idle
+			b.csEvltr.RecordTransition(connectivity.Shutdown, connectivity.Idle)
+			sc.Connect()
+		}
+	}
+	for a, sc := range b.subConns.All() {
+		// a was removed by resolver.
+		if _, ok := addrsSet.Get(a); !ok {
+			sc.Shutdown()
+			b.subConns.Delete(a)
+			// Keep the state of this sc in b.scStates until sc's state becomes Shutdown.
+			// The entry will be deleted in updateSubConnState.
+		}
+	}
+	// If resolver state contains no addresses, return an error so ClientConn
+	// will trigger re-resolve. Also records this as a resolver error, so when
+	// the overall state turns transient failure, the error message will have
+	// the zero address information.
+	if len(s.ResolverState.Addresses) == 0 {
+		b.ResolverError(errors.New("produced zero addresses"))
+		return balancer.ErrBadResolverState
+	}
+
+	b.regeneratePicker()
+	b.cc.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
+	return nil
+}
+
+// mergeErrors builds an error from the last connection error and the last
+// resolver error.  Must only be called if b.state is TransientFailure.
+func (b *baseBalancer) mergeErrors() error {
+	// connErr must always be non-nil unless there are no SubConns, in which
+	// case resolverErr must be non-nil.
+	if b.connErr == nil {
+		return fmt.Errorf("last resolver error: %v", b.resolverErr)
+	}
+	if b.resolverErr == nil {
+		return fmt.Errorf("last connection error: %v", b.connErr)
+	}
+	return fmt.Errorf("last connection error: %v; last resolver error: %v", b.connErr, b.resolverErr)
+}
+
+// regeneratePicker takes a snapshot of the balancer, and generates a picker
+// from it. The picker is
+//   - errPicker if the balancer is in TransientFailure,
+//   - built by the pickerBuilder with all READY SubConns otherwise.
+func (b *baseBalancer) regeneratePicker() {
+	if b.state == connectivity.TransientFailure {
+		b.picker = NewErrPicker(b.mergeErrors())
+		return
+	}
+	readySCs := make(map[balancer.SubConn]SubConnInfo)
+
+	// Filter out all ready SCs from full subConn map.
+	for addr, sc := range b.subConns.All() {
+		if st, ok := b.scStates[sc]; ok && st == connectivity.Ready {
+			readySCs[sc] = SubConnInfo{Address: addr}
+		}
+	}
+	b.picker = b.pickerBuilder.Build(PickerBuildInfo{ReadySCs: readySCs})
+}
+
+// UpdateSubConnState is a nop because a StateListener is always set in NewSubConn.
+func (b *baseBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	logger.Errorf("base.baseBalancer: UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+func (b *baseBalancer) updateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	s := state.ConnectivityState
+	if logger.V(2) {
+		logger.Infof("base.baseBalancer: handle SubConn state change: %p, %v", sc, s)
+	}
+	oldS, ok := b.scStates[sc]
+	if !ok {
+		if logger.V(2) {
+			logger.Infof("base.baseBalancer: got state changes for an unknown SubConn: %p, %v", sc, s)
+		}
+		return
+	}
+	if oldS == connectivity.TransientFailure &&
+		(s == connectivity.Connecting || s == connectivity.Idle) {
+		// Once a subconn enters TRANSIENT_FAILURE, ignore subsequent IDLE or
+		// CONNECTING transitions to prevent the aggregated state from being
+		// always CONNECTING when many backends exist but are all down.
+		if s == connectivity.Idle {
+			sc.Connect()
+		}
+		return
+	}
+	b.scStates[sc] = s
+	switch s {
+	case connectivity.Idle:
+		sc.Connect()
+	case connectivity.Shutdown:
+		// When an address was removed by resolver, b called Shutdown but kept
+		// the sc's state in scStates. Remove state for this sc here.
+		delete(b.scStates, sc)
+	case connectivity.TransientFailure:
+		// Save error to be reported via picker.
+		b.connErr = state.ConnectionError
+	}
+
+	b.state = b.csEvltr.RecordTransition(oldS, s)
+
+	// Regenerate picker when one of the following happens:
+	//  - this sc entered or left ready
+	//  - the aggregated state of balancer is TransientFailure
+	//    (may need to update error message)
+	if (s == connectivity.Ready) != (oldS == connectivity.Ready) ||
+		b.state == connectivity.TransientFailure {
+		b.regeneratePicker()
+	}
+	b.cc.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
+}
+
+// Close is a nop because base balancer doesn't have internal state to clean up,
+// and it doesn't need to call Shutdown for the SubConns.
+func (b *baseBalancer) Close() {
+}
+
+// ExitIdle is a nop because the base balancer attempts to stay connected to
+// all SubConns at all times.
+func (b *baseBalancer) ExitIdle() {
+}
+
+// NewErrPicker returns a Picker that always returns err on Pick().
+func NewErrPicker(err error) balancer.Picker {
+	return &errPicker{err: err}
+}
+
+// NewErrPickerV2 is temporarily defined for backward compatibility reasons.
+//
+// Deprecated: use NewErrPicker instead.
+var NewErrPickerV2 = NewErrPicker
+
+type errPicker struct {
+	err error // Pick() always returns this err.
+}
+
+func (p *errPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	return balancer.PickResult{}, p.err
+}

balancer/base/balancer_test.go

@@ -0,0 +1,122 @@
+/*
+ *
+ * Copyright 2020 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 base
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/resolver"
+)
+
+type testClientConn struct {
+	balancer.ClientConn
+	newSubConn func([]resolver.Address, balancer.NewSubConnOptions) (balancer.SubConn, error)
+}
+
+func (c *testClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	return c.newSubConn(addrs, opts)
+}
+
+func (c *testClientConn) UpdateState(balancer.State) {}
+
+type testSubConn struct {
+	balancer.SubConn
+	updateState func(balancer.SubConnState)
+}
+
+func (sc *testSubConn) UpdateAddresses([]resolver.Address) {}
+
+func (sc *testSubConn) Connect() {}
+
+func (sc *testSubConn) Shutdown() {}
+
+func (sc *testSubConn) GetOrBuildProducer(balancer.ProducerBuilder) (balancer.Producer, func()) {
+	return nil, nil
+}
+
+// RegisterHealthListener is a no-op.
+func (*testSubConn) RegisterHealthListener(func(balancer.SubConnState)) {}
+
+// testPickBuilder creates balancer.Picker for test.
+type testPickBuilder struct {
+	validate func(info PickerBuildInfo)
+}
+
+func (p *testPickBuilder) Build(info PickerBuildInfo) balancer.Picker {
+	p.validate(info)
+	return nil
+}
+
+func TestBaseBalancerReserveAttributes(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	validated := make(chan struct{}, 1)
+	v := func(info PickerBuildInfo) {
+		defer func() { validated <- struct{}{} }()
+		for _, sc := range info.ReadySCs {
+			if sc.Address.Addr == "1.1.1.1" {
+				if sc.Address.Attributes == nil {
+					t.Errorf("in picker.validate, got address %+v with nil attributes, want not nil", sc.Address)
+				}
+				foo, ok := sc.Address.Attributes.Value("foo").(string)
+				if !ok || foo != "2233niang" {
+					t.Errorf("in picker.validate, got address[1.1.1.1] with invalid attributes value %v, want 2233niang", sc.Address.Attributes.Value("foo"))
+				}
+			} else if sc.Address.Addr == "2.2.2.2" {
+				if sc.Address.Attributes != nil {
+					t.Error("in b.subConns, got address[2.2.2.2] with not nil attributes, want nil")
+				}
+			}
+		}
+	}
+	pickBuilder := &testPickBuilder{validate: v}
+	b := (&baseBuilder{pickerBuilder: pickBuilder}).Build(&testClientConn{
+		newSubConn: func(_ []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+			return &testSubConn{updateState: opts.StateListener}, nil
+		},
+	}, balancer.BuildOptions{}).(*baseBalancer)
+
+	b.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState: resolver.State{
+			Addresses: []resolver.Address{
+				{Addr: "1.1.1.1", Attributes: attributes.New("foo", "2233niang")},
+				{Addr: "2.2.2.2", Attributes: nil},
+			},
+		},
+	})
+	select {
+	case <-validated:
+	case <-ctx.Done():
+		t.Fatalf("timed out waiting for UpdateClientConnState to call picker.Build")
+	}
+
+	for sc := range b.scStates {
+		sc.(*testSubConn).updateState(balancer.SubConnState{ConnectivityState: connectivity.Ready, ConnectionError: nil})
+		select {
+		case <-validated:
+		case <-ctx.Done():
+			t.Fatalf("timed out waiting for UpdateClientConnState to call picker.Build")
+		}
+	}
+}

balancer/base/base.go

@@ -0,0 +1,71 @@
+/*
+ *
+ * Copyright 2017 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 base defines a balancer base that can be used to build balancers with
+// different picking algorithms.
+//
+// The base balancer creates a new SubConn for each resolved address. The
+// provided picker will only be notified about READY SubConns.
+//
+// This package is the base of round_robin balancer, its purpose is to be used
+// to build round_robin like balancers with complex picking algorithms.
+// Balancers with more complicated logic should try to implement a balancer
+// builder from scratch.
+//
+// All APIs in this package are experimental.
+package base
+
+import (
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/resolver"
+)
+
+// PickerBuilder creates balancer.Picker.
+type PickerBuilder interface {
+	// Build returns a picker that will be used by gRPC to pick a SubConn.
+	Build(info PickerBuildInfo) balancer.Picker
+}
+
+// PickerBuildInfo contains information needed by the picker builder to
+// construct a picker.
+type PickerBuildInfo struct {
+	// ReadySCs is a map from all ready SubConns to the Addresses used to
+	// create them.
+	ReadySCs map[balancer.SubConn]SubConnInfo
+}
+
+// SubConnInfo contains information about a SubConn created by the base
+// balancer.
+type SubConnInfo struct {
+	Address resolver.Address // the address used to create this SubConn
+}
+
+// Config contains the config info about the base balancer builder.
+type Config struct {
+	// HealthCheck indicates whether health checking should be enabled for this specific balancer.
+	HealthCheck bool
+}
+
+// NewBalancerBuilder returns a base balancer builder configured by the provided config.
+func NewBalancerBuilder(name string, pb PickerBuilder, config Config) balancer.Builder {
+	return &baseBuilder{
+		name:          name,
+		pickerBuilder: pb,
+		config:        config,
+	}
+}

balancer/conn_state_evaluator.go

@@ -0,0 +1,74 @@
+/*
+ *
+ * Copyright 2022 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 balancer
+
+import "google.golang.org/grpc/connectivity"
+
+// ConnectivityStateEvaluator takes the connectivity states of multiple SubConns
+// and returns one aggregated connectivity state.
+//
+// It's not thread safe.
+type ConnectivityStateEvaluator struct {
+	numReady            uint64 // Number of addrConns in ready state.
+	numConnecting       uint64 // Number of addrConns in connecting state.
+	numTransientFailure uint64 // Number of addrConns in transient failure state.
+	numIdle             uint64 // Number of addrConns in idle state.
+}
+
+// RecordTransition records state change happening in subConn and based on that
+// it evaluates what aggregated state should be.
+//
+//   - If at least one SubConn in Ready, the aggregated state is Ready;
+//   - Else if at least one SubConn in Connecting, the aggregated state is Connecting;
+//   - Else if at least one SubConn is Idle, the aggregated state is Idle;
+//   - Else if at least one SubConn is TransientFailure (or there are no SubConns), the aggregated state is Transient Failure.
+//
+// Shutdown is not considered.
+func (cse *ConnectivityStateEvaluator) RecordTransition(oldState, newState connectivity.State) connectivity.State {
+	// Update counters.
+	for idx, state := range []connectivity.State{oldState, newState} {
+		updateVal := 2*uint64(idx) - 1 // -1 for oldState and +1 for new.
+		switch state {
+		case connectivity.Ready:
+			cse.numReady += updateVal
+		case connectivity.Connecting:
+			cse.numConnecting += updateVal
+		case connectivity.TransientFailure:
+			cse.numTransientFailure += updateVal
+		case connectivity.Idle:
+			cse.numIdle += updateVal
+		}
+	}
+	return cse.CurrentState()
+}
+
+// CurrentState returns the current aggregate conn state by evaluating the counters
+func (cse *ConnectivityStateEvaluator) CurrentState() connectivity.State {
+	// Evaluate.
+	if cse.numReady > 0 {
+		return connectivity.Ready
+	}
+	if cse.numConnecting > 0 {
+		return connectivity.Connecting
+	}
+	if cse.numIdle > 0 {
+		return connectivity.Idle
+	}
+	return connectivity.TransientFailure
+}

balancer/conn_state_evaluator_test.go

@@ -0,0 +1,245 @@
+/*
+ *
+ * Copyright 2022 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 balancer
+
+import (
+	"testing"
+
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// TestRecordTransition_FirstStateChange tests the first call to
+// RecordTransition where the `oldState` is usually set to `Shutdown` (a state
+// that the ConnectivityStateEvaluator is set to ignore).
+func (s) TestRecordTransition_FirstStateChange(t *testing.T) {
+	tests := []struct {
+		newState  connectivity.State
+		wantState connectivity.State
+	}{
+		{
+			newState:  connectivity.Idle,
+			wantState: connectivity.Idle,
+		},
+		{
+			newState:  connectivity.Connecting,
+			wantState: connectivity.Connecting,
+		},
+		{
+			newState:  connectivity.Ready,
+			wantState: connectivity.Ready,
+		},
+		{
+			newState:  connectivity.TransientFailure,
+			wantState: connectivity.TransientFailure,
+		},
+		{
+			newState:  connectivity.Shutdown,
+			wantState: connectivity.TransientFailure,
+		},
+	}
+	for _, test := range tests {
+		cse := &ConnectivityStateEvaluator{}
+		if gotState := cse.RecordTransition(connectivity.Shutdown, test.newState); gotState != test.wantState {
+			t.Fatalf("RecordTransition(%v, %v) = %v, want %v", connectivity.Shutdown, test.newState, gotState, test.wantState)
+		}
+	}
+}
+
+// TestRecordTransition_SameState tests the scenario where state transitions to
+// the same state are recorded multiple times.
+func (s) TestRecordTransition_SameState(t *testing.T) {
+	tests := []struct {
+		newState  connectivity.State
+		wantState connectivity.State
+	}{
+		{
+			newState:  connectivity.Idle,
+			wantState: connectivity.Idle,
+		},
+		{
+			newState:  connectivity.Connecting,
+			wantState: connectivity.Connecting,
+		},
+		{
+			newState:  connectivity.Ready,
+			wantState: connectivity.Ready,
+		},
+		{
+			newState:  connectivity.TransientFailure,
+			wantState: connectivity.TransientFailure,
+		},
+		{
+			newState:  connectivity.Shutdown,
+			wantState: connectivity.TransientFailure,
+		},
+	}
+	const numStateChanges = 5
+	for _, test := range tests {
+		cse := &ConnectivityStateEvaluator{}
+		var prevState, gotState connectivity.State
+		prevState = connectivity.Shutdown
+		for i := 0; i < numStateChanges; i++ {
+			gotState = cse.RecordTransition(prevState, test.newState)
+			prevState = test.newState
+		}
+		if gotState != test.wantState {
+			t.Fatalf("RecordTransition() = %v, want %v", gotState, test.wantState)
+		}
+	}
+}
+
+// TestRecordTransition_SingleSubConn_DifferentStates tests some common
+// connectivity state change scenarios, on a single subConn.
+func (s) TestRecordTransition_SingleSubConn_DifferentStates(t *testing.T) {
+	tests := []struct {
+		name      string
+		states    []connectivity.State
+		wantState connectivity.State
+	}{
+		{
+			name:      "regular transition to ready",
+			states:    []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready},
+			wantState: connectivity.Ready,
+		},
+		{
+			name:      "regular transition to transient failure",
+			states:    []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure},
+			wantState: connectivity.TransientFailure,
+		},
+		{
+			name:      "regular transition to ready",
+			states:    []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.Idle},
+			wantState: connectivity.Idle,
+		},
+		{
+			name:      "transition from ready to transient failure",
+			states:    []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure},
+			wantState: connectivity.TransientFailure,
+		},
+		{
+			name:      "transition from transient failure back to ready",
+			states:    []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure, connectivity.Ready},
+			wantState: connectivity.Ready,
+		},
+		{
+			// This state transition is usually suppressed at the LB policy level, by
+			// not calling RecordTransition.
+			name:      "transition from transient failure back to idle",
+			states:    []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure, connectivity.Idle},
+			wantState: connectivity.Idle,
+		},
+		{
+			// This state transition is usually suppressed at the LB policy level, by
+			// not calling RecordTransition.
+			name:      "transition from transient failure back to connecting",
+			states:    []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure, connectivity.Connecting},
+			wantState: connectivity.Connecting,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			cse := &ConnectivityStateEvaluator{}
+			var prevState, gotState connectivity.State
+			prevState = connectivity.Shutdown
+			for _, newState := range test.states {
+				gotState = cse.RecordTransition(prevState, newState)
+				prevState = newState
+			}
+			if gotState != test.wantState {
+				t.Fatalf("RecordTransition() = %v, want %v", gotState, test.wantState)
+			}
+		})
+	}
+}
+
+// TestRecordTransition_MultipleSubConns_DifferentStates tests state transitions
+// among multiple subConns, and verifies that the connectivity state aggregation
+// algorithm produces the expected aggregate connectivity state.
+func (s) TestRecordTransition_MultipleSubConns_DifferentStates(t *testing.T) {
+	tests := []struct {
+		name string
+		// Each entry in this slice corresponds to the state changes happening on an
+		// individual subConn.
+		subConnStates [][]connectivity.State
+		wantState     connectivity.State
+	}{
+		{
+			name: "atleast one ready",
+			subConnStates: [][]connectivity.State{
+				{connectivity.Idle, connectivity.Connecting, connectivity.Ready},
+				{connectivity.Idle},
+				{connectivity.Idle, connectivity.Connecting},
+				{connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure},
+			},
+			wantState: connectivity.Ready,
+		},
+		{
+			name: "atleast one connecting",
+			subConnStates: [][]connectivity.State{
+				{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.Connecting},
+				{connectivity.Idle},
+				{connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure},
+			},
+			wantState: connectivity.Connecting,
+		},
+		{
+			name: "atleast one idle",
+			subConnStates: [][]connectivity.State{
+				{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.Idle},
+				{connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure},
+			},
+			wantState: connectivity.Idle,
+		},
+		{
+			name: "atleast one transient failure",
+			subConnStates: [][]connectivity.State{
+				{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure},
+				{connectivity.TransientFailure},
+			},
+			wantState: connectivity.TransientFailure,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			cse := &ConnectivityStateEvaluator{}
+			var prevState, gotState connectivity.State
+			for _, scStates := range test.subConnStates {
+				prevState = connectivity.Shutdown
+				for _, newState := range scStates {
+					gotState = cse.RecordTransition(prevState, newState)
+					prevState = newState
+				}
+			}
+			if gotState != test.wantState {
+				t.Fatalf("RecordTransition() = %v, want %v", gotState, test.wantState)
+			}
+		})
+	}
+}

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()
+}

balancer/endpointsharding/endpointsharding_ext_test.go

@@ -0,0 +1,353 @@
+/*
+ *
+ * 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_test
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/backoff"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/stubserver"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/internal/testutils/roundrobin"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/grpc/serviceconfig"
+	"google.golang.org/grpc/status"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+)
+
+var (
+	defaultTestTimeout      = time.Second * 10
+	defaultTestShortTimeout = time.Millisecond * 10
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+var logger = grpclog.Component("endpoint-sharding-test")
+
+func init() {
+	balancer.Register(fakePetioleBuilder{})
+}
+
+const fakePetioleName = "fake_petiole"
+
+type fakePetioleBuilder struct{}
+
+func (fakePetioleBuilder) Name() string {
+	return fakePetioleName
+}
+
+func (fakePetioleBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	fp := &fakePetiole{
+		ClientConn: cc,
+		bOpts:      opts,
+	}
+	fp.Balancer = endpointsharding.NewBalancer(fp, opts, balancer.Get(pickfirst.Name).Build, endpointsharding.Options{})
+	return fp
+}
+
+func (fakePetioleBuilder) ParseConfig(json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	return nil, nil
+}
+
+// fakePetiole is a load balancer that wraps the endpointShardingBalancer, and
+// forwards ClientConnUpdates with a child config of graceful switch that wraps
+// pick first. It also intercepts UpdateState to make sure it can access the
+// child state maintained by EndpointSharding.
+type fakePetiole struct {
+	balancer.Balancer
+	balancer.ClientConn
+	bOpts balancer.BuildOptions
+}
+
+func (fp *fakePetiole) UpdateClientConnState(state balancer.ClientConnState) error {
+	if el := state.ResolverState.Endpoints; len(el) != 2 {
+		return fmt.Errorf("UpdateClientConnState wants two endpoints, got: %v", el)
+	}
+
+	return fp.Balancer.UpdateClientConnState(state)
+}
+
+func (fp *fakePetiole) UpdateState(state balancer.State) {
+	childStates := endpointsharding.ChildStatesFromPicker(state.Picker)
+	// Both child states should be present in the child picker. States and
+	// picker change over the lifecycle of test, but there should always be two.
+	if len(childStates) != 2 {
+		logger.Fatal(fmt.Errorf("length of child states received: %v, want 2", len(childStates)))
+	}
+
+	fp.ClientConn.UpdateState(state)
+}
+
+// TestEndpointShardingBasic tests the basic functionality of the endpoint
+// sharding balancer. It specifies a petiole policy that is essentially a
+// wrapper around the endpoint sharder. Two backends are started, with each
+// backend's address specified in an endpoint. The petiole does not have a
+// special picker, so it should fallback to the default behavior, which is to
+// round_robin amongst the endpoint children that are in the aggregated state.
+// It also verifies the petiole has access to the raw child state in case it
+// wants to implement a custom picker. The test sends a resolver error to the
+// endpointsharding balancer and verifies an error picker from the children
+// is used while making an RPC.
+func (s) TestEndpointShardingBasic(t *testing.T) {
+	backend1 := stubserver.StartTestService(t, nil)
+	defer backend1.Stop()
+	backend2 := stubserver.StartTestService(t, nil)
+	defer backend2.Stop()
+
+	mr := manual.NewBuilderWithScheme("e2e-test")
+	defer mr.Close()
+
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, fakePetioleName)
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(json)
+	mr.InitialState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend1.Address}}},
+			{Addresses: []resolver.Address{{Addr: backend2.Address}}},
+		},
+		ServiceConfig: sc,
+	})
+
+	dOpts := []grpc.DialOption{
+		grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()),
+		// Use a large backoff delay to avoid the error picker being updated
+		// too quickly.
+		grpc.WithConnectParams(grpc.ConnectParams{
+			Backoff: backoff.Config{
+				BaseDelay:  2 * defaultTestTimeout,
+				Multiplier: float64(0),
+				Jitter:     float64(0),
+				MaxDelay:   2 * defaultTestTimeout,
+			},
+		}),
+	}
+	cc, err := grpc.NewClient(mr.Scheme()+":///", dOpts...)
+	if err != nil {
+		t.Fatalf("Failed to create new client: %v", err)
+	}
+	defer cc.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	client := testgrpc.NewTestServiceClient(cc)
+	// Assert a round robin distribution between the two spun up backends. This
+	// requires a poll and eventual consistency as both endpoint children do not
+	// start in state READY.
+	if err = roundrobin.CheckRoundRobinRPCs(ctx, client, []resolver.Address{{Addr: backend1.Address}, {Addr: backend2.Address}}); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// Stopping both the backends should make the channel enter
+	// TransientFailure.
+	backend1.Stop()
+	backend2.Stop()
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+
+	// When the resolver reports an error, the picker should get updated to
+	// return the resolver error.
+	mr.CC().ReportError(errors.New("test error"))
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+	for ; ctx.Err() == nil; <-time.After(time.Millisecond) {
+		_, err := client.EmptyCall(ctx, &testpb.Empty{})
+		if err == nil {
+			t.Fatalf("EmptyCall succeeded when expected to fail with %q", "test error")
+		}
+		if strings.Contains(err.Error(), "test error") {
+			break
+		}
+	}
+	if ctx.Err() != nil {
+		t.Fatalf("Context timed out waiting for picker with resolver error.")
+	}
+}
+
+// Tests that endpointsharding doesn't automatically re-connect IDLE children.
+// The test creates an endpoint with two servers and another with a single
+// server. The active service in endpoint 1 is closed to make the child
+// pickfirst enter IDLE state. The test verifies that the child pickfirst
+// doesn't connect to the second address in the endpoint.
+func (s) TestEndpointShardingReconnectDisabled(t *testing.T) {
+	backend1 := stubserver.StartTestService(t, nil)
+	defer backend1.Stop()
+	backend2 := stubserver.StartTestService(t, nil)
+	defer backend2.Stop()
+	backend3 := stubserver.StartTestService(t, nil)
+	defer backend3.Stop()
+
+	mr := manual.NewBuilderWithScheme("e2e-test")
+	defer mr.Close()
+
+	name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
+	bf := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			epOpts := endpointsharding.Options{DisableAutoReconnect: true}
+			bd.ChildBalancer = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build, epOpts)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+	}
+	stub.Register(name, bf)
+
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, name)
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(json)
+	mr.InitialState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend1.Address}, {Addr: backend2.Address}}},
+			{Addresses: []resolver.Address{{Addr: backend3.Address}}},
+		},
+		ServiceConfig: sc,
+	})
+
+	cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("Failed to create new client: %v", err)
+	}
+	defer cc.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	client := testgrpc.NewTestServiceClient(cc)
+	// Assert a round robin distribution between the two spun up backends. This
+	// requires a poll and eventual consistency as both endpoint children do not
+	// start in state READY.
+	if err = roundrobin.CheckRoundRobinRPCs(ctx, client, []resolver.Address{{Addr: backend1.Address}, {Addr: backend3.Address}}); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// On closing the first server, the first child balancer should enter
+	// IDLE. Since endpointsharding is configured not to auto-reconnect, it will
+	// remain IDLE and will not try to connect to the second backend in the same
+	// endpoint.
+	backend1.Stop()
+	// CheckRoundRobinRPCs waits for all the backends to become reachable, we
+	// call it to ensure the picker no longer sends RPCs to closed backend.
+	if err = roundrobin.CheckRoundRobinRPCs(ctx, client, []resolver.Address{{Addr: backend3.Address}}); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// Verify requests go only to backend3 for a short time.
+	shortCtx, cancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer cancel()
+	for ; shortCtx.Err() == nil; <-time.After(time.Millisecond) {
+		var peer peer.Peer
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
+			if status.Code(err) != codes.DeadlineExceeded {
+				t.Fatalf("EmptyCall() returned unexpected error %v", err)
+			}
+			break
+		}
+		if got, want := peer.Addr.String(), backend3.Address; got != want {
+			t.Fatalf("EmptyCall() went to unexpected backend: got %q, want %q", got, want)
+		}
+	}
+}
+
+// Tests that endpointsharding doesn't automatically re-connect IDLE children
+// until cc.Connect() is called. The test creates an endpoint with a single
+// address. The client is connected and the active server is closed to make the
+// child pickfirst enter IDLE state. The test verifies that the child pickfirst
+// doesn't re-connect automatically. The test calls cc.Connect() and verified
+// that the balancer connects causing the channel to enter TransientFailure.
+func (s) TestEndpointShardingExitIdle(t *testing.T) {
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+
+	mr := manual.NewBuilderWithScheme("e2e-test")
+	defer mr.Close()
+
+	name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
+	bf := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			epOpts := endpointsharding.Options{DisableAutoReconnect: true}
+			bd.ChildBalancer = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build, epOpts)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+		ExitIdle: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.ExitIdle()
+		},
+	}
+	stub.Register(name, bf)
+
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, name)
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(json)
+	mr.InitialState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend.Address}}},
+		},
+		ServiceConfig: sc,
+	})
+
+	cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("Failed to create new client: %v", err)
+	}
+	defer cc.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
+	}
+
+	// On closing the first server, the first child balancer should enter
+	// IDLE. Since endpointsharding is configured not to auto-reconnect, it will
+	// remain IDLE and will not try to re-connect
+	backend.Stop()
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+	shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer shortCancel()
+	testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
+
+	// The balancer should try to re-connect and fail.
+	cc.Connect()
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+}

balancer/endpointsharding/endpointsharding_test.go

@@ -0,0 +1,83 @@
+/*
+ *
+ * Copyright 2025 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
+
+import (
+	"fmt"
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/resolver"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestRotateEndpoints(t *testing.T) {
+	ep := func(addr string) resolver.Endpoint {
+		return resolver.Endpoint{Addresses: []resolver.Address{{Addr: addr}}}
+	}
+	endpoints := []resolver.Endpoint{ep("1"), ep("2"), ep("3"), ep("4"), ep("5")}
+	testCases := []struct {
+		rval int
+		want []resolver.Endpoint
+	}{
+		{
+			rval: 0,
+			want: []resolver.Endpoint{ep("1"), ep("2"), ep("3"), ep("4"), ep("5")},
+		},
+		{
+			rval: 1,
+			want: []resolver.Endpoint{ep("2"), ep("3"), ep("4"), ep("5"), ep("1")},
+		},
+		{
+			rval: 2,
+			want: []resolver.Endpoint{ep("3"), ep("4"), ep("5"), ep("1"), ep("2")},
+		},
+		{
+			rval: 3,
+			want: []resolver.Endpoint{ep("4"), ep("5"), ep("1"), ep("2"), ep("3")},
+		},
+		{
+			rval: 4,
+			want: []resolver.Endpoint{ep("5"), ep("1"), ep("2"), ep("3"), ep("4")},
+		},
+	}
+
+	defer func(r func(int) int) {
+		randIntN = r
+	}(randIntN)
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprint(tc.rval), func(t *testing.T) {
+			randIntN = func(int) int {
+				return tc.rval
+			}
+			got := rotateEndpoints(endpoints)
+			if fmt.Sprint(got) != fmt.Sprint(tc.want) {
+				t.Fatalf("rand=%v; rotateEndpoints(%v) = %v; want %v", tc.rval, endpoints, got, tc.want)
+			}
+		})
+	}
+}

balancer/grpclb/grpc_lb_v1/load_balancer.pb.go

@@ -0,0 +1,763 @@
+// Copyright 2015 The 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.
+
+// This file defines the GRPCLB LoadBalancing protocol.
+//
+// The canonical version of this proto can be found at
+// https://github.com/grpc/grpc-proto/blob/master/grpc/lb/v1/load_balancer.proto
+
+// Code generated by protoc-gen-go. DO NOT EDIT.
+// versions:
+// 	protoc-gen-go v1.36.12
+// 	protoc        v5.27.1
+// source: grpc/lb/v1/load_balancer.proto
+
+package grpc_lb_v1
+
+import (
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
+	durationpb "google.golang.org/protobuf/types/known/durationpb"
+	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
+	reflect "reflect"
+	sync "sync"
+	unsafe "unsafe"
+)
+
+const (
+	// Verify that this generated code is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
+	// Verify that runtime/protoimpl is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
+)
+
+type LoadBalanceRequest struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// Types that are valid to be assigned to LoadBalanceRequestType:
+	//
+	//	*LoadBalanceRequest_InitialRequest
+	//	*LoadBalanceRequest_ClientStats
+	LoadBalanceRequestType isLoadBalanceRequest_LoadBalanceRequestType `protobuf_oneof:"load_balance_request_type"`
+	unknownFields          protoimpl.UnknownFields
+	sizeCache              protoimpl.SizeCache
+}
+
+func (x *LoadBalanceRequest) Reset() {
+	*x = LoadBalanceRequest{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[0]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *LoadBalanceRequest) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*LoadBalanceRequest) ProtoMessage() {}
+
+func (x *LoadBalanceRequest) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[0]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use LoadBalanceRequest.ProtoReflect.Descriptor instead.
+func (*LoadBalanceRequest) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{0}
+}
+
+func (x *LoadBalanceRequest) GetLoadBalanceRequestType() isLoadBalanceRequest_LoadBalanceRequestType {
+	if x != nil {
+		return x.LoadBalanceRequestType
+	}
+	return nil
+}
+
+func (x *LoadBalanceRequest) GetInitialRequest() *InitialLoadBalanceRequest {
+	if x != nil {
+		if x, ok := x.LoadBalanceRequestType.(*LoadBalanceRequest_InitialRequest); ok {
+			return x.InitialRequest
+		}
+	}
+	return nil
+}
+
+func (x *LoadBalanceRequest) GetClientStats() *ClientStats {
+	if x != nil {
+		if x, ok := x.LoadBalanceRequestType.(*LoadBalanceRequest_ClientStats); ok {
+			return x.ClientStats
+		}
+	}
+	return nil
+}
+
+type isLoadBalanceRequest_LoadBalanceRequestType interface {
+	isLoadBalanceRequest_LoadBalanceRequestType()
+}
+
+type LoadBalanceRequest_InitialRequest struct {
+	// This message should be sent on the first request to the load balancer.
+	InitialRequest *InitialLoadBalanceRequest `protobuf:"bytes,1,opt,name=initial_request,json=initialRequest,proto3,oneof"`
+}
+
+type LoadBalanceRequest_ClientStats struct {
+	// The client stats should be periodically reported to the load balancer
+	// based on the duration defined in the InitialLoadBalanceResponse.
+	ClientStats *ClientStats `protobuf:"bytes,2,opt,name=client_stats,json=clientStats,proto3,oneof"`
+}
+
+func (*LoadBalanceRequest_InitialRequest) isLoadBalanceRequest_LoadBalanceRequestType() {}
+
+func (*LoadBalanceRequest_ClientStats) isLoadBalanceRequest_LoadBalanceRequestType() {}
+
+type InitialLoadBalanceRequest struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// The name of the load balanced service (e.g., service.googleapis.com). Its
+	// length should be less than 256 bytes.
+	// The name might include a port number. How to handle the port number is up
+	// to the balancer.
+	Name          string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *InitialLoadBalanceRequest) Reset() {
+	*x = InitialLoadBalanceRequest{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[1]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *InitialLoadBalanceRequest) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*InitialLoadBalanceRequest) ProtoMessage() {}
+
+func (x *InitialLoadBalanceRequest) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[1]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use InitialLoadBalanceRequest.ProtoReflect.Descriptor instead.
+func (*InitialLoadBalanceRequest) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{1}
+}
+
+func (x *InitialLoadBalanceRequest) GetName() string {
+	if x != nil {
+		return x.Name
+	}
+	return ""
+}
+
+// Contains the number of calls finished for a particular load balance token.
+type ClientStatsPerToken struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// See Server.load_balance_token.
+	LoadBalanceToken string `protobuf:"bytes,1,opt,name=load_balance_token,json=loadBalanceToken,proto3" json:"load_balance_token,omitempty"`
+	// The total number of RPCs that finished associated with the token.
+	NumCalls      int64 `protobuf:"varint,2,opt,name=num_calls,json=numCalls,proto3" json:"num_calls,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *ClientStatsPerToken) Reset() {
+	*x = ClientStatsPerToken{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[2]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *ClientStatsPerToken) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*ClientStatsPerToken) ProtoMessage() {}
+
+func (x *ClientStatsPerToken) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[2]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use ClientStatsPerToken.ProtoReflect.Descriptor instead.
+func (*ClientStatsPerToken) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{2}
+}
+
+func (x *ClientStatsPerToken) GetLoadBalanceToken() string {
+	if x != nil {
+		return x.LoadBalanceToken
+	}
+	return ""
+}
+
+func (x *ClientStatsPerToken) GetNumCalls() int64 {
+	if x != nil {
+		return x.NumCalls
+	}
+	return 0
+}
+
+// Contains client level statistics that are useful to load balancing. Each
+// count except the timestamp should be reset to zero after reporting the stats.
+type ClientStats struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// The timestamp of generating the report.
+	Timestamp *timestamppb.Timestamp `protobuf:"bytes,1,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
+	// The total number of RPCs that started.
+	NumCallsStarted int64 `protobuf:"varint,2,opt,name=num_calls_started,json=numCallsStarted,proto3" json:"num_calls_started,omitempty"`
+	// The total number of RPCs that finished.
+	NumCallsFinished int64 `protobuf:"varint,3,opt,name=num_calls_finished,json=numCallsFinished,proto3" json:"num_calls_finished,omitempty"`
+	// The total number of RPCs that failed to reach a server except dropped RPCs.
+	NumCallsFinishedWithClientFailedToSend int64 `protobuf:"varint,6,opt,name=num_calls_finished_with_client_failed_to_send,json=numCallsFinishedWithClientFailedToSend,proto3" json:"num_calls_finished_with_client_failed_to_send,omitempty"`
+	// The total number of RPCs that finished and are known to have been received
+	// by a server.
+	NumCallsFinishedKnownReceived int64 `protobuf:"varint,7,opt,name=num_calls_finished_known_received,json=numCallsFinishedKnownReceived,proto3" json:"num_calls_finished_known_received,omitempty"`
+	// The list of dropped calls.
+	CallsFinishedWithDrop []*ClientStatsPerToken `protobuf:"bytes,8,rep,name=calls_finished_with_drop,json=callsFinishedWithDrop,proto3" json:"calls_finished_with_drop,omitempty"`
+	unknownFields         protoimpl.UnknownFields
+	sizeCache             protoimpl.SizeCache
+}
+
+func (x *ClientStats) Reset() {
+	*x = ClientStats{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[3]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *ClientStats) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*ClientStats) ProtoMessage() {}
+
+func (x *ClientStats) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[3]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use ClientStats.ProtoReflect.Descriptor instead.
+func (*ClientStats) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{3}
+}
+
+func (x *ClientStats) GetTimestamp() *timestamppb.Timestamp {
+	if x != nil {
+		return x.Timestamp
+	}
+	return nil
+}
+
+func (x *ClientStats) GetNumCallsStarted() int64 {
+	if x != nil {
+		return x.NumCallsStarted
+	}
+	return 0
+}
+
+func (x *ClientStats) GetNumCallsFinished() int64 {
+	if x != nil {
+		return x.NumCallsFinished
+	}
+	return 0
+}
+
+func (x *ClientStats) GetNumCallsFinishedWithClientFailedToSend() int64 {
+	if x != nil {
+		return x.NumCallsFinishedWithClientFailedToSend
+	}
+	return 0
+}
+
+func (x *ClientStats) GetNumCallsFinishedKnownReceived() int64 {
+	if x != nil {
+		return x.NumCallsFinishedKnownReceived
+	}
+	return 0
+}
+
+func (x *ClientStats) GetCallsFinishedWithDrop() []*ClientStatsPerToken {
+	if x != nil {
+		return x.CallsFinishedWithDrop
+	}
+	return nil
+}
+
+type LoadBalanceResponse struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// Types that are valid to be assigned to LoadBalanceResponseType:
+	//
+	//	*LoadBalanceResponse_InitialResponse
+	//	*LoadBalanceResponse_ServerList
+	//	*LoadBalanceResponse_FallbackResponse
+	LoadBalanceResponseType isLoadBalanceResponse_LoadBalanceResponseType `protobuf_oneof:"load_balance_response_type"`
+	unknownFields           protoimpl.UnknownFields
+	sizeCache               protoimpl.SizeCache
+}
+
+func (x *LoadBalanceResponse) Reset() {
+	*x = LoadBalanceResponse{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[4]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *LoadBalanceResponse) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*LoadBalanceResponse) ProtoMessage() {}
+
+func (x *LoadBalanceResponse) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[4]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use LoadBalanceResponse.ProtoReflect.Descriptor instead.
+func (*LoadBalanceResponse) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{4}
+}
+
+func (x *LoadBalanceResponse) GetLoadBalanceResponseType() isLoadBalanceResponse_LoadBalanceResponseType {
+	if x != nil {
+		return x.LoadBalanceResponseType
+	}
+	return nil
+}
+
+func (x *LoadBalanceResponse) GetInitialResponse() *InitialLoadBalanceResponse {
+	if x != nil {
+		if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_InitialResponse); ok {
+			return x.InitialResponse
+		}
+	}
+	return nil
+}
+
+func (x *LoadBalanceResponse) GetServerList() *ServerList {
+	if x != nil {
+		if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_ServerList); ok {
+			return x.ServerList
+		}
+	}
+	return nil
+}
+
+func (x *LoadBalanceResponse) GetFallbackResponse() *FallbackResponse {
+	if x != nil {
+		if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_FallbackResponse); ok {
+			return x.FallbackResponse
+		}
+	}
+	return nil
+}
+
+type isLoadBalanceResponse_LoadBalanceResponseType interface {
+	isLoadBalanceResponse_LoadBalanceResponseType()
+}
+
+type LoadBalanceResponse_InitialResponse struct {
+	// This message should be sent on the first response to the client.
+	InitialResponse *InitialLoadBalanceResponse `protobuf:"bytes,1,opt,name=initial_response,json=initialResponse,proto3,oneof"`
+}
+
+type LoadBalanceResponse_ServerList struct {
+	// Contains the list of servers selected by the load balancer. The client
+	// should send requests to these servers in the specified order.
+	ServerList *ServerList `protobuf:"bytes,2,opt,name=server_list,json=serverList,proto3,oneof"`
+}
+
+type LoadBalanceResponse_FallbackResponse struct {
+	// If this field is set, then the client should eagerly enter fallback
+	// mode (even if there are existing, healthy connections to backends).
+	FallbackResponse *FallbackResponse `protobuf:"bytes,3,opt,name=fallback_response,json=fallbackResponse,proto3,oneof"`
+}
+
+func (*LoadBalanceResponse_InitialResponse) isLoadBalanceResponse_LoadBalanceResponseType() {}
+
+func (*LoadBalanceResponse_ServerList) isLoadBalanceResponse_LoadBalanceResponseType() {}
+
+func (*LoadBalanceResponse_FallbackResponse) isLoadBalanceResponse_LoadBalanceResponseType() {}
+
+type FallbackResponse struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *FallbackResponse) Reset() {
+	*x = FallbackResponse{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[5]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *FallbackResponse) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*FallbackResponse) ProtoMessage() {}
+
+func (x *FallbackResponse) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[5]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use FallbackResponse.ProtoReflect.Descriptor instead.
+func (*FallbackResponse) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{5}
+}
+
+type InitialLoadBalanceResponse struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// This interval defines how often the client should send the client stats
+	// to the load balancer. Stats should only be reported when the duration is
+	// positive.
+	ClientStatsReportInterval *durationpb.Duration `protobuf:"bytes,2,opt,name=client_stats_report_interval,json=clientStatsReportInterval,proto3" json:"client_stats_report_interval,omitempty"`
+	unknownFields             protoimpl.UnknownFields
+	sizeCache                 protoimpl.SizeCache
+}
+
+func (x *InitialLoadBalanceResponse) Reset() {
+	*x = InitialLoadBalanceResponse{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[6]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *InitialLoadBalanceResponse) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*InitialLoadBalanceResponse) ProtoMessage() {}
+
+func (x *InitialLoadBalanceResponse) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[6]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use InitialLoadBalanceResponse.ProtoReflect.Descriptor instead.
+func (*InitialLoadBalanceResponse) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{6}
+}
+
+func (x *InitialLoadBalanceResponse) GetClientStatsReportInterval() *durationpb.Duration {
+	if x != nil {
+		return x.ClientStatsReportInterval
+	}
+	return nil
+}
+
+type ServerList struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// Contains a list of servers selected by the load balancer. The list will
+	// be updated when server resolutions change or as needed to balance load
+	// across more servers. The client should consume the server list in order
+	// unless instructed otherwise via the client_config.
+	Servers       []*Server `protobuf:"bytes,1,rep,name=servers,proto3" json:"servers,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *ServerList) Reset() {
+	*x = ServerList{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[7]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *ServerList) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*ServerList) ProtoMessage() {}
+
+func (x *ServerList) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[7]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use ServerList.ProtoReflect.Descriptor instead.
+func (*ServerList) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{7}
+}
+
+func (x *ServerList) GetServers() []*Server {
+	if x != nil {
+		return x.Servers
+	}
+	return nil
+}
+
+// Contains server information. When the drop field is not true, use the other
+// fields.
+type Server struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// A resolved address for the server, serialized in network-byte-order. It may
+	// either be an IPv4 or IPv6 address.
+	IpAddress []byte `protobuf:"bytes,1,opt,name=ip_address,json=ipAddress,proto3" json:"ip_address,omitempty"`
+	// A resolved port number for the server.
+	Port int32 `protobuf:"varint,2,opt,name=port,proto3" json:"port,omitempty"`
+	// An opaque but printable token for load reporting. The client must include
+	// the token of the picked server into the initial metadata when it starts a
+	// call to that server. The token is used by the server to verify the request
+	// and to allow the server to report load to the gRPC LB system. The token is
+	// also used in client stats for reporting dropped calls.
+	//
+	// Its length can be variable but must be less than 50 bytes.
+	LoadBalanceToken string `protobuf:"bytes,3,opt,name=load_balance_token,json=loadBalanceToken,proto3" json:"load_balance_token,omitempty"`
+	// Indicates whether this particular request should be dropped by the client.
+	// If the request is dropped, there will be a corresponding entry in
+	// ClientStats.calls_finished_with_drop.
+	Drop          bool `protobuf:"varint,4,opt,name=drop,proto3" json:"drop,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *Server) Reset() {
+	*x = Server{}
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[8]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *Server) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*Server) ProtoMessage() {}
+
+func (x *Server) ProtoReflect() protoreflect.Message {
+	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[8]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use Server.ProtoReflect.Descriptor instead.
+func (*Server) Descriptor() ([]byte, []int) {
+	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{8}
+}
+
+func (x *Server) GetIpAddress() []byte {
+	if x != nil {
+		return x.IpAddress
+	}
+	return nil
+}
+
+func (x *Server) GetPort() int32 {
+	if x != nil {
+		return x.Port
+	}
+	return 0
+}
+
+func (x *Server) GetLoadBalanceToken() string {
+	if x != nil {
+		return x.LoadBalanceToken
+	}
+	return ""
+}
+
+func (x *Server) GetDrop() bool {
+	if x != nil {
+		return x.Drop
+	}
+	return false
+}
+
+var File_grpc_lb_v1_load_balancer_proto protoreflect.FileDescriptor
+
+const file_grpc_lb_v1_load_balancer_proto_rawDesc = "" +
+	"\n" +
+	"\x1egrpc/lb/v1/load_balancer.proto\x12\n" +
+	"grpc.lb.v1\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\xc1\x01\n" +
+	"\x12LoadBalanceRequest\x12P\n" +
+	"\x0finitial_request\x18\x01 \x01(\v2%.grpc.lb.v1.InitialLoadBalanceRequestH\x00R\x0einitialRequest\x12<\n" +
+	"\fclient_stats\x18\x02 \x01(\v2\x17.grpc.lb.v1.ClientStatsH\x00R\vclientStatsB\x1b\n" +
+	"\x19load_balance_request_type\"/\n" +
+	"\x19InitialLoadBalanceRequest\x12\x12\n" +
+	"\x04name\x18\x01 \x01(\tR\x04name\"`\n" +
+	"\x13ClientStatsPerToken\x12,\n" +
+	"\x12load_balance_token\x18\x01 \x01(\tR\x10loadBalanceToken\x12\x1b\n" +
+	"\tnum_calls\x18\x02 \x01(\x03R\bnumCalls\"\xb0\x03\n" +
+	"\vClientStats\x128\n" +
+	"\ttimestamp\x18\x01 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12*\n" +
+	"\x11num_calls_started\x18\x02 \x01(\x03R\x0fnumCallsStarted\x12,\n" +
+	"\x12num_calls_finished\x18\x03 \x01(\x03R\x10numCallsFinished\x12]\n" +
+	"-num_calls_finished_with_client_failed_to_send\x18\x06 \x01(\x03R&numCallsFinishedWithClientFailedToSend\x12H\n" +
+	"!num_calls_finished_known_received\x18\a \x01(\x03R\x1dnumCallsFinishedKnownReceived\x12X\n" +
+	"\x18calls_finished_with_drop\x18\b \x03(\v2\x1f.grpc.lb.v1.ClientStatsPerTokenR\x15callsFinishedWithDropJ\x04\b\x04\x10\x05J\x04\b\x05\x10\x06\"\x90\x02\n" +
+	"\x13LoadBalanceResponse\x12S\n" +
+	"\x10initial_response\x18\x01 \x01(\v2&.grpc.lb.v1.InitialLoadBalanceResponseH\x00R\x0finitialResponse\x129\n" +
+	"\vserver_list\x18\x02 \x01(\v2\x16.grpc.lb.v1.ServerListH\x00R\n" +
+	"serverList\x12K\n" +
+	"\x11fallback_response\x18\x03 \x01(\v2\x1c.grpc.lb.v1.FallbackResponseH\x00R\x10fallbackResponseB\x1c\n" +
+	"\x1aload_balance_response_type\"\x12\n" +
+	"\x10FallbackResponse\"~\n" +
+	"\x1aInitialLoadBalanceResponse\x12Z\n" +
+	"\x1cclient_stats_report_interval\x18\x02 \x01(\v2\x19.google.protobuf.DurationR\x19clientStatsReportIntervalJ\x04\b\x01\x10\x02\"@\n" +
+	"\n" +
+	"ServerList\x12,\n" +
+	"\aservers\x18\x01 \x03(\v2\x12.grpc.lb.v1.ServerR\aserversJ\x04\b\x03\x10\x04\"\x83\x01\n" +
+	"\x06Server\x12\x1d\n" +
+	"\n" +
+	"ip_address\x18\x01 \x01(\fR\tipAddress\x12\x12\n" +
+	"\x04port\x18\x02 \x01(\x05R\x04port\x12,\n" +
+	"\x12load_balance_token\x18\x03 \x01(\tR\x10loadBalanceToken\x12\x12\n" +
+	"\x04drop\x18\x04 \x01(\bR\x04dropJ\x04\b\x05\x10\x062b\n" +
+	"\fLoadBalancer\x12R\n" +
+	"\vBalanceLoad\x12\x1e.grpc.lb.v1.LoadBalanceRequest\x1a\x1f.grpc.lb.v1.LoadBalanceResponse(\x010\x01BW\n" +
+	"\rio.grpc.lb.v1B\x11LoadBalancerProtoP\x01Z1google.golang.org/grpc/balancer/grpclb/grpc_lb_v1b\x06proto3"
+
+var (
+	file_grpc_lb_v1_load_balancer_proto_rawDescOnce sync.Once
+	file_grpc_lb_v1_load_balancer_proto_rawDescData []byte
+)
+
+func file_grpc_lb_v1_load_balancer_proto_rawDescGZIP() []byte {
+	file_grpc_lb_v1_load_balancer_proto_rawDescOnce.Do(func() {
+		file_grpc_lb_v1_load_balancer_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_grpc_lb_v1_load_balancer_proto_rawDesc), len(file_grpc_lb_v1_load_balancer_proto_rawDesc)))
+	})
+	return file_grpc_lb_v1_load_balancer_proto_rawDescData
+}
+
+var file_grpc_lb_v1_load_balancer_proto_msgTypes = make([]protoimpl.MessageInfo, 9)
+var file_grpc_lb_v1_load_balancer_proto_goTypes = []any{
+	(*LoadBalanceRequest)(nil),         // 0: grpc.lb.v1.LoadBalanceRequest
+	(*InitialLoadBalanceRequest)(nil),  // 1: grpc.lb.v1.InitialLoadBalanceRequest
+	(*ClientStatsPerToken)(nil),        // 2: grpc.lb.v1.ClientStatsPerToken
+	(*ClientStats)(nil),                // 3: grpc.lb.v1.ClientStats
+	(*LoadBalanceResponse)(nil),        // 4: grpc.lb.v1.LoadBalanceResponse
+	(*FallbackResponse)(nil),           // 5: grpc.lb.v1.FallbackResponse
+	(*InitialLoadBalanceResponse)(nil), // 6: grpc.lb.v1.InitialLoadBalanceResponse
+	(*ServerList)(nil),                 // 7: grpc.lb.v1.ServerList
+	(*Server)(nil),                     // 8: grpc.lb.v1.Server
+	(*timestamppb.Timestamp)(nil),      // 9: google.protobuf.Timestamp
+	(*durationpb.Duration)(nil),        // 10: google.protobuf.Duration
+}
+var file_grpc_lb_v1_load_balancer_proto_depIdxs = []int32{
+	1,  // 0: grpc.lb.v1.LoadBalanceRequest.initial_request:type_name -> grpc.lb.v1.InitialLoadBalanceRequest
+	3,  // 1: grpc.lb.v1.LoadBalanceRequest.client_stats:type_name -> grpc.lb.v1.ClientStats
+	9,  // 2: grpc.lb.v1.ClientStats.timestamp:type_name -> google.protobuf.Timestamp
+	2,  // 3: grpc.lb.v1.ClientStats.calls_finished_with_drop:type_name -> grpc.lb.v1.ClientStatsPerToken
+	6,  // 4: grpc.lb.v1.LoadBalanceResponse.initial_response:type_name -> grpc.lb.v1.InitialLoadBalanceResponse
+	7,  // 5: grpc.lb.v1.LoadBalanceResponse.server_list:type_name -> grpc.lb.v1.ServerList
+	5,  // 6: grpc.lb.v1.LoadBalanceResponse.fallback_response:type_name -> grpc.lb.v1.FallbackResponse
+	10, // 7: grpc.lb.v1.InitialLoadBalanceResponse.client_stats_report_interval:type_name -> google.protobuf.Duration
+	8,  // 8: grpc.lb.v1.ServerList.servers:type_name -> grpc.lb.v1.Server
+	0,  // 9: grpc.lb.v1.LoadBalancer.BalanceLoad:input_type -> grpc.lb.v1.LoadBalanceRequest
+	4,  // 10: grpc.lb.v1.LoadBalancer.BalanceLoad:output_type -> grpc.lb.v1.LoadBalanceResponse
+	10, // [10:11] is the sub-list for method output_type
+	9,  // [9:10] is the sub-list for method input_type
+	9,  // [9:9] is the sub-list for extension type_name
+	9,  // [9:9] is the sub-list for extension extendee
+	0,  // [0:9] is the sub-list for field type_name
+}
+
+func init() { file_grpc_lb_v1_load_balancer_proto_init() }
+func file_grpc_lb_v1_load_balancer_proto_init() {
+	if File_grpc_lb_v1_load_balancer_proto != nil {
+		return
+	}
+	file_grpc_lb_v1_load_balancer_proto_msgTypes[0].OneofWrappers = []any{
+		(*LoadBalanceRequest_InitialRequest)(nil),
+		(*LoadBalanceRequest_ClientStats)(nil),
+	}
+	file_grpc_lb_v1_load_balancer_proto_msgTypes[4].OneofWrappers = []any{
+		(*LoadBalanceResponse_InitialResponse)(nil),
+		(*LoadBalanceResponse_ServerList)(nil),
+		(*LoadBalanceResponse_FallbackResponse)(nil),
+	}
+	type x struct{}
+	out := protoimpl.TypeBuilder{
+		File: protoimpl.DescBuilder{
+			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
+			RawDescriptor: unsafe.Slice(unsafe.StringData(file_grpc_lb_v1_load_balancer_proto_rawDesc), len(file_grpc_lb_v1_load_balancer_proto_rawDesc)),
+			NumEnums:      0,
+			NumMessages:   9,
+			NumExtensions: 0,
+			NumServices:   1,
+		},
+		GoTypes:           file_grpc_lb_v1_load_balancer_proto_goTypes,
+		DependencyIndexes: file_grpc_lb_v1_load_balancer_proto_depIdxs,
+		MessageInfos:      file_grpc_lb_v1_load_balancer_proto_msgTypes,
+	}.Build()
+	File_grpc_lb_v1_load_balancer_proto = out.File
+	file_grpc_lb_v1_load_balancer_proto_goTypes = nil
+	file_grpc_lb_v1_load_balancer_proto_depIdxs = nil
+}

balancer/grpclb/grpc_lb_v1/load_balancer_grpc.pb.go

@@ -0,0 +1,134 @@
+// Copyright 2015 The 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.
+
+// This file defines the GRPCLB LoadBalancing protocol.
+//
+// The canonical version of this proto can be found at
+// https://github.com/grpc/grpc-proto/blob/master/grpc/lb/v1/load_balancer.proto
+
+// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
+// versions:
+// - protoc-gen-go-grpc v1.6.2
+// - protoc             v5.27.1
+// source: grpc/lb/v1/load_balancer.proto
+
+package grpc_lb_v1
+
+import (
+	context "context"
+	grpc "google.golang.org/grpc"
+	codes "google.golang.org/grpc/codes"
+	status "google.golang.org/grpc/status"
+)
+
+// This is a compile-time assertion to ensure that this generated file
+// is compatible with the grpc package it is being compiled against.
+// Requires gRPC-Go v1.64.0 or later.
+const _ = grpc.SupportPackageIsVersion9
+
+const (
+	LoadBalancer_BalanceLoad_FullMethodName = "/grpc.lb.v1.LoadBalancer/BalanceLoad"
+)
+
+// LoadBalancerClient is the client API for LoadBalancer service.
+//
+// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
+type LoadBalancerClient interface {
+	// Bidirectional rpc to get a list of servers.
+	BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[LoadBalanceRequest, LoadBalanceResponse], error)
+}
+
+type loadBalancerClient struct {
+	cc grpc.ClientConnInterface
+}
+
+func NewLoadBalancerClient(cc grpc.ClientConnInterface) LoadBalancerClient {
+	return &loadBalancerClient{cc}
+}
+
+func (c *loadBalancerClient) BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[LoadBalanceRequest, LoadBalanceResponse], error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	stream, err := c.cc.NewStream(ctx, &LoadBalancer_ServiceDesc.Streams[0], LoadBalancer_BalanceLoad_FullMethodName, cOpts...)
+	if err != nil {
+		return nil, err
+	}
+	x := &grpc.GenericClientStream[LoadBalanceRequest, LoadBalanceResponse]{ClientStream: stream}
+	return x, nil
+}
+
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type LoadBalancer_BalanceLoadClient = grpc.BidiStreamingClient[LoadBalanceRequest, LoadBalanceResponse]
+
+// LoadBalancerServer is the server API for LoadBalancer service.
+// All implementations should embed UnimplementedLoadBalancerServer
+// for forward compatibility.
+type LoadBalancerServer interface {
+	// Bidirectional rpc to get a list of servers.
+	BalanceLoad(grpc.BidiStreamingServer[LoadBalanceRequest, LoadBalanceResponse]) error
+}
+
+// UnimplementedLoadBalancerServer should be embedded to have
+// forward compatible implementations.
+//
+// NOTE: this should be embedded by value instead of pointer to avoid a nil
+// pointer dereference when methods are called.
+type UnimplementedLoadBalancerServer struct{}
+
+func (UnimplementedLoadBalancerServer) BalanceLoad(grpc.BidiStreamingServer[LoadBalanceRequest, LoadBalanceResponse]) error {
+	return status.Error(codes.Unimplemented, "method BalanceLoad not implemented")
+}
+func (UnimplementedLoadBalancerServer) testEmbeddedByValue() {}
+
+// UnsafeLoadBalancerServer may be embedded to opt out of forward compatibility for this service.
+// Use of this interface is not recommended, as added methods to LoadBalancerServer will
+// result in compilation errors.
+type UnsafeLoadBalancerServer interface {
+	mustEmbedUnimplementedLoadBalancerServer()
+}
+
+func RegisterLoadBalancerServer(s grpc.ServiceRegistrar, srv LoadBalancerServer) {
+	// If the following call panics, it indicates UnimplementedLoadBalancerServer was
+	// embedded by pointer and is nil.  This will cause panics if an
+	// unimplemented method is ever invoked, so we test this at initialization
+	// time to prevent it from happening at runtime later due to I/O.
+	if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
+		t.testEmbeddedByValue()
+	}
+	s.RegisterService(&LoadBalancer_ServiceDesc, srv)
+}
+
+func _LoadBalancer_BalanceLoad_Handler(srv interface{}, stream grpc.ServerStream) error {
+	return srv.(LoadBalancerServer).BalanceLoad(&grpc.GenericServerStream[LoadBalanceRequest, LoadBalanceResponse]{ServerStream: stream})
+}
+
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type LoadBalancer_BalanceLoadServer = grpc.BidiStreamingServer[LoadBalanceRequest, LoadBalanceResponse]
+
+// LoadBalancer_ServiceDesc is the grpc.ServiceDesc for LoadBalancer service.
+// It's only intended for direct use with grpc.RegisterService,
+// and not to be introspected or modified (even as a copy)
+var LoadBalancer_ServiceDesc = grpc.ServiceDesc{
+	ServiceName: "grpc.lb.v1.LoadBalancer",
+	HandlerType: (*LoadBalancerServer)(nil),
+	Methods:     []grpc.MethodDesc{},
+	Streams: []grpc.StreamDesc{
+		{
+			StreamName:    "BalanceLoad",
+			Handler:       _LoadBalancer_BalanceLoad_Handler,
+			ServerStreams: true,
+			ClientStreams: true,
+		},
+	},
+	Metadata: "grpc/lb/v1/load_balancer.proto",
+}

balancer/grpclb/grpclb.go

@@ -0,0 +1,535 @@
+/*
+ *
+ * Copyright 2016 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 grpclb defines a grpclb balancer.
+//
+// To install grpclb balancer, import this package as:
+//
+//	import _ "google.golang.org/grpc/balancer/grpclb"
+package grpclb
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	grpclbstate "google.golang.org/grpc/balancer/grpclb/state"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/backoff"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/pretty"
+	"google.golang.org/grpc/internal/resolver/dns"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/protobuf/types/known/durationpb"
+
+	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
+)
+
+const (
+	lbTokenKey             = "lb-token"
+	defaultFallbackTimeout = 10 * time.Second
+	grpclbName             = "grpclb"
+)
+
+var errServerTerminatedConnection = errors.New("grpclb: failed to recv server list: server terminated connection")
+var logger = grpclog.Component("grpclb")
+
+func convertDuration(d *durationpb.Duration) time.Duration {
+	if d == nil {
+		return 0
+	}
+	return time.Duration(d.Seconds)*time.Second + time.Duration(d.Nanos)*time.Nanosecond
+}
+
+// Client API for LoadBalancer service.
+// Mostly copied from generated pb.go file.
+// To avoid circular dependency.
+type loadBalancerClient struct {
+	cc *grpc.ClientConn
+}
+
+func (c *loadBalancerClient) BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (*balanceLoadClientStream, error) {
+	desc := &grpc.StreamDesc{
+		StreamName:    "BalanceLoad",
+		ServerStreams: true,
+		ClientStreams: true,
+	}
+	stream, err := c.cc.NewStream(ctx, desc, "/grpc.lb.v1.LoadBalancer/BalanceLoad", opts...)
+	if err != nil {
+		return nil, err
+	}
+	x := &balanceLoadClientStream{stream}
+	return x, nil
+}
+
+type balanceLoadClientStream struct {
+	grpc.ClientStream
+}
+
+func (x *balanceLoadClientStream) Send(m *lbpb.LoadBalanceRequest) error {
+	return x.ClientStream.SendMsg(m)
+}
+
+func (x *balanceLoadClientStream) Recv() (*lbpb.LoadBalanceResponse, error) {
+	m := new(lbpb.LoadBalanceResponse)
+	if err := x.ClientStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+func init() {
+	balancer.Register(newLBBuilder())
+	dns.EnableSRVLookups = true
+}
+
+// newLBBuilder creates a builder for grpclb.
+func newLBBuilder() balancer.Builder {
+	return newLBBuilderWithFallbackTimeout(defaultFallbackTimeout)
+}
+
+// newLBBuilderWithFallbackTimeout creates a grpclb builder with the given
+// fallbackTimeout. If no response is received from the remote balancer within
+// fallbackTimeout, the backend addresses from the resolved address list will be
+// used.
+//
+// Only call this function when a non-default fallback timeout is needed.
+func newLBBuilderWithFallbackTimeout(fallbackTimeout time.Duration) balancer.Builder {
+	return &lbBuilder{
+		fallbackTimeout: fallbackTimeout,
+	}
+}
+
+type lbBuilder struct {
+	fallbackTimeout time.Duration
+}
+
+func (b *lbBuilder) Name() string {
+	return grpclbName
+}
+
+func (b *lbBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) balancer.Balancer {
+	// This generates a manual resolver builder with a fixed scheme. This
+	// scheme will be used to dial to remote LB, so we can send filtered
+	// address updates to remote LB ClientConn using this manual resolver.
+	mr := manual.NewBuilderWithScheme("grpclb-internal")
+	// ResolveNow() on this manual resolver is forwarded to the parent
+	// ClientConn, so when grpclb client loses contact with the remote balancer,
+	// the parent ClientConn's resolver will re-resolve.
+	mr.ResolveNowCallback = cc.ResolveNow
+
+	lb := &lbBalancer{
+		cc:              newLBCacheClientConn(cc),
+		dialTarget:      opt.Target.Endpoint(),
+		target:          opt.Target.Endpoint(),
+		opt:             opt,
+		fallbackTimeout: b.fallbackTimeout,
+		doneCh:          make(chan struct{}),
+
+		manualResolver: mr,
+		subConns:       make(map[resolver.Address]balancer.SubConn),
+		scStates:       make(map[balancer.SubConn]connectivity.State),
+		picker:         base.NewErrPicker(balancer.ErrNoSubConnAvailable),
+		clientStats:    newRPCStats(),
+		backoff:        backoff.DefaultExponential, // TODO: make backoff configurable.
+	}
+	lb.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[grpclb %p] ", lb))
+
+	var err error
+	if opt.CredsBundle != nil {
+		lb.grpclbClientConnCreds, err = opt.CredsBundle.NewWithMode(internal.CredsBundleModeBalancer)
+		if err != nil {
+			lb.logger.Warningf("Failed to create credentials used for connecting to grpclb: %v", err)
+		}
+		lb.grpclbBackendCreds, err = opt.CredsBundle.NewWithMode(internal.CredsBundleModeBackendFromBalancer)
+		if err != nil {
+			lb.logger.Warningf("Failed to create credentials used for connecting to backends returned by grpclb: %v", err)
+		}
+	}
+
+	return lb
+}
+
+type lbBalancer struct {
+	cc         *lbCacheClientConn
+	dialTarget string // user's dial target
+	target     string // same as dialTarget unless overridden in service config
+	opt        balancer.BuildOptions
+	logger     *internalgrpclog.PrefixLogger
+
+	usePickFirst bool
+
+	// grpclbClientConnCreds is the creds bundle to be used to connect to grpclb
+	// servers. If it's nil, use the TransportCredentials from BuildOptions
+	// instead.
+	grpclbClientConnCreds credentials.Bundle
+	// grpclbBackendCreds is the creds bundle to be used for addresses that are
+	// returned by grpclb server. If it's nil, don't set anything when creating
+	// SubConns.
+	grpclbBackendCreds credentials.Bundle
+
+	fallbackTimeout time.Duration
+	doneCh          chan struct{}
+
+	// manualResolver is used in the remote LB ClientConn inside grpclb. When
+	// resolved address updates are received by grpclb, filtered updates will be
+	// sent to remote LB ClientConn through this resolver.
+	manualResolver *manual.Resolver
+	// The ClientConn to talk to the remote balancer.
+	ccRemoteLB *remoteBalancerCCWrapper
+	// backoff for calling remote balancer.
+	backoff backoff.Strategy
+
+	// Support client side load reporting. Each picker gets a reference to this,
+	// and will update its content.
+	clientStats *rpcStats
+
+	mu sync.Mutex // guards everything following.
+	// The full server list including drops, used to check if the newly received
+	// serverList contains anything new. Each generate picker will also have
+	// reference to this list to do the first layer pick.
+	fullServerList []*lbpb.Server
+	// Backend addresses. It's kept so the addresses are available when
+	// switching between round_robin and pickfirst.
+	backendAddrs []resolver.Address
+	// All backends addresses, with metadata set to nil. This list contains all
+	// backend addresses in the same order and with the same duplicates as in
+	// serverlist. When generating picker, a SubConn slice with the same order
+	// but with only READY SCs will be generated.
+	backendAddrsWithoutMetadata []resolver.Address
+	// Roundrobin functionalities.
+	state    connectivity.State
+	subConns map[resolver.Address]balancer.SubConn   // Used to new/shutdown SubConn.
+	scStates map[balancer.SubConn]connectivity.State // Used to filter READY SubConns.
+	picker   balancer.Picker
+	// Support fallback to resolved backend addresses if there's no response
+	// from remote balancer within fallbackTimeout.
+	remoteBalancerConnected bool
+	serverListReceived      bool
+	inFallback              bool
+	// resolvedBackendAddrs is resolvedAddrs minus remote balancers. It's set
+	// when resolved address updates are received, and read in the goroutine
+	// handling fallback.
+	resolvedBackendAddrs []resolver.Address
+	connErr              error // the last connection error
+}
+
+// regeneratePicker takes a snapshot of the balancer, and generates a picker from
+// it. The picker
+//   - always returns ErrTransientFailure if the balancer is in TransientFailure,
+//   - does two layer roundrobin pick otherwise.
+//
+// Caller must hold lb.mu.
+func (lb *lbBalancer) regeneratePicker(resetDrop bool) {
+	if lb.state == connectivity.TransientFailure {
+		lb.picker = base.NewErrPicker(fmt.Errorf("all SubConns are in TransientFailure, last connection error: %v", lb.connErr))
+		return
+	}
+
+	if lb.state == connectivity.Connecting {
+		lb.picker = base.NewErrPicker(balancer.ErrNoSubConnAvailable)
+		return
+	}
+
+	var readySCs []balancer.SubConn
+	if lb.usePickFirst {
+		for _, sc := range lb.subConns {
+			readySCs = append(readySCs, sc)
+			break
+		}
+	} else {
+		for _, a := range lb.backendAddrsWithoutMetadata {
+			if sc, ok := lb.subConns[a]; ok {
+				if st, ok := lb.scStates[sc]; ok && st == connectivity.Ready {
+					readySCs = append(readySCs, sc)
+				}
+			}
+		}
+	}
+
+	if len(readySCs) <= 0 {
+		// If there's no ready SubConns, always re-pick. This is to avoid drops
+		// unless at least one SubConn is ready. Otherwise we may drop more
+		// often than want because of drops + re-picks(which become re-drops).
+		//
+		// This doesn't seem to be necessary after the connecting check above.
+		// Kept for safety.
+		lb.picker = base.NewErrPicker(balancer.ErrNoSubConnAvailable)
+		return
+	}
+	if lb.inFallback {
+		lb.picker = newRRPicker(readySCs)
+		return
+	}
+	if resetDrop {
+		lb.picker = newLBPicker(lb.fullServerList, readySCs, lb.clientStats)
+		return
+	}
+	prevLBPicker, ok := lb.picker.(*lbPicker)
+	if !ok {
+		lb.picker = newLBPicker(lb.fullServerList, readySCs, lb.clientStats)
+		return
+	}
+	prevLBPicker.updateReadySCs(readySCs)
+}
+
+// aggregateSubConnStats calculate the aggregated state of SubConns in
+// lb.SubConns. These SubConns are subconns in use (when switching between
+// fallback and grpclb). lb.scState contains states for all SubConns, including
+// those in cache (SubConns are cached for 10 seconds after shutdown).
+//
+//	The aggregated state is:
+//	- If at least one SubConn in Ready, the aggregated state is Ready;
+//	- Else if at least one SubConn in Connecting or IDLE, the aggregated state is Connecting;
+//	  - It's OK to consider IDLE as Connecting. SubConns never stay in IDLE,
+//	    they start to connect immediately. But there's a race between the overall
+//	    state is reported, and when the new SubConn state arrives. And SubConns
+//	    never go back to IDLE.
+//	- Else the aggregated state is TransientFailure.
+func (lb *lbBalancer) aggregateSubConnStates() connectivity.State {
+	var numConnecting uint64
+
+	for _, sc := range lb.subConns {
+		if state, ok := lb.scStates[sc]; ok {
+			switch state {
+			case connectivity.Ready:
+				return connectivity.Ready
+			case connectivity.Connecting, connectivity.Idle:
+				numConnecting++
+			}
+		}
+	}
+	if numConnecting > 0 {
+		return connectivity.Connecting
+	}
+	return connectivity.TransientFailure
+}
+
+// UpdateSubConnState is unused; NewSubConn's options always specifies
+// updateSubConnState as the listener.
+func (lb *lbBalancer) UpdateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) {
+	lb.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, scs)
+}
+
+func (lb *lbBalancer) updateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) {
+	s := scs.ConnectivityState
+	if lb.logger.V(2) {
+		lb.logger.Infof("SubConn state change: %p, %v", sc, s)
+	}
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	oldS, ok := lb.scStates[sc]
+	if !ok {
+		if lb.logger.V(2) {
+			lb.logger.Infof("Received state change for an unknown SubConn: %p, %v", sc, s)
+		}
+		return
+	}
+	lb.scStates[sc] = s
+	switch s {
+	case connectivity.Idle:
+		sc.Connect()
+	case connectivity.Shutdown:
+		// When an address was removed by resolver, b called Shutdown but kept
+		// the sc's state in scStates. Remove state for this sc here.
+		delete(lb.scStates, sc)
+	case connectivity.TransientFailure:
+		lb.connErr = scs.ConnectionError
+	}
+	// Force regenerate picker if
+	//  - this sc became ready from not-ready
+	//  - this sc became not-ready from ready
+	lb.updateStateAndPicker((oldS == connectivity.Ready) != (s == connectivity.Ready), false)
+
+	// Enter fallback when the aggregated state is not Ready and the connection
+	// to remote balancer is lost.
+	if lb.state != connectivity.Ready {
+		if !lb.inFallback && !lb.remoteBalancerConnected {
+			// Enter fallback.
+			lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
+		}
+	}
+}
+
+// updateStateAndPicker re-calculate the aggregated state, and regenerate picker
+// if overall state is changed.
+//
+// If forceRegeneratePicker is true, picker will be regenerated.
+func (lb *lbBalancer) updateStateAndPicker(forceRegeneratePicker bool, resetDrop bool) {
+	oldAggrState := lb.state
+	lb.state = lb.aggregateSubConnStates()
+	// Regenerate picker when one of the following happens:
+	//  - caller wants to regenerate
+	//  - the aggregated state changed
+	if forceRegeneratePicker || (lb.state != oldAggrState) {
+		lb.regeneratePicker(resetDrop)
+	}
+	var cc balancer.ClientConn = lb.cc
+	if lb.usePickFirst {
+		// Bypass the caching layer that would wrap the picker.
+		cc = lb.cc.ClientConn
+	}
+
+	cc.UpdateState(balancer.State{ConnectivityState: lb.state, Picker: lb.picker})
+}
+
+// fallbackToBackendsAfter blocks for fallbackTimeout and falls back to use
+// resolved backends (backends received from resolver, not from remote balancer)
+// if no connection to remote balancers was successful.
+func (lb *lbBalancer) fallbackToBackendsAfter(fallbackTimeout time.Duration) {
+	timer := time.NewTimer(fallbackTimeout)
+	defer timer.Stop()
+	select {
+	case <-timer.C:
+	case <-lb.doneCh:
+		return
+	}
+	lb.mu.Lock()
+	if lb.inFallback || lb.serverListReceived {
+		lb.mu.Unlock()
+		return
+	}
+	// Enter fallback.
+	lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
+	lb.mu.Unlock()
+}
+
+func (lb *lbBalancer) handleServiceConfig(gc *grpclbServiceConfig) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	// grpclb uses the user's dial target to populate the `Name` field of the
+	// `InitialLoadBalanceRequest` message sent to the remote balancer. But when
+	// grpclb is used a child policy in the context of RLS, we want the `Name`
+	// field to be populated with the value received from the RLS server. To
+	// support this use case, an optional "target_name" field has been added to
+	// the grpclb LB policy's config.  If specified, it overrides the name of
+	// the target to be sent to the remote balancer; if not, the target to be
+	// sent to the balancer will continue to be obtained from the target URI
+	// passed to the gRPC client channel. Whenever that target to be sent to the
+	// balancer is updated, we need to restart the stream to the balancer as
+	// this target is sent in the first message on the stream.
+	if gc != nil {
+		target := lb.dialTarget
+		if gc.ServiceName != "" {
+			target = gc.ServiceName
+		}
+		if target != lb.target {
+			lb.target = target
+			if lb.ccRemoteLB != nil {
+				lb.ccRemoteLB.cancelRemoteBalancerCall()
+			}
+		}
+	}
+
+	newUsePickFirst := childIsPickFirst(gc)
+	if lb.usePickFirst == newUsePickFirst {
+		return
+	}
+	if lb.logger.V(2) {
+		lb.logger.Infof("Switching mode. Is pick_first used for backends? %v", newUsePickFirst)
+	}
+	lb.refreshSubConns(lb.backendAddrs, lb.inFallback, newUsePickFirst)
+}
+
+func (lb *lbBalancer) ResolverError(error) {
+	// Ignore resolver errors.  GRPCLB is not selected unless the resolver
+	// works at least once.
+}
+
+func (lb *lbBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	if lb.logger.V(2) {
+		lb.logger.Infof("UpdateClientConnState: %s", pretty.ToJSON(ccs))
+	}
+	gc, _ := ccs.BalancerConfig.(*grpclbServiceConfig)
+	lb.handleServiceConfig(gc)
+
+	backendAddrs := ccs.ResolverState.Addresses
+
+	var remoteBalancerAddrs []resolver.Address
+	if sd := grpclbstate.Get(ccs.ResolverState); sd != nil {
+		// Override any balancer addresses provided via
+		// ccs.ResolverState.Addresses.
+		remoteBalancerAddrs = sd.BalancerAddresses
+	}
+
+	if len(backendAddrs)+len(remoteBalancerAddrs) == 0 {
+		// There should be at least one address, either grpclb server or
+		// fallback. Empty address is not valid.
+		return balancer.ErrBadResolverState
+	}
+
+	if len(remoteBalancerAddrs) == 0 {
+		if lb.ccRemoteLB != nil {
+			lb.ccRemoteLB.close()
+			lb.ccRemoteLB = nil
+		}
+	} else if lb.ccRemoteLB == nil {
+		// First time receiving resolved addresses, create a cc to remote
+		// balancers.
+		if err := lb.newRemoteBalancerCCWrapper(); err != nil {
+			return err
+		}
+		// Start the fallback goroutine.
+		go lb.fallbackToBackendsAfter(lb.fallbackTimeout)
+	}
+
+	if lb.ccRemoteLB != nil {
+		// cc to remote balancers uses lb.manualResolver. Send the updated remote
+		// balancer addresses to it through manualResolver.
+		lb.manualResolver.UpdateState(resolver.State{Addresses: remoteBalancerAddrs})
+	}
+
+	lb.mu.Lock()
+	lb.resolvedBackendAddrs = backendAddrs
+	if len(remoteBalancerAddrs) == 0 || lb.inFallback {
+		// If there's no remote balancer address in ClientConn update, grpclb
+		// enters fallback mode immediately.
+		//
+		// If a new update is received while grpclb is in fallback, update the
+		// list of backends being used to the new fallback backends.
+		lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
+	}
+	lb.mu.Unlock()
+	return nil
+}
+
+func (lb *lbBalancer) Close() {
+	select {
+	case <-lb.doneCh:
+		return
+	default:
+	}
+	close(lb.doneCh)
+	if lb.ccRemoteLB != nil {
+		lb.ccRemoteLB.close()
+	}
+	lb.cc.close()
+}
+
+func (lb *lbBalancer) ExitIdle() {}

balancer/grpclb/grpclb_config.go

@@ -0,0 +1,67 @@
+/*
+ *
+ * Copyright 2019 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 grpclb
+
+import (
+	"encoding/json"
+
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/balancer/roundrobin"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+const (
+	roundRobinName = roundrobin.Name
+	pickFirstName  = pickfirst.Name
+)
+
+type grpclbServiceConfig struct {
+	serviceconfig.LoadBalancingConfig
+	ChildPolicy *[]map[string]json.RawMessage
+	ServiceName string
+}
+
+func (b *lbBuilder) ParseConfig(lbConfig json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	ret := &grpclbServiceConfig{}
+	if err := json.Unmarshal(lbConfig, ret); err != nil {
+		return nil, err
+	}
+	return ret, nil
+}
+
+func childIsPickFirst(sc *grpclbServiceConfig) bool {
+	if sc == nil {
+		return false
+	}
+	childConfigs := sc.ChildPolicy
+	if childConfigs == nil {
+		return false
+	}
+	for _, childC := range *childConfigs {
+		// If round_robin exists before pick_first, return false
+		if _, ok := childC[roundRobinName]; ok {
+			return false
+		}
+		// If pick_first is before round_robin, return true
+		if _, ok := childC[pickFirstName]; ok {
+			return true
+		}
+	}
+	return false
+}

balancer/grpclb/grpclb_config_test.go

@@ -0,0 +1,123 @@
+/*
+ *
+ * Copyright 2019 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 grpclb
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+func (s) TestParse(t *testing.T) {
+	tests := []struct {
+		name    string
+		sc      string
+		want    serviceconfig.LoadBalancingConfig
+		wantErr bool
+	}{
+		{
+			name:    "empty",
+			sc:      "",
+			want:    nil,
+			wantErr: true,
+		},
+		{
+			name: "success1",
+			sc: `
+{
+	"childPolicy": [
+		{"pick_first":{}}
+	],
+	"serviceName": "foo-service"
+}`,
+			want: &grpclbServiceConfig{
+				ChildPolicy: &[]map[string]json.RawMessage{
+					{"pick_first": json.RawMessage("{}")},
+				},
+				ServiceName: "foo-service",
+			},
+		},
+		{
+			name: "success2",
+			sc: `
+{
+	"childPolicy": [
+		{"round_robin":{}},
+		{"pick_first":{}}
+	],
+	"serviceName": "foo-service"
+}`,
+			want: &grpclbServiceConfig{
+				ChildPolicy: &[]map[string]json.RawMessage{
+					{"round_robin": json.RawMessage("{}")},
+					{"pick_first": json.RawMessage("{}")},
+				},
+				ServiceName: "foo-service",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := (&lbBuilder{}).ParseConfig(json.RawMessage(tt.sc))
+			if (err != nil) != (tt.wantErr) {
+				t.Fatalf("ParseConfig(%q) returned error: %v, wantErr: %v", tt.sc, err, tt.wantErr)
+			}
+			if diff := cmp.Diff(tt.want, got); diff != "" {
+				t.Fatalf("ParseConfig(%q) returned unexpected difference (-want +got):\n%s", tt.sc, diff)
+			}
+		})
+	}
+}
+
+func (s) TestChildIsPickFirst(t *testing.T) {
+	tests := []struct {
+		name string
+		s    string
+		want bool
+	}{
+		{
+			name: "pickfirst_only",
+			s:    `{"childPolicy":[{"pick_first":{}}]}`,
+			want: true,
+		},
+		{
+			name: "pickfirst_before_rr",
+			s:    `{"childPolicy":[{"pick_first":{}},{"round_robin":{}}]}`,
+			want: true,
+		},
+		{
+			name: "rr_before_pickfirst",
+			s:    `{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}`,
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gc, err := (&lbBuilder{}).ParseConfig(json.RawMessage(tt.s))
+			if err != nil {
+				t.Fatalf("Parse(%v) = _, %v; want _, nil", tt.s, err)
+			}
+			if got := childIsPickFirst(gc.(*grpclbServiceConfig)); got != tt.want {
+				t.Errorf("childIsPickFirst() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

balancer/grpclb/grpclb_picker.go

@@ -0,0 +1,193 @@
+/*
+ *
+ * Copyright 2017 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 grpclb
+
+import (
+	rand "math/rand/v2"
+	"sync"
+	"sync/atomic"
+
+	"google.golang.org/grpc/balancer"
+	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+// rpcStats is same as lbpb.ClientStats, except that numCallsDropped is a map
+// instead of a slice.
+type rpcStats struct {
+	// Only access the following fields atomically.
+	numCallsStarted                        int64
+	numCallsFinished                       int64
+	numCallsFinishedWithClientFailedToSend int64
+	numCallsFinishedKnownReceived          int64
+
+	mu sync.Mutex
+	// map load_balance_token -> num_calls_dropped
+	numCallsDropped map[string]int64
+}
+
+func newRPCStats() *rpcStats {
+	return &rpcStats{
+		numCallsDropped: make(map[string]int64),
+	}
+}
+
+func isZeroStats(stats *lbpb.ClientStats) bool {
+	return len(stats.CallsFinishedWithDrop) == 0 &&
+		stats.NumCallsStarted == 0 &&
+		stats.NumCallsFinished == 0 &&
+		stats.NumCallsFinishedWithClientFailedToSend == 0 &&
+		stats.NumCallsFinishedKnownReceived == 0
+}
+
+// toClientStats converts rpcStats to lbpb.ClientStats, and clears rpcStats.
+func (s *rpcStats) toClientStats() *lbpb.ClientStats {
+	stats := &lbpb.ClientStats{
+		NumCallsStarted:                        atomic.SwapInt64(&s.numCallsStarted, 0),
+		NumCallsFinished:                       atomic.SwapInt64(&s.numCallsFinished, 0),
+		NumCallsFinishedWithClientFailedToSend: atomic.SwapInt64(&s.numCallsFinishedWithClientFailedToSend, 0),
+		NumCallsFinishedKnownReceived:          atomic.SwapInt64(&s.numCallsFinishedKnownReceived, 0),
+	}
+	s.mu.Lock()
+	dropped := s.numCallsDropped
+	s.numCallsDropped = make(map[string]int64)
+	s.mu.Unlock()
+	for token, count := range dropped {
+		stats.CallsFinishedWithDrop = append(stats.CallsFinishedWithDrop, &lbpb.ClientStatsPerToken{
+			LoadBalanceToken: token,
+			NumCalls:         count,
+		})
+	}
+	return stats
+}
+
+func (s *rpcStats) drop(token string) {
+	atomic.AddInt64(&s.numCallsStarted, 1)
+	s.mu.Lock()
+	s.numCallsDropped[token]++
+	s.mu.Unlock()
+	atomic.AddInt64(&s.numCallsFinished, 1)
+}
+
+func (s *rpcStats) failedToSend() {
+	atomic.AddInt64(&s.numCallsStarted, 1)
+	atomic.AddInt64(&s.numCallsFinishedWithClientFailedToSend, 1)
+	atomic.AddInt64(&s.numCallsFinished, 1)
+}
+
+func (s *rpcStats) knownReceived() {
+	atomic.AddInt64(&s.numCallsStarted, 1)
+	atomic.AddInt64(&s.numCallsFinishedKnownReceived, 1)
+	atomic.AddInt64(&s.numCallsFinished, 1)
+}
+
+// rrPicker does roundrobin on subConns. It's typically used when there's no
+// response from remote balancer, and grpclb falls back to the resolved
+// backends.
+//
+// It guaranteed that len(subConns) > 0.
+type rrPicker struct {
+	mu           sync.Mutex
+	subConns     []balancer.SubConn // The subConns that were READY when taking the snapshot.
+	subConnsNext int
+}
+
+func newRRPicker(readySCs []balancer.SubConn) *rrPicker {
+	return &rrPicker{
+		subConns:     readySCs,
+		subConnsNext: rand.IntN(len(readySCs)),
+	}
+}
+
+func (p *rrPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	sc := p.subConns[p.subConnsNext]
+	p.subConnsNext = (p.subConnsNext + 1) % len(p.subConns)
+	return balancer.PickResult{SubConn: sc}, nil
+}
+
+// lbPicker does two layers of picks:
+//
+// First layer: roundrobin on all servers in serverList, including drops and backends.
+// - If it picks a drop, the RPC will fail as being dropped.
+// - If it picks a backend, do a second layer pick to pick the real backend.
+//
+// Second layer: roundrobin on all READY backends.
+//
+// It's guaranteed that len(serverList) > 0.
+type lbPicker struct {
+	mu             sync.Mutex
+	serverList     []*lbpb.Server
+	serverListNext int
+	subConns       []balancer.SubConn // The subConns that were READY when taking the snapshot.
+	subConnsNext   int
+
+	stats *rpcStats
+}
+
+func newLBPicker(serverList []*lbpb.Server, readySCs []balancer.SubConn, stats *rpcStats) *lbPicker {
+	return &lbPicker{
+		serverList:   serverList,
+		subConns:     readySCs,
+		subConnsNext: rand.IntN(len(readySCs)),
+		stats:        stats,
+	}
+}
+
+func (p *lbPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	// Layer one roundrobin on serverList.
+	s := p.serverList[p.serverListNext]
+	p.serverListNext = (p.serverListNext + 1) % len(p.serverList)
+
+	// If it's a drop, return an error and fail the RPC.
+	if s.Drop {
+		p.stats.drop(s.LoadBalanceToken)
+		return balancer.PickResult{}, status.Errorf(codes.Unavailable, "request dropped by grpclb")
+	}
+
+	// If not a drop but there's no ready subConns.
+	if len(p.subConns) <= 0 {
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+	}
+
+	// Return the next ready subConn in the list, also collect rpc stats.
+	sc := p.subConns[p.subConnsNext]
+	p.subConnsNext = (p.subConnsNext + 1) % len(p.subConns)
+	done := func(info balancer.DoneInfo) {
+		if !info.BytesSent {
+			p.stats.failedToSend()
+		} else if info.BytesReceived {
+			p.stats.knownReceived()
+		}
+	}
+	return balancer.PickResult{SubConn: sc, Done: done}, nil
+}
+
+func (p *lbPicker) updateReadySCs(readySCs []balancer.SubConn) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	p.subConns = readySCs
+	p.subConnsNext = p.subConnsNext % len(readySCs)
+}

balancer/grpclb/grpclb_remote_balancer.go

@@ -0,0 +1,459 @@
+/*
+ *
+ * Copyright 2017 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 grpclb
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net"
+	"net/netip"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/backoff"
+	imetadata "google.golang.org/grpc/internal/metadata"
+	"google.golang.org/grpc/keepalive"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
+)
+
+func serverListEqual(a, b []*lbpb.Server) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i := 0; i < len(a); i++ {
+		if !proto.Equal(a[i], b[i]) {
+			return false
+		}
+	}
+	return true
+}
+
+// processServerList updates balancer's internal state, create/remove SubConns
+// and regenerates picker using the received serverList.
+func (lb *lbBalancer) processServerList(l *lbpb.ServerList) {
+	if lb.logger.V(2) {
+		lb.logger.Infof("Processing server list: %#v", l)
+	}
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	// Set serverListReceived to true so fallback will not take effect if it has
+	// not hit timeout.
+	lb.serverListReceived = true
+
+	// If the new server list == old server list, do nothing.
+	if serverListEqual(lb.fullServerList, l.Servers) {
+		if lb.logger.V(2) {
+			lb.logger.Infof("Ignoring new server list as it is the same as the previous one")
+		}
+		return
+	}
+	lb.fullServerList = l.Servers
+
+	var backendAddrs []resolver.Address
+	for i, s := range l.Servers {
+		if s.Drop {
+			continue
+		}
+
+		md := metadata.Pairs(lbTokenKey, s.LoadBalanceToken)
+		var ipStr string
+		if ip, ok := netip.AddrFromSlice(s.IpAddress); ok {
+			ipStr = ip.String()
+		} else {
+			ipStr = fmt.Sprintf("? %x", s.IpAddress)
+		}
+		addr := imetadata.Set(resolver.Address{Addr: net.JoinHostPort(ipStr, fmt.Sprintf("%d", s.Port))}, md)
+		if lb.logger.V(2) {
+			lb.logger.Infof("Server list entry:|%d|, ipStr:|%s|, port:|%d|, load balancer token:|%v|", i, ipStr, s.Port, s.LoadBalanceToken)
+		}
+		backendAddrs = append(backendAddrs, addr)
+	}
+
+	// Call refreshSubConns to create/remove SubConns.  If we are in fallback,
+	// this is also exiting fallback.
+	lb.refreshSubConns(backendAddrs, false, lb.usePickFirst)
+}
+
+// refreshSubConns creates/removes SubConns with backendAddrs, and refreshes
+// balancer state and picker.
+//
+// Caller must hold lb.mu.
+func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback bool, pickFirst bool) {
+	opts := balancer.NewSubConnOptions{}
+	if !fallback {
+		opts.CredsBundle = lb.grpclbBackendCreds
+	}
+
+	lb.backendAddrs = backendAddrs
+	lb.backendAddrsWithoutMetadata = nil
+
+	fallbackModeChanged := lb.inFallback != fallback
+	lb.inFallback = fallback
+	if fallbackModeChanged && lb.inFallback {
+		// Clear previous received list when entering fallback, so if the server
+		// comes back and sends the same list again, the new addresses will be
+		// used.
+		lb.fullServerList = nil
+	}
+
+	balancingPolicyChanged := lb.usePickFirst != pickFirst
+	lb.usePickFirst = pickFirst
+
+	if fallbackModeChanged || balancingPolicyChanged {
+		// Remove all SubConns when switching balancing policy or switching
+		// fallback mode.
+		//
+		// For fallback mode switching with pickfirst, we want to recreate the
+		// SubConn because the creds could be different.
+		for a, sc := range lb.subConns {
+			sc.Shutdown()
+			delete(lb.subConns, a)
+		}
+	}
+
+	if lb.usePickFirst {
+		var (
+			scKey resolver.Address
+			sc    balancer.SubConn
+		)
+		for scKey, sc = range lb.subConns {
+			break
+		}
+		if sc != nil {
+			if len(backendAddrs) == 0 {
+				sc.Shutdown()
+				delete(lb.subConns, scKey)
+				return
+			}
+			lb.cc.ClientConn.UpdateAddresses(sc, backendAddrs)
+			sc.Connect()
+			return
+		}
+		opts.StateListener = func(scs balancer.SubConnState) { lb.updateSubConnState(sc, scs) }
+		// This bypasses the cc wrapper with SubConn cache.
+		sc, err := lb.cc.ClientConn.NewSubConn(backendAddrs, opts)
+		if err != nil {
+			lb.logger.Warningf("Failed to create new SubConn: %v", err)
+			return
+		}
+		sc.Connect()
+		lb.subConns[backendAddrs[0]] = sc
+		lb.scStates[sc] = connectivity.Idle
+		return
+	}
+
+	// addrsSet is the set converted from backendAddrsWithoutMetadata, it's used to quick
+	// lookup for an address.
+	addrsSet := make(map[resolver.Address]struct{})
+	// Create new SubConns.
+	for _, addr := range backendAddrs {
+		addrWithoutAttrs := addr
+		addrWithoutAttrs.Attributes = nil
+		addrsSet[addrWithoutAttrs] = struct{}{}
+		lb.backendAddrsWithoutMetadata = append(lb.backendAddrsWithoutMetadata, addrWithoutAttrs)
+
+		if _, ok := lb.subConns[addrWithoutAttrs]; !ok {
+			// Use addrWithMD to create the SubConn.
+			var sc balancer.SubConn
+			opts.StateListener = func(scs balancer.SubConnState) { lb.updateSubConnState(sc, scs) }
+			sc, err := lb.cc.NewSubConn([]resolver.Address{addr}, opts)
+			if err != nil {
+				lb.logger.Warningf("Failed to create new SubConn: %v", err)
+				continue
+			}
+			lb.subConns[addrWithoutAttrs] = sc // Use the addr without MD as key for the map.
+			if _, ok := lb.scStates[sc]; !ok {
+				// Only set state of new sc to IDLE. The state could already be
+				// READY for cached SubConns.
+				lb.scStates[sc] = connectivity.Idle
+			}
+			sc.Connect()
+		}
+	}
+
+	for a, sc := range lb.subConns {
+		// a was removed by resolver.
+		if _, ok := addrsSet[a]; !ok {
+			sc.Shutdown()
+			delete(lb.subConns, a)
+			// Keep the state of this sc in b.scStates until sc's state becomes Shutdown.
+			// The entry will be deleted in UpdateSubConnState.
+		}
+	}
+
+	// Regenerate and update picker after refreshing subconns because with
+	// cache, even if SubConn was newed/removed, there might be no state
+	// changes (the subconn will be kept in cache, not actually
+	// newed/removed).
+	lb.updateStateAndPicker(true, true)
+}
+
+type remoteBalancerCCWrapper struct {
+	cc      *grpc.ClientConn
+	lb      *lbBalancer
+	backoff backoff.Strategy
+	done    chan struct{}
+
+	streamMu     sync.Mutex
+	streamCancel func()
+
+	// waitgroup to wait for all goroutines to exit.
+	wg sync.WaitGroup
+}
+
+func (lb *lbBalancer) newRemoteBalancerCCWrapper() error {
+	var dopts []grpc.DialOption
+	if creds := lb.opt.DialCreds; creds != nil {
+		dopts = append(dopts, grpc.WithTransportCredentials(creds))
+	} else if bundle := lb.grpclbClientConnCreds; bundle != nil {
+		dopts = append(dopts, grpc.WithCredentialsBundle(bundle))
+	} else {
+		dopts = append(dopts, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	}
+	if lb.opt.Dialer != nil {
+		dopts = append(dopts, grpc.WithContextDialer(lb.opt.Dialer))
+	}
+	if lb.opt.CustomUserAgent != "" {
+		dopts = append(dopts, grpc.WithUserAgent(lb.opt.CustomUserAgent))
+	}
+	// Explicitly set pickfirst as the balancer.
+	dopts = append(dopts, grpc.WithDefaultServiceConfig(`{"loadBalancingPolicy":"pick_first"}`))
+	dopts = append(dopts, grpc.WithResolvers(lb.manualResolver))
+	dopts = append(dopts, grpc.WithChannelzParentID(lb.opt.ChannelzParent))
+
+	// Enable Keepalive for grpclb client.
+	dopts = append(dopts, grpc.WithKeepaliveParams(keepalive.ClientParameters{
+		Time:                20 * time.Second,
+		Timeout:             10 * time.Second,
+		PermitWithoutStream: true,
+	}))
+
+	// The dial target is not important.
+	//
+	// The grpclb server addresses will set field ServerName, and creds will
+	// receive ServerName as authority.
+	target := lb.manualResolver.Scheme() + ":///grpclb.subClientConn"
+	cc, err := grpc.NewClient(target, dopts...)
+	if err != nil {
+		return fmt.Errorf("grpc.NewClient(%s): %v", target, err)
+	}
+	cc.Connect()
+	ccw := &remoteBalancerCCWrapper{
+		cc:      cc,
+		lb:      lb,
+		backoff: lb.backoff,
+		done:    make(chan struct{}),
+	}
+	lb.ccRemoteLB = ccw
+	ccw.wg.Add(1)
+	go ccw.watchRemoteBalancer()
+	return nil
+}
+
+// close closed the ClientConn to remote balancer, and waits until all
+// goroutines to finish.
+func (ccw *remoteBalancerCCWrapper) close() {
+	close(ccw.done)
+	ccw.cc.Close()
+	ccw.wg.Wait()
+}
+
+func (ccw *remoteBalancerCCWrapper) readServerList(s *balanceLoadClientStream) error {
+	for {
+		reply, err := s.Recv()
+		if err != nil {
+			if err == io.EOF {
+				return errServerTerminatedConnection
+			}
+			return fmt.Errorf("grpclb: failed to recv server list: %v", err)
+		}
+		if serverList := reply.GetServerList(); serverList != nil {
+			ccw.lb.processServerList(serverList)
+		}
+		if reply.GetFallbackResponse() != nil {
+			// Eagerly enter fallback
+			ccw.lb.mu.Lock()
+			ccw.lb.refreshSubConns(ccw.lb.resolvedBackendAddrs, true, ccw.lb.usePickFirst)
+			ccw.lb.mu.Unlock()
+		}
+	}
+}
+
+func (ccw *remoteBalancerCCWrapper) sendLoadReport(s *balanceLoadClientStream, interval time.Duration) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+	lastZero := false
+	for {
+		select {
+		case <-ticker.C:
+		case <-s.Context().Done():
+			return
+		}
+		stats := ccw.lb.clientStats.toClientStats()
+		zero := isZeroStats(stats)
+		if zero && lastZero {
+			// Quash redundant empty load reports.
+			continue
+		}
+		lastZero = zero
+		t := time.Now()
+		stats.Timestamp = &timestamppb.Timestamp{
+			Seconds: t.Unix(),
+			Nanos:   int32(t.Nanosecond()),
+		}
+		if err := s.Send(&lbpb.LoadBalanceRequest{
+			LoadBalanceRequestType: &lbpb.LoadBalanceRequest_ClientStats{
+				ClientStats: stats,
+			},
+		}); err != nil {
+			return
+		}
+	}
+}
+
+func (ccw *remoteBalancerCCWrapper) callRemoteBalancer(ctx context.Context) (backoff bool, _ error) {
+	lbClient := &loadBalancerClient{cc: ccw.cc}
+	stream, err := lbClient.BalanceLoad(ctx, grpc.WaitForReady(true))
+	if err != nil {
+		return true, fmt.Errorf("grpclb: failed to perform RPC to the remote balancer: %v", err)
+	}
+	ccw.lb.mu.Lock()
+	ccw.lb.remoteBalancerConnected = true
+	ccw.lb.mu.Unlock()
+
+	// grpclb handshake on the stream.
+	initReq := &lbpb.LoadBalanceRequest{
+		LoadBalanceRequestType: &lbpb.LoadBalanceRequest_InitialRequest{
+			InitialRequest: &lbpb.InitialLoadBalanceRequest{
+				Name: ccw.lb.target,
+			},
+		},
+	}
+	if err := stream.Send(initReq); err != nil {
+		return true, fmt.Errorf("grpclb: failed to send init request: %v", err)
+	}
+	reply, err := stream.Recv()
+	if err != nil {
+		return true, fmt.Errorf("grpclb: failed to recv init response: %v", err)
+	}
+	initResp := reply.GetInitialResponse()
+	if initResp == nil {
+		return true, fmt.Errorf("grpclb: reply from remote balancer did not include initial response")
+	}
+
+	ccw.wg.Add(1)
+	go func() {
+		defer ccw.wg.Done()
+		if d := convertDuration(initResp.ClientStatsReportInterval); d > 0 {
+			ccw.sendLoadReport(stream, d)
+		}
+	}()
+	// No backoff if init req/resp handshake was successful.
+	return false, ccw.readServerList(stream)
+}
+
+// cancelRemoteBalancerCall cancels the context used by the stream to the remote
+// balancer. watchRemoteBalancer() takes care of restarting this call after the
+// stream fails.
+func (ccw *remoteBalancerCCWrapper) cancelRemoteBalancerCall() {
+	ccw.streamMu.Lock()
+	if ccw.streamCancel != nil {
+		ccw.streamCancel()
+		ccw.streamCancel = nil
+	}
+	ccw.streamMu.Unlock()
+}
+
+func (ccw *remoteBalancerCCWrapper) watchRemoteBalancer() {
+	defer func() {
+		ccw.wg.Done()
+		ccw.streamMu.Lock()
+		if ccw.streamCancel != nil {
+			// This is to make sure that we don't leak the context when we are
+			// directly returning from inside of the below `for` loop.
+			ccw.streamCancel()
+			ccw.streamCancel = nil
+		}
+		ccw.streamMu.Unlock()
+	}()
+
+	var retryCount int
+	var ctx context.Context
+	for {
+		ccw.streamMu.Lock()
+		if ccw.streamCancel != nil {
+			ccw.streamCancel()
+			ccw.streamCancel = nil
+		}
+		ctx, ccw.streamCancel = context.WithCancel(context.Background())
+		ccw.streamMu.Unlock()
+
+		doBackoff, err := ccw.callRemoteBalancer(ctx)
+		select {
+		case <-ccw.done:
+			return
+		default:
+			if err != nil {
+				if err == errServerTerminatedConnection {
+					ccw.lb.logger.Infof("Call to remote balancer failed: %v", err)
+				} else {
+					ccw.lb.logger.Warningf("Call to remote balancer failed: %v", err)
+				}
+			}
+		}
+		// Trigger a re-resolve when the stream errors.
+		ccw.lb.cc.ClientConn.ResolveNow(resolver.ResolveNowOptions{})
+
+		ccw.lb.mu.Lock()
+		ccw.lb.remoteBalancerConnected = false
+		ccw.lb.fullServerList = nil
+		// Enter fallback when connection to remote balancer is lost, and the
+		// aggregated state is not Ready.
+		if !ccw.lb.inFallback && ccw.lb.state != connectivity.Ready {
+			// Entering fallback.
+			ccw.lb.refreshSubConns(ccw.lb.resolvedBackendAddrs, true, ccw.lb.usePickFirst)
+		}
+		ccw.lb.mu.Unlock()
+
+		if !doBackoff {
+			retryCount = 0
+			continue
+		}
+
+		timer := time.NewTimer(ccw.backoff.Backoff(retryCount)) // Copy backoff
+		select {
+		case <-timer.C:
+		case <-ccw.done:
+			timer.Stop()
+			return
+		}
+		retryCount++
+	}
+}

balancer/grpclb/grpclb_util.go

@@ -0,0 +1,173 @@
+/*
+ *
+ * Copyright 2016 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 grpclb
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/resolver"
+)
+
+const subConnCacheTime = time.Second * 10
+
+// lbCacheClientConn is a wrapper balancer.ClientConn with a SubConn cache.
+// SubConns will be kept in cache for subConnCacheTime before being shut down.
+//
+// Its NewSubconn and SubConn.Shutdown methods are updated to do cache first.
+type lbCacheClientConn struct {
+	balancer.ClientConn
+
+	timeout time.Duration
+
+	mu sync.Mutex
+	// subConnCache only keeps subConns that are being deleted.
+	subConnCache  map[resolver.Address]*subConnCacheEntry
+	subConnToAddr map[balancer.SubConn]resolver.Address
+}
+
+type subConnCacheEntry struct {
+	sc balancer.SubConn
+
+	cancel        func()
+	abortDeleting bool
+}
+
+func newLBCacheClientConn(cc balancer.ClientConn) *lbCacheClientConn {
+	return &lbCacheClientConn{
+		ClientConn:    cc,
+		timeout:       subConnCacheTime,
+		subConnCache:  make(map[resolver.Address]*subConnCacheEntry),
+		subConnToAddr: make(map[balancer.SubConn]resolver.Address),
+	}
+}
+
+func (ccc *lbCacheClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	if len(addrs) != 1 {
+		return nil, fmt.Errorf("grpclb calling NewSubConn with addrs of length %v", len(addrs))
+	}
+	addrWithoutAttrs := addrs[0]
+	addrWithoutAttrs.Attributes = nil
+
+	ccc.mu.Lock()
+	defer ccc.mu.Unlock()
+	if entry, ok := ccc.subConnCache[addrWithoutAttrs]; ok {
+		// If entry is in subConnCache, the SubConn was being deleted.
+		// cancel function will never be nil.
+		entry.cancel()
+		delete(ccc.subConnCache, addrWithoutAttrs)
+		return entry.sc, nil
+	}
+
+	scNew, err := ccc.ClientConn.NewSubConn(addrs, opts)
+	if err != nil {
+		return nil, err
+	}
+	scNew = &lbCacheSubConn{SubConn: scNew, ccc: ccc}
+
+	ccc.subConnToAddr[scNew] = addrWithoutAttrs
+	return scNew, nil
+}
+
+func (ccc *lbCacheClientConn) RemoveSubConn(sc balancer.SubConn) {
+	logger.Errorf("RemoveSubConn(%v) called unexpectedly", sc)
+}
+
+type lbCacheSubConn struct {
+	balancer.SubConn
+	ccc *lbCacheClientConn
+}
+
+func (sc *lbCacheSubConn) Shutdown() {
+	ccc := sc.ccc
+	ccc.mu.Lock()
+	defer ccc.mu.Unlock()
+	addr, ok := ccc.subConnToAddr[sc]
+	if !ok {
+		return
+	}
+
+	if entry, ok := ccc.subConnCache[addr]; ok {
+		if entry.sc != sc {
+			// This could happen if NewSubConn was called multiple times for
+			// the same address, and those SubConns are all shut down. We
+			// remove sc immediately here.
+			delete(ccc.subConnToAddr, sc)
+			sc.SubConn.Shutdown()
+		}
+		return
+	}
+
+	entry := &subConnCacheEntry{
+		sc: sc,
+	}
+	ccc.subConnCache[addr] = entry
+
+	timer := time.AfterFunc(ccc.timeout, func() {
+		ccc.mu.Lock()
+		defer ccc.mu.Unlock()
+		if entry.abortDeleting {
+			return
+		}
+		sc.SubConn.Shutdown()
+		delete(ccc.subConnToAddr, sc)
+		delete(ccc.subConnCache, addr)
+	})
+	entry.cancel = func() {
+		if !timer.Stop() {
+			// If stop was not successful, the timer has fired (this can only
+			// happen in a race). But the deleting function is blocked on ccc.mu
+			// because the mutex was held by the caller of this function.
+			//
+			// Set abortDeleting to true to abort the deleting function. When
+			// the lock is released, the deleting function will acquire the
+			// lock, check the value of abortDeleting and return.
+			entry.abortDeleting = true
+		}
+	}
+}
+
+func (ccc *lbCacheClientConn) UpdateState(s balancer.State) {
+	s.Picker = &lbCachePicker{Picker: s.Picker}
+	ccc.ClientConn.UpdateState(s)
+}
+
+func (ccc *lbCacheClientConn) close() {
+	ccc.mu.Lock()
+	defer ccc.mu.Unlock()
+	// Only cancel all existing timers. There's no need to shut down SubConns.
+	for _, entry := range ccc.subConnCache {
+		entry.cancel()
+	}
+}
+
+type lbCachePicker struct {
+	balancer.Picker
+}
+
+func (cp *lbCachePicker) Pick(i balancer.PickInfo) (balancer.PickResult, error) {
+	res, err := cp.Picker.Pick(i)
+	if err != nil {
+		return res, err
+	}
+	res.SubConn = res.SubConn.(*lbCacheSubConn).SubConn
+	return res, nil
+}

balancer/grpclb/grpclb_util_test.go

@@ -0,0 +1,267 @@
+/*
+ *
+ * Copyright 2018 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 grpclb
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/resolver"
+)
+
+type mockSubConn struct {
+	balancer.SubConn
+	mcc *mockClientConn
+}
+
+func (msc *mockSubConn) Shutdown() {
+	msc.mcc.mu.Lock()
+	defer msc.mcc.mu.Unlock()
+	delete(msc.mcc.subConns, msc)
+}
+
+type mockClientConn struct {
+	balancer.ClientConn
+
+	mu       sync.Mutex
+	subConns map[balancer.SubConn]resolver.Address
+}
+
+func newMockClientConn() *mockClientConn {
+	return &mockClientConn{
+		subConns: make(map[balancer.SubConn]resolver.Address),
+	}
+}
+
+func (mcc *mockClientConn) NewSubConn(addrs []resolver.Address, _ balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	sc := &mockSubConn{mcc: mcc}
+	mcc.mu.Lock()
+	defer mcc.mu.Unlock()
+	mcc.subConns[sc] = addrs[0]
+	return sc, nil
+}
+
+func (mcc *mockClientConn) RemoveSubConn(sc balancer.SubConn) {
+	panic(fmt.Sprintf("RemoveSubConn(%v) called unexpectedly", sc))
+}
+
+const testCacheTimeout = 100 * time.Millisecond
+
+func checkMockCC(mcc *mockClientConn, scLen int) error {
+	mcc.mu.Lock()
+	defer mcc.mu.Unlock()
+	if len(mcc.subConns) != scLen {
+		return fmt.Errorf("mcc = %+v, want len(mcc.subConns) = %v", mcc.subConns, scLen)
+	}
+	return nil
+}
+
+func checkCacheCC(ccc *lbCacheClientConn, sccLen, sctaLen int) error {
+	ccc.mu.Lock()
+	defer ccc.mu.Unlock()
+	if len(ccc.subConnCache) != sccLen {
+		return fmt.Errorf("ccc = %+v, want len(ccc.subConnCache) = %v", ccc.subConnCache, sccLen)
+	}
+	if len(ccc.subConnToAddr) != sctaLen {
+		return fmt.Errorf("ccc = %+v, want len(ccc.subConnToAddr) = %v", ccc.subConnToAddr, sctaLen)
+	}
+	return nil
+}
+
+// Test that SubConn won't be immediately shut down.
+func (s) TestLBCacheClientConnExpire(t *testing.T) {
+	mcc := newMockClientConn()
+	if err := checkMockCC(mcc, 0); err != nil {
+		t.Fatal(err)
+	}
+
+	ccc := newLBCacheClientConn(mcc)
+	ccc.timeout = testCacheTimeout
+	if err := checkCacheCC(ccc, 0, 0); err != nil {
+		t.Fatal(err)
+	}
+
+	sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
+	// One subconn in MockCC.
+	if err := checkMockCC(mcc, 1); err != nil {
+		t.Fatal(err)
+	}
+	// No subconn being deleted, and one in CacheCC.
+	if err := checkCacheCC(ccc, 0, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	sc.Shutdown()
+	// One subconn in MockCC before timeout.
+	if err := checkMockCC(mcc, 1); err != nil {
+		t.Fatal(err)
+	}
+	// One subconn being deleted, and one in CacheCC.
+	if err := checkCacheCC(ccc, 1, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	// Should all become empty after timeout.
+	var err error
+	for i := 0; i < 2; i++ {
+		time.Sleep(testCacheTimeout)
+		err = checkMockCC(mcc, 0)
+		if err != nil {
+			continue
+		}
+		err = checkCacheCC(ccc, 0, 0)
+		if err != nil {
+			continue
+		}
+	}
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+// Test that NewSubConn with the same address of a SubConn being shut down will
+// reuse the SubConn and cancel the removing.
+func (s) TestLBCacheClientConnReuse(t *testing.T) {
+	mcc := newMockClientConn()
+	if err := checkMockCC(mcc, 0); err != nil {
+		t.Fatal(err)
+	}
+
+	ccc := newLBCacheClientConn(mcc)
+	ccc.timeout = testCacheTimeout
+	if err := checkCacheCC(ccc, 0, 0); err != nil {
+		t.Fatal(err)
+	}
+
+	sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
+	// One subconn in MockCC.
+	if err := checkMockCC(mcc, 1); err != nil {
+		t.Fatal(err)
+	}
+	// No subconn being deleted, and one in CacheCC.
+	if err := checkCacheCC(ccc, 0, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	sc.Shutdown()
+	// One subconn in MockCC before timeout.
+	if err := checkMockCC(mcc, 1); err != nil {
+		t.Fatal(err)
+	}
+	// One subconn being deleted, and one in CacheCC.
+	if err := checkCacheCC(ccc, 1, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	// Recreate the old subconn, this should cancel the deleting process.
+	sc, _ = ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
+	// One subconn in MockCC.
+	if err := checkMockCC(mcc, 1); err != nil {
+		t.Fatal(err)
+	}
+	// No subconn being deleted, and one in CacheCC.
+	if err := checkCacheCC(ccc, 0, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	var err error
+	// Should not become empty after 2*timeout.
+	time.Sleep(2 * testCacheTimeout)
+	err = checkMockCC(mcc, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = checkCacheCC(ccc, 0, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Call Shutdown again, will delete after timeout.
+	sc.Shutdown()
+	// One subconn in MockCC before timeout.
+	if err := checkMockCC(mcc, 1); err != nil {
+		t.Fatal(err)
+	}
+	// One subconn being deleted, and one in CacheCC.
+	if err := checkCacheCC(ccc, 1, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	// Should all become empty after timeout.
+	for i := 0; i < 2; i++ {
+		time.Sleep(testCacheTimeout)
+		err = checkMockCC(mcc, 0)
+		if err != nil {
+			continue
+		}
+		err = checkCacheCC(ccc, 0, 0)
+		if err != nil {
+			continue
+		}
+	}
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+// Test that if the timer to shut down a SubConn fires at the same time
+// NewSubConn cancels the timer, it doesn't cause deadlock.
+func (s) TestLBCache_ShutdownTimer_New_Race(t *testing.T) {
+	mcc := newMockClientConn()
+	if err := checkMockCC(mcc, 0); err != nil {
+		t.Fatal(err)
+	}
+
+	ccc := newLBCacheClientConn(mcc)
+	ccc.timeout = time.Nanosecond
+	if err := checkCacheCC(ccc, 0, 0); err != nil {
+		t.Fatal(err)
+	}
+
+	sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
+	// One subconn in MockCC.
+	if err := checkMockCC(mcc, 1); err != nil {
+		t.Fatal(err)
+	}
+	// No subconn being deleted, and one in CacheCC.
+	if err := checkCacheCC(ccc, 0, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	done := make(chan struct{})
+
+	go func() {
+		for i := 0; i < 1000; i++ {
+			// Shutdown starts a timer with 1 ns timeout, the NewSubConn will
+			// race with the timer.
+			sc.Shutdown()
+			sc, _ = ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
+		}
+		close(done)
+	}()
+
+	select {
+	case <-time.After(time.Second):
+		t.Fatalf("Test didn't finish within 1 second. Deadlock")
+	case <-done:
+	}
+}

balancer/grpclb/state/state.go

@@ -0,0 +1,51 @@
+/*
+ *
+ * Copyright 2020 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 state declares grpclb types to be set by resolvers wishing to pass
+// information to grpclb via resolver.State Attributes.
+package state
+
+import (
+	"google.golang.org/grpc/resolver"
+)
+
+// keyType is the key to use for storing State in Attributes.
+type keyType string
+
+const key = keyType("grpc.grpclb.state")
+
+// State contains gRPCLB-relevant data passed from the name resolver.
+type State struct {
+	// BalancerAddresses contains the remote load balancer address(es).  If
+	// set, overrides any resolver-provided addresses with Type of GRPCLB.
+	BalancerAddresses []resolver.Address
+}
+
+// Set returns a copy of the provided state with attributes containing s.  s's
+// data should not be mutated after calling Set.
+func Set(state resolver.State, s *State) resolver.State {
+	state.Attributes = state.Attributes.WithValue(key, s)
+	return state
+}
+
+// Get returns the grpclb State in the resolver.State, or nil if not present.
+// The returned data should not be mutated.
+func Get(state resolver.State) *State {
+	s, _ := state.Attributes.Value(key).(*State)
+	return s
+}

balancer/lazy/lazy.go

@@ -0,0 +1,157 @@
+/*
+ *
+ * Copyright 2025 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 lazy contains a load balancer that starts in IDLE instead of
+// CONNECTING. Once it starts connecting, it instantiates its delegate.
+//
+// # Experimental
+//
+// Notice: This package is EXPERIMENTAL and may be changed or removed in a
+// later release.
+package lazy
+
+import (
+	"fmt"
+	"sync"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/resolver"
+
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+)
+
+var (
+	logger = grpclog.Component("lazy-lb")
+)
+
+const (
+	logPrefix = "[lazy-lb %p] "
+)
+
+// 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 is the constructor for the lazy balancer.
+func NewBalancer(cc balancer.ClientConn, bOpts balancer.BuildOptions, childBuilder ChildBuilderFunc) balancer.Balancer {
+	b := &lazyBalancer{
+		cc:           cc,
+		buildOptions: bOpts,
+		childBuilder: childBuilder,
+	}
+	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
+	cc.UpdateState(balancer.State{
+		ConnectivityState: connectivity.Idle,
+		Picker: &idlePicker{exitIdle: sync.OnceFunc(func() {
+			// Call ExitIdle in a new goroutine to avoid deadlocks while calling
+			// back into the channel synchronously.
+			go b.ExitIdle()
+		})},
+	})
+	return b
+}
+
+type lazyBalancer struct {
+	// The following fields are initialized at build time and read-only after
+	// that and therefore do not need to be guarded by a mutex.
+	cc           balancer.ClientConn
+	buildOptions balancer.BuildOptions
+	logger       *internalgrpclog.PrefixLogger
+	childBuilder ChildBuilderFunc
+
+	// The following fields are accessed while handling calls to the idlePicker
+	// and when handling ClientConn state updates. They are guarded by a mutex.
+
+	mu                    sync.Mutex
+	delegate              balancer.Balancer
+	latestClientConnState *balancer.ClientConnState
+	latestResolverError   error
+}
+
+func (lb *lazyBalancer) Close() {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+	if lb.delegate != nil {
+		lb.delegate.Close()
+		lb.delegate = nil
+	}
+}
+
+func (lb *lazyBalancer) ResolverError(err error) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+	if lb.delegate != nil {
+		lb.delegate.ResolverError(err)
+		return
+	}
+	lb.latestResolverError = err
+}
+
+func (lb *lazyBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+	if lb.delegate != nil {
+		return lb.delegate.UpdateClientConnState(ccs)
+	}
+
+	lb.latestClientConnState = &ccs
+	lb.latestResolverError = nil
+	return nil
+}
+
+// UpdateSubConnState implements balancer.Balancer.
+func (lb *lazyBalancer) UpdateSubConnState(balancer.SubConn, balancer.SubConnState) {
+	// UpdateSubConnState is deprecated.
+}
+
+func (lb *lazyBalancer) ExitIdle() {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+	if lb.delegate != nil {
+		lb.delegate.ExitIdle()
+		return
+	}
+	lb.delegate = lb.childBuilder(lb.cc, lb.buildOptions)
+	if lb.latestClientConnState != nil {
+		if err := lb.delegate.UpdateClientConnState(*lb.latestClientConnState); err != nil {
+			if err == balancer.ErrBadResolverState {
+				lb.cc.ResolveNow(resolver.ResolveNowOptions{})
+			} else {
+				lb.logger.Warningf("Error from child policy on receiving initial state: %v", err)
+			}
+		}
+		lb.latestClientConnState = nil
+	}
+	if lb.latestResolverError != nil {
+		lb.delegate.ResolverError(lb.latestResolverError)
+		lb.latestResolverError = nil
+	}
+}
+
+// idlePicker is used when the SubConn is IDLE and kicks the SubConn into
+// CONNECTING when Pick is called.
+type idlePicker struct {
+	exitIdle func()
+}
+
+func (i *idlePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	i.exitIdle()
+	return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+}

balancer/lazy/lazy_ext_test.go

@@ -0,0 +1,466 @@
+/*
+ *
+ * Copyright 2025 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 lazy_test
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/lazy"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/stubserver"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+)
+
+const (
+	// Default timeout for tests in this package.
+	defaultTestTimeout = 10 * time.Second
+	// Default short timeout, to be used when waiting for events which are not
+	// expected to happen.
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// TestExitIdle creates a lazy balancer than manages a pickfirst child. The test
+// calls Connect() on the channel which in turn calls ExitIdle on the lazy
+// balancer. The test verifies that the channel enters READY.
+func (s) TestExitIdle(t *testing.T) {
+	backend1 := stubserver.StartTestService(t, nil)
+	defer backend1.Stop()
+
+	mr := manual.NewBuilderWithScheme("e2e-test")
+	defer mr.Close()
+
+	mr.InitialState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend1.Address}}},
+		},
+	})
+
+	bf := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build)
+		},
+		ExitIdle: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.ExitIdle()
+		},
+		ResolverError: func(bd *stub.BalancerData, err error) {
+			bd.ChildBalancer.ResolverError(err)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+	}
+	stub.Register(t.Name(), bf)
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s": {}}]}`, t.Name())
+	opts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultServiceConfig(json),
+		grpc.WithResolvers(mr),
+	}
+	cc, err := grpc.NewClient(mr.Scheme()+":///", opts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient(_) failed: %v", err)
+	}
+	defer cc.Close()
+
+	cc.Connect()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// Send a resolver update to verify that the resolver state is correctly
+	// passed through to the leaf pickfirst balancer.
+	backend2 := stubserver.StartTestService(t, nil)
+	defer backend2.Stop()
+
+	mr.UpdateState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend2.Address}}},
+		},
+	})
+
+	var peer peer.Peer
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
+		t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
+	}
+	if got, want := peer.Addr.String(), backend2.Address; got != want {
+		t.Errorf("EmptyCall() went to unexpected backend: got %q, want %q", got, want)
+	}
+}
+
+// TestPicker creates a lazy balancer under a stub balancer which block all
+// calls to ExitIdle. This ensures the only way to trigger lazy to exit idle is
+// through the picker. The test makes an RPC and ensures it succeeds.
+func (s) TestPicker(t *testing.T) {
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	bf := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build)
+		},
+		ExitIdle: func(*stub.BalancerData) {
+			t.Log("Ignoring call to ExitIdle, calling the picker should make the lazy balancer exit IDLE state.")
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+	}
+
+	name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
+	stub.Register(name, bf)
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{%q: {}}]}`, name)
+
+	opts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultServiceConfig(json),
+	}
+	cc, err := grpc.NewClient(backend.Address, opts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient(_) failed: %v", err)
+	}
+	defer cc.Close()
+
+	// The channel should remain in IDLE as the ExitIdle calls are not
+	// propagated to the lazy balancer from the stub balancer.
+	cc.Connect()
+	shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer shortCancel()
+	testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
+
+	// The picker from the lazy balancer should be send to the channel when the
+	// first resolver update is received by lazy. Making an RPC should trigger
+	// child creation.
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
+	}
+}
+
+// Tests the scenario when a resolver produces a good state followed by a
+// resolver error. The test verifies that the child balancer receives the good
+// update followed by the error.
+func (s) TestGoodUpdateThenResolverError(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	resolverStateReceived := false
+	resolverErrorReceived := grpcsync.NewEvent()
+
+	childBF := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			if resolverErrorReceived.HasFired() {
+				t.Error("Received resolver error before resolver state.")
+			}
+			resolverStateReceived = true
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		ResolverError: func(bd *stub.BalancerData, err error) {
+			if !resolverStateReceived {
+				t.Error("Received resolver error before resolver state.")
+			}
+			resolverErrorReceived.Fire()
+			bd.ChildBalancer.ResolverError(err)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+	}
+
+	childBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_child", "/", "")
+	stub.Register(childBalName, childBF)
+
+	topLevelBF := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
+		},
+		ExitIdle: func(*stub.BalancerData) {
+			t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
+		},
+		ResolverError: func(bd *stub.BalancerData, err error) {
+			bd.ChildBalancer.ResolverError(err)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+	}
+
+	topLevelBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_top_level", "/", "")
+	stub.Register(topLevelBalName, topLevelBF)
+
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{%q: {}}]}`, topLevelBalName)
+
+	mr := manual.NewBuilderWithScheme("e2e-test")
+	defer mr.Close()
+
+	mr.InitialState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend.Address}}},
+		},
+	})
+
+	opts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(mr),
+		grpc.WithDefaultServiceConfig(json),
+	}
+	cc, err := grpc.NewClient(mr.Scheme()+":///whatever", opts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient(_) failed: %v", err)
+
+	}
+
+	defer cc.Close()
+	cc.Connect()
+
+	mr.CC().ReportError(errors.New("test error"))
+	// The channel should remain in IDLE as the ExitIdle calls are not
+	// propagated to the lazy balancer from the stub balancer.
+	shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer shortCancel()
+	testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
+
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
+	}
+
+	if !resolverStateReceived {
+		t.Fatalf("Child balancer did not receive resolver state.")
+	}
+
+	select {
+	case <-resolverErrorReceived.Done():
+	case <-ctx.Done():
+		t.Fatal("Context timed out waiting for resolver error to be delivered to child balancer.")
+	}
+}
+
+// Tests the scenario when a resolver produces a list of endpoints followed by
+// a resolver error. The test verifies that the child balancer receives only the
+// good update.
+func (s) TestResolverErrorThenGoodUpdate(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+
+	childBF := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		ResolverError: func(bd *stub.BalancerData, err error) {
+			t.Error("Received unexpected resolver error.")
+			bd.ChildBalancer.ResolverError(err)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+	}
+
+	childBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_child", "/", "")
+	stub.Register(childBalName, childBF)
+
+	topLevelBF := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
+		},
+		ExitIdle: func(*stub.BalancerData) {
+			t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+	}
+
+	topLevelBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_top_level", "/", "")
+	stub.Register(topLevelBalName, topLevelBF)
+
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{%q: {}}]}`, topLevelBalName)
+
+	mr := manual.NewBuilderWithScheme("e2e-test")
+	defer mr.Close()
+
+	mr.InitialState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend.Address}}},
+		},
+	})
+
+	opts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(mr),
+		grpc.WithDefaultServiceConfig(json),
+	}
+	cc, err := grpc.NewClient(mr.Scheme()+":///whatever", opts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient(_) failed: %v", err)
+
+	}
+
+	defer cc.Close()
+	cc.Connect()
+
+	// Send an error followed by a good update.
+	mr.CC().ReportError(errors.New("test error"))
+	mr.UpdateState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend.Address}}},
+		},
+	})
+
+	// The channel should remain in IDLE as the ExitIdle calls are not
+	// propagated to the lazy balancer from the stub balancer.
+	shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer shortCancel()
+	testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
+
+	// An RPC would succeed only if the leaf pickfirst receives the endpoint
+	// list.
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
+	}
+}
+
+// Tests that ExitIdle calls are correctly passed through to the child balancer.
+// It starts a backend and ensures the channel connects to it. The test then
+// stops the backend, making the channel enter IDLE. The test calls Connect on
+// the channel and verifies that the child balancer exits idle.
+func (s) TestExitIdlePassthrough(t *testing.T) {
+	backend1 := stubserver.StartTestService(t, nil)
+	defer backend1.Stop()
+
+	mr := manual.NewBuilderWithScheme("e2e-test")
+	defer mr.Close()
+
+	mr.InitialState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend1.Address}}},
+		},
+	})
+
+	bf := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build)
+		},
+		ExitIdle: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.ExitIdle()
+		},
+		ResolverError: func(bd *stub.BalancerData, err error) {
+			bd.ChildBalancer.ResolverError(err)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+	}
+	stub.Register(t.Name(), bf)
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s": {}}]}`, t.Name())
+	opts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultServiceConfig(json),
+		grpc.WithResolvers(mr),
+	}
+	cc, err := grpc.NewClient(mr.Scheme()+":///", opts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient(_) failed: %v", err)
+
+	}
+	defer cc.Close()
+
+	cc.Connect()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// Stopping the active backend should put the channel in IDLE.
+	backend1.Stop()
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+
+	// Sending a new backend address should not kick the channel out of IDLE.
+	// On calling cc.Connect(), the channel should call ExitIdle on the lazy
+	// balancer which passes through the call to the leaf pickfirst.
+	backend2 := stubserver.StartTestService(t, nil)
+	defer backend2.Stop()
+
+	mr.UpdateState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: backend2.Address}}},
+		},
+	})
+
+	shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer shortCancel()
+	testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
+
+	cc.Connect()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+}

balancer/leastrequest/leastrequest.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 leastrequest implements a least request load balancer.
+package leastrequest
+
+import (
+	"encoding/json"
+	"fmt"
+	rand "math/rand/v2"
+	"sync"
+	"sync/atomic"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/grpclog"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// Name is the name of the least request balancer.
+const Name = "least_request_experimental"
+
+var (
+	// randuint32 is a global to stub out in tests.
+	randuint32 = rand.Uint32
+	logger     = grpclog.Component("least-request")
+)
+
+func init() {
+	balancer.Register(bb{})
+}
+
+// LBConfig is the balancer config for least_request_experimental balancer.
+type LBConfig struct {
+	serviceconfig.LoadBalancingConfig `json:"-"`
+
+	// ChoiceCount is the number of random SubConns to sample to find the one
+	// with the fewest outstanding requests. If unset, defaults to 2. If set to
+	// < 2, the config will be rejected, and if set to > 10, will become 10.
+	ChoiceCount uint32 `json:"choiceCount,omitempty"`
+}
+
+type bb struct{}
+
+func (bb) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	lbConfig := &LBConfig{
+		ChoiceCount: 2,
+	}
+	if err := json.Unmarshal(s, lbConfig); err != nil {
+		return nil, fmt.Errorf("least-request: unable to unmarshal LBConfig: %v", err)
+	}
+	// "If `choice_count < 2`, the config will be rejected." - A48
+	if lbConfig.ChoiceCount < 2 { // sweet
+		return nil, fmt.Errorf("least-request: lbConfig.choiceCount: %v, must be >= 2", lbConfig.ChoiceCount)
+	}
+	// "If a LeastRequestLoadBalancingConfig with a choice_count > 10 is
+	// received, the least_request_experimental policy will set choice_count =
+	// 10." - A48
+	if lbConfig.ChoiceCount > 10 {
+		lbConfig.ChoiceCount = 10
+	}
+	return lbConfig, nil
+}
+
+func (bb) Name() string {
+	return Name
+}
+
+func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+	b := &leastRequestBalancer{
+		ClientConn:        cc,
+		endpointRPCCounts: resolver.NewEndpointMap[*atomic.Int32](),
+	}
+	b.child = endpointsharding.NewBalancer(b, bOpts, balancer.Get(pickfirst.Name).Build, endpointsharding.Options{})
+	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[%p] ", b))
+	b.logger.Infof("Created")
+	return b
+}
+
+type leastRequestBalancer struct {
+	// Embeds balancer.ClientConn because we need to intercept UpdateState
+	// calls from the child balancer.
+	balancer.ClientConn
+	child  balancer.Balancer
+	logger *internalgrpclog.PrefixLogger
+
+	mu          sync.Mutex
+	choiceCount uint32
+	// endpointRPCCounts holds RPC counts to keep track for subsequent picker
+	// updates.
+	endpointRPCCounts *resolver.EndpointMap[*atomic.Int32]
+}
+
+func (lrb *leastRequestBalancer) Close() {
+	lrb.child.Close()
+	lrb.endpointRPCCounts = nil
+}
+
+func (lrb *leastRequestBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	lrb.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+func (lrb *leastRequestBalancer) ResolverError(err error) {
+	// Will cause inline picker update from endpoint sharding.
+	lrb.child.ResolverError(err)
+}
+
+func (lrb *leastRequestBalancer) ExitIdle() {
+	lrb.child.ExitIdle()
+}
+
+func (lrb *leastRequestBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	lrCfg, ok := ccs.BalancerConfig.(*LBConfig)
+	if !ok {
+		logger.Errorf("least-request: received config with unexpected type %T: %v", ccs.BalancerConfig, ccs.BalancerConfig)
+		return balancer.ErrBadResolverState
+	}
+
+	lrb.mu.Lock()
+	lrb.choiceCount = lrCfg.ChoiceCount
+	lrb.mu.Unlock()
+	return lrb.child.UpdateClientConnState(balancer.ClientConnState{
+		// Enable the health listener in pickfirst children for client side health
+		// checks and outlier detection, if configured.
+		ResolverState: pickfirst.EnableHealthListener(ccs.ResolverState),
+	})
+}
+
+type endpointState struct {
+	picker  balancer.Picker
+	numRPCs *atomic.Int32
+}
+
+func (lrb *leastRequestBalancer) UpdateState(state balancer.State) {
+	var readyEndpoints []endpointsharding.ChildState
+	for _, child := range endpointsharding.ChildStatesFromPicker(state.Picker) {
+		if child.State.ConnectivityState == connectivity.Ready {
+			readyEndpoints = append(readyEndpoints, child)
+		}
+	}
+
+	// If no ready pickers are present, simply defer to the round robin picker
+	// from endpoint sharding, which will round robin across the most relevant
+	// pick first children in the highest precedence connectivity state.
+	if len(readyEndpoints) == 0 {
+		lrb.ClientConn.UpdateState(state)
+		return
+	}
+
+	lrb.mu.Lock()
+	defer lrb.mu.Unlock()
+
+	if logger.V(2) {
+		lrb.logger.Infof("UpdateState called with ready endpoints: %v", readyEndpoints)
+	}
+
+	// Reconcile endpoints.
+	newEndpoints := resolver.NewEndpointMap[any]()
+	for _, child := range readyEndpoints {
+		newEndpoints.Set(child.Endpoint, nil)
+	}
+
+	// If endpoints are no longer ready, no need to count their active RPCs.
+	for endpoint := range lrb.endpointRPCCounts.All() {
+		if _, ok := newEndpoints.Get(endpoint); !ok {
+			lrb.endpointRPCCounts.Delete(endpoint)
+		}
+	}
+
+	// Copy refs to counters into picker.
+	endpointStates := make([]endpointState, 0, len(readyEndpoints))
+	for _, child := range readyEndpoints {
+		counter, ok := lrb.endpointRPCCounts.Get(child.Endpoint)
+		if !ok {
+			// Create new counts if needed.
+			counter = new(atomic.Int32)
+			lrb.endpointRPCCounts.Set(child.Endpoint, counter)
+		}
+		endpointStates = append(endpointStates, endpointState{
+			picker:  child.State.Picker,
+			numRPCs: counter,
+		})
+	}
+
+	lrb.ClientConn.UpdateState(balancer.State{
+		Picker: &picker{
+			choiceCount:    lrb.choiceCount,
+			endpointStates: endpointStates,
+		},
+		ConnectivityState: connectivity.Ready,
+	})
+}
+
+type picker struct {
+	// choiceCount is the number of random endpoints to sample for choosing the
+	// one with the least requests.
+	choiceCount    uint32
+	endpointStates []endpointState
+}
+
+func (p *picker) Pick(pInfo balancer.PickInfo) (balancer.PickResult, error) {
+	var pickedEndpointState *endpointState
+	var pickedEndpointNumRPCs int32
+	for i := 0; i < int(p.choiceCount); i++ {
+		index := randuint32() % uint32(len(p.endpointStates))
+		endpointState := p.endpointStates[index]
+		n := endpointState.numRPCs.Load()
+		if pickedEndpointState == nil || n < pickedEndpointNumRPCs {
+			pickedEndpointState = &endpointState
+			pickedEndpointNumRPCs = n
+		}
+	}
+	result, err := pickedEndpointState.picker.Pick(pInfo)
+	if err != nil {
+		return result, err
+	}
+	// "The counter for a subchannel should be atomically incremented by one
+	// after it has been successfully picked by the picker." - A48
+	pickedEndpointState.numRPCs.Add(1)
+	// "the picker should add a callback for atomically decrementing the
+	// subchannel counter once the RPC finishes (regardless of Status code)." -
+	// A48.
+	originalDone := result.Done
+	result.Done = func(info balancer.DoneInfo) {
+		pickedEndpointState.numRPCs.Add(-1)
+		if originalDone != nil {
+			originalDone(info)
+		}
+	}
+	return result, nil
+}

balancer/leastrequest/leastrequest_test.go

@@ -0,0 +1,770 @@
+/*
+ *
+ * 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 leastrequest
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/stubserver"
+	"google.golang.org/grpc/internal/testutils"
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 10 * time.Millisecond
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestParseConfig(t *testing.T) {
+	parser := bb{}
+	tests := []struct {
+		name    string
+		input   string
+		wantCfg serviceconfig.LoadBalancingConfig
+		wantErr string
+	}{
+		{
+			name:  "happy-case-default",
+			input: `{}`,
+			wantCfg: &LBConfig{
+				ChoiceCount: 2,
+			},
+		},
+		{
+			name:  "happy-case-choice-count-set",
+			input: `{"choiceCount": 3}`,
+			wantCfg: &LBConfig{
+				ChoiceCount: 3,
+			},
+		},
+		{
+			name:  "happy-case-choice-count-greater-than-ten",
+			input: `{"choiceCount": 11}`,
+			wantCfg: &LBConfig{
+				ChoiceCount: 10,
+			},
+		},
+		{
+			name:    "choice-count-less-than-2",
+			input:   `{"choiceCount": 1}`,
+			wantErr: "must be >= 2",
+		},
+		{
+			name:    "invalid-json",
+			input:   "{{invalidjson{{",
+			wantErr: "invalid character",
+		},
+	}
+	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)
+			}
+		})
+	}
+}
+
+func startBackends(t *testing.T, numBackends int) []*stubserver.StubServer {
+	backends := make([]*stubserver.StubServer, 0, numBackends)
+	// Construct and start working backends.
+	for i := 0; i < numBackends; i++ {
+		backend := &stubserver.StubServer{
+			EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+				return &testpb.Empty{}, nil
+			},
+			FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
+				<-stream.Context().Done()
+				return nil
+			},
+		}
+		if err := backend.StartServer(); err != nil {
+			t.Fatalf("Failed to start backend: %v", err)
+		}
+		t.Logf("Started good TestService backend at: %q", backend.Address)
+		t.Cleanup(func() { backend.Stop() })
+		backends = append(backends, backend)
+	}
+	return backends
+}
+
+// setupBackends spins up three test backends, each listening on a port on
+// localhost. The three backends always reply with an empty response with no
+// error, and for streaming receive until hitting an EOF error.
+func setupBackends(t *testing.T, numBackends int) []string {
+	t.Helper()
+	addresses := make([]string, numBackends)
+	backends := startBackends(t, numBackends)
+	// Construct and start working backends.
+	for i := 0; i < numBackends; i++ {
+		addresses[i] = backends[i].Address
+	}
+	return addresses
+}
+
+// checkRoundRobinRPCs verifies that EmptyCall RPCs on the given ClientConn,
+// connected to a server exposing the test.grpc_testing.TestService, are
+// roundrobined across the given backend addresses.
+//
+// Returns a non-nil error if context deadline expires before RPCs start to get
+// roundrobined across the given backends.
+func checkRoundRobinRPCs(ctx context.Context, client testgrpc.TestServiceClient, addrs []resolver.Address) error {
+	wantAddrCount := make(map[string]int)
+	for _, addr := range addrs {
+		wantAddrCount[addr.Addr]++
+	}
+	gotAddrCount := make(map[string]int)
+	for ; ctx.Err() == nil; <-time.After(time.Millisecond) {
+		gotAddrCount = make(map[string]int)
+		// Perform 3 iterations.
+		var iterations [][]string
+		for i := 0; i < 3; i++ {
+			iteration := make([]string, len(addrs))
+			for c := 0; c < len(addrs); c++ {
+				var peer peer.Peer
+				client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer))
+				iteration[c] = peer.Addr.String()
+			}
+			iterations = append(iterations, iteration)
+		}
+		// Ensure the first iteration contains all addresses in addrs.
+		for _, addr := range iterations[0] {
+			gotAddrCount[addr]++
+		}
+		if !cmp.Equal(gotAddrCount, wantAddrCount) {
+			continue
+		}
+		// Ensure all three iterations contain the same addresses.
+		if !cmp.Equal(iterations[0], iterations[1]) || !cmp.Equal(iterations[0], iterations[2]) {
+			continue
+		}
+		return nil
+	}
+	return fmt.Errorf("timeout when waiting for roundrobin distribution of RPCs across addresses: %v; got: %v", addrs, gotAddrCount)
+}
+
+// TestLeastRequestE2E tests the Least Request LB policy in an e2e style. The
+// Least Request balancer is configured as the top level balancer of the
+// channel, and is passed three addresses. Eventually, the test creates three
+// streams, which should be on certain backends according to the least request
+// algorithm. The randomness in the picker is injected in the test to be
+// deterministic, allowing the test to make assertions on the distribution.
+func (s) TestLeastRequestE2E(t *testing.T) {
+	defer func(u func() uint32) {
+		randuint32 = u
+	}(randuint32)
+	var index int
+	indexes := []uint32{
+		0, 0, 1, 1, 2, 2, // Triggers a round robin distribution.
+	}
+	randuint32 = func() uint32 {
+		ret := indexes[index%len(indexes)]
+		index++
+		return ret
+	}
+	addresses := setupBackends(t, 3)
+
+	mr := manual.NewBuilderWithScheme("lr-e2e")
+	defer mr.Close()
+
+	// Configure least request as top level balancer of channel.
+	lrscJSON := `
+{
+  "loadBalancingConfig": [
+    {
+      "least_request_experimental": {
+        "choiceCount": 2
+      }
+    }
+  ]
+}`
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
+	firstThreeAddresses := []resolver.Address{
+		{Addr: addresses[0]},
+		{Addr: addresses[1]},
+		{Addr: addresses[2]},
+	}
+	mr.InitialState(resolver.State{
+		Addresses:     firstThreeAddresses,
+		ServiceConfig: sc,
+	})
+
+	cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testServiceClient := testgrpc.NewTestServiceClient(cc)
+
+	// Wait for all 3 backends to round robin across. The happens because a
+	// SubConn transitioning into READY causes a new picker update. Once the
+	// picker update with all 3 backends is present, this test can start to make
+	// assertions based on those backends.
+	if err := checkRoundRobinRPCs(ctx, testServiceClient, firstThreeAddresses); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// Map ordering of READY SubConns is non deterministic. Thus, perform 3 RPCs
+	// mocked from the random to each index to learn the addresses of SubConns
+	// at each index.
+	index = 0
+	peerAtIndex := make([]string, 3)
+	var peer0 peer.Peer
+	if _, err := testServiceClient.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer0)); err != nil {
+		t.Fatalf("testServiceClient.EmptyCall failed: %v", err)
+	}
+	peerAtIndex[0] = peer0.Addr.String()
+	if _, err := testServiceClient.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer0)); err != nil {
+		t.Fatalf("testServiceClient.EmptyCall failed: %v", err)
+	}
+	peerAtIndex[1] = peer0.Addr.String()
+	if _, err := testServiceClient.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer0)); err != nil {
+		t.Fatalf("testServiceClient.EmptyCall failed: %v", err)
+	}
+	peerAtIndex[2] = peer0.Addr.String()
+
+	// Start streaming RPCs, but do not finish them. Each subsequent stream
+	// should be started according to the least request algorithm, and chosen
+	// between the indexes provided.
+	index = 0
+	indexes = []uint32{
+		0, 0, // Causes first stream to be on first address.
+		0, 1, // Compares first address (one RPC) to second (no RPCs), so choose second.
+		1, 2, // Compares second address (one RPC) to third (no RPCs), so choose third.
+		0, 3, // Causes another stream on first address.
+		1, 0, // Compares second address (one RPC) to first (two RPCs), so choose second.
+		2, 0, // Compares third address (one RPC) to first (two RPCs), so choose third.
+		0, 0, // Causes another stream on first address.
+		2, 2, // Causes a stream on third address.
+		2, 1, // Compares third address (three RPCs) to second (two RPCs), so choose third.
+	}
+	wantIndex := []uint32{0, 1, 2, 0, 1, 2, 0, 2, 1}
+
+	// Start streaming RPC's, but do not finish them. Each created stream should
+	// be started based on the least request algorithm and injected randomness
+	// (see indexes slice above for exact expectations).
+	for _, wantIndex := range wantIndex {
+		stream, err := testServiceClient.FullDuplexCall(ctx)
+		if err != nil {
+			t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
+		}
+		p, ok := peer.FromContext(stream.Context())
+		if !ok {
+			t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
+		}
+		if p.Addr.String() != peerAtIndex[wantIndex] {
+			t.Fatalf("testServiceClient.FullDuplexCall's Peer got: %v, want: %v", p.Addr.String(), peerAtIndex[wantIndex])
+		}
+	}
+}
+
+// TestLeastRequestPersistsCounts tests that the Least Request Balancer persists
+// counts once it gets a new picker update. It first updates the Least Request
+// Balancer with two backends, and creates a bunch of streams on them. Then, it
+// updates the Least Request Balancer with three backends, including the two
+// previous. Any created streams should then be started on the new backend.
+func (s) TestLeastRequestPersistsCounts(t *testing.T) {
+	defer func(u func() uint32) {
+		randuint32 = u
+	}(randuint32)
+	var index int
+	indexes := []uint32{
+		0, 0, 1, 1,
+	}
+	randuint32 = func() uint32 {
+		ret := indexes[index%len(indexes)]
+		index++
+		return ret
+	}
+	addresses := setupBackends(t, 3)
+
+	mr := manual.NewBuilderWithScheme("lr-e2e")
+	defer mr.Close()
+
+	// Configure least request as top level balancer of channel.
+	lrscJSON := `
+{
+  "loadBalancingConfig": [
+    {
+      "least_request_experimental": {
+        "choiceCount": 2
+      }
+    }
+  ]
+}`
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
+	firstTwoAddresses := []resolver.Address{
+		{Addr: addresses[0]},
+		{Addr: addresses[1]},
+	}
+	mr.InitialState(resolver.State{
+		Addresses:     firstTwoAddresses,
+		ServiceConfig: sc,
+	})
+
+	cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testServiceClient := testgrpc.NewTestServiceClient(cc)
+
+	// Wait for the two backends to round robin across. The happens because a
+	// SubConn transitioning into READY causes a new picker update. Once the
+	// picker update with the two backends is present, this test can start to
+	// populate those backends with streams.
+	if err := checkRoundRobinRPCs(ctx, testServiceClient, firstTwoAddresses); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// Start 50 streaming RPCs, and leave them unfinished for the duration of
+	// the test. This will populate the first two addresses with many active
+	// RPCs.
+	for i := 0; i < 50; i++ {
+		_, err := testServiceClient.FullDuplexCall(ctx)
+		if err != nil {
+			t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
+		}
+	}
+
+	// Update the least request balancer to choice count 3. Also update the
+	// address list adding a third address. Alongside the injected randomness,
+	// this should trigger the least request balancer to search all created
+	// SubConns. Thus, since address 3 is the new address and the first two
+	// addresses are populated with RPCs, once the picker update of all 3 READY
+	// SubConns takes effect, all new streams should be started on address 3.
+	index = 0
+	indexes = []uint32{
+		0, 1, 2, 3, 4, 5,
+	}
+	lrscJSON = `
+{
+  "loadBalancingConfig": [
+    {
+      "least_request_experimental": {
+        "choiceCount": 3
+      }
+    }
+  ]
+}`
+	sc = internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
+	fullAddresses := []resolver.Address{
+		{Addr: addresses[0]},
+		{Addr: addresses[1]},
+		{Addr: addresses[2]},
+	}
+	mr.UpdateState(resolver.State{
+		Addresses:     fullAddresses,
+		ServiceConfig: sc,
+	})
+	newAddress := fullAddresses[2]
+	// Poll for only address 3 to show up. This requires a polling loop because
+	// picker update with all three SubConns doesn't take into effect
+	// immediately, needs the third SubConn to become READY.
+	if err := checkRoundRobinRPCs(ctx, testServiceClient, []resolver.Address{newAddress}); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// Start 25 rpcs, but don't finish them. They should all start on address 3,
+	// since the first two addresses both have 25 RPCs (and randomness
+	// injection/choiceCount causes all 3 to be compared every iteration).
+	for i := 0; i < 25; i++ {
+		stream, err := testServiceClient.FullDuplexCall(ctx)
+		if err != nil {
+			t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
+		}
+		p, ok := peer.FromContext(stream.Context())
+		if !ok {
+			t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
+		}
+		if p.Addr.String() != addresses[2] {
+			t.Fatalf("testServiceClient.FullDuplexCall's Peer got: %v, want: %v", p.Addr.String(), addresses[2])
+		}
+	}
+
+	// Now 25 RPC's are active on each address, the next three RPC's should
+	// round robin, since choiceCount is three and the injected random indexes
+	// cause it to search all three addresses for fewest outstanding requests on
+	// each iteration.
+	wantAddrCount := map[string]int{
+		addresses[0]: 1,
+		addresses[1]: 1,
+		addresses[2]: 1,
+	}
+	gotAddrCount := make(map[string]int)
+	for i := 0; i < len(addresses); i++ {
+		stream, err := testServiceClient.FullDuplexCall(ctx)
+		if err != nil {
+			t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
+		}
+		p, ok := peer.FromContext(stream.Context())
+		if !ok {
+			t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
+		}
+		if p.Addr != nil {
+			gotAddrCount[p.Addr.String()]++
+		}
+	}
+	if diff := cmp.Diff(gotAddrCount, wantAddrCount); diff != "" {
+		t.Fatalf("addr count (-got:, +want): %v", diff)
+	}
+}
+
+// TestConcurrentRPCs tests concurrent RPCs on the least request balancer. It
+// configures a channel with a least request balancer as the top level balancer,
+// and makes 100 RPCs asynchronously. This makes sure no race conditions happen
+// in this scenario.
+func (s) TestConcurrentRPCs(t *testing.T) {
+	addresses := setupBackends(t, 3)
+
+	mr := manual.NewBuilderWithScheme("lr-e2e")
+	defer mr.Close()
+
+	// Configure least request as top level balancer of channel.
+	lrscJSON := `
+{
+  "loadBalancingConfig": [
+    {
+      "least_request_experimental": {
+        "choiceCount": 2
+      }
+    }
+  ]
+}`
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
+	firstTwoAddresses := []resolver.Address{
+		{Addr: addresses[0]},
+		{Addr: addresses[1]},
+	}
+	mr.InitialState(resolver.State{
+		Addresses:     firstTwoAddresses,
+		ServiceConfig: sc,
+	})
+
+	cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testServiceClient := testgrpc.NewTestServiceClient(cc)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < 5; j++ {
+				testServiceClient.EmptyCall(ctx, &testpb.Empty{})
+			}
+		}()
+	}
+	wg.Wait()
+}
+
+// Test tests that the least request balancer persists RPC counts once it gets
+// new picker updates and backends within an endpoint go down. It first updates
+// the balancer with two endpoints having two addresses each. It verifies the
+// requests are round robined across the first address of each endpoint. It then
+// stops the active backend in endpoint[0]. It verified that the balancer starts
+// using the second address in endpoint[0]. The test then creates a bunch of
+// streams on two endpoints. Then, it updates the balancer with three endpoints,
+// including the two previous. Any created streams should then be started on the
+// new endpoint. The test shuts down the active backed in endpoint[1] and
+// endpoint[2]. The test verifies that new RPCs are round robined across the
+// active backends in endpoint[1] and endpoint[2].
+func (s) TestLeastRequestEndpoints_MultipleAddresses(t *testing.T) {
+	defer func(u func() uint32) {
+		randuint32 = u
+	}(randuint32)
+	var index int
+	indexes := []uint32{
+		0, 0, 1, 1,
+	}
+	randuint32 = func() uint32 {
+		ret := indexes[index%len(indexes)]
+		index++
+		return ret
+	}
+	backends := startBackends(t, 6)
+	mr := manual.NewBuilderWithScheme("lr-e2e")
+	defer mr.Close()
+
+	// Configure least request as top level balancer of channel.
+	lrscJSON := `
+{
+  "loadBalancingConfig": [
+    {
+      "least_request_experimental": {
+        "choiceCount": 2
+      }
+    }
+  ]
+}`
+	endpoints := []resolver.Endpoint{
+		{Addresses: []resolver.Address{{Addr: backends[0].Address}, {Addr: backends[1].Address}}},
+		{Addresses: []resolver.Address{{Addr: backends[2].Address}, {Addr: backends[3].Address}}},
+		{Addresses: []resolver.Address{{Addr: backends[4].Address}, {Addr: backends[5].Address}}},
+	}
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
+	firstTwoEndpoints := []resolver.Endpoint{endpoints[0], endpoints[1]}
+	mr.InitialState(resolver.State{
+		Endpoints:     firstTwoEndpoints,
+		ServiceConfig: sc,
+	})
+
+	cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testServiceClient := testgrpc.NewTestServiceClient(cc)
+
+	// Wait for the two backends to round robin across. The happens because a
+	// child pickfirst transitioning into READY causes a new picker update. Once
+	// the picker update with the two backends is present, this test can start
+	// to populate those backends with streams.
+	wantAddrs := []resolver.Address{
+		endpoints[0].Addresses[0],
+		endpoints[1].Addresses[0],
+	}
+	if err := checkRoundRobinRPCs(ctx, testServiceClient, wantAddrs); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// Shut down one of the addresses in endpoints[0], the child pickfirst
+	// should fallback to the next address in endpoints[0].
+	backends[0].Stop()
+	wantAddrs = []resolver.Address{
+		endpoints[0].Addresses[1],
+		endpoints[1].Addresses[0],
+	}
+	if err := checkRoundRobinRPCs(ctx, testServiceClient, wantAddrs); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// Start 50 streaming RPCs, and leave them unfinished for the duration of
+	// the test. This will populate the first two endpoints with many active
+	// RPCs.
+	for i := 0; i < 50; i++ {
+		_, err := testServiceClient.FullDuplexCall(ctx)
+		if err != nil {
+			t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
+		}
+	}
+
+	// Update the least request balancer to choice count 3. Also update the
+	// address list adding a third endpoint. Alongside the injected randomness,
+	// this should trigger the least request balancer to search all created
+	// endpoints. Thus, since endpoint 3 is the new endpoint and the first two
+	// endpoint are populated with RPCs, once the picker update of all 3 READY
+	// pickfirsts takes effect, all new streams should be started on endpoint 3.
+	index = 0
+	indexes = []uint32{
+		0, 1, 2, 3, 4, 5,
+	}
+	lrscJSON = `
+{
+  "loadBalancingConfig": [
+    {
+      "least_request_experimental": {
+        "choiceCount": 3
+      }
+    }
+  ]
+}`
+	sc = internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
+	mr.UpdateState(resolver.State{
+		Endpoints:     endpoints,
+		ServiceConfig: sc,
+	})
+	newAddress := endpoints[2].Addresses[0]
+	// Poll for only endpoint 3 to show up. This requires a polling loop because
+	// picker update with all three endpoints doesn't take into effect
+	// immediately, needs the third pickfirst to become READY.
+	if err := checkRoundRobinRPCs(ctx, testServiceClient, []resolver.Address{newAddress}); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+
+	// Start 25 rpcs, but don't finish them. They should all start on endpoint 3,
+	// since the first two endpoints both have 25 RPCs (and randomness
+	// injection/choiceCount causes all 3 to be compared every iteration).
+	for i := 0; i < 25; i++ {
+		stream, err := testServiceClient.FullDuplexCall(ctx)
+		if err != nil {
+			t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
+		}
+		p, ok := peer.FromContext(stream.Context())
+		if !ok {
+			t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
+		}
+		if p.Addr.String() != newAddress.Addr {
+			t.Fatalf("testServiceClient.FullDuplexCall's Peer got: %v, want: %v", p.Addr.String(), newAddress)
+		}
+	}
+
+	// Now 25 RPC's are active on each endpoint, the next three RPC's should
+	// round robin, since choiceCount is three and the injected random indexes
+	// cause it to search all three endpoints for fewest outstanding requests on
+	// each iteration.
+	wantAddrCount := map[string]int{
+		endpoints[0].Addresses[1].Addr: 1,
+		endpoints[1].Addresses[0].Addr: 1,
+		endpoints[2].Addresses[0].Addr: 1,
+	}
+	gotAddrCount := make(map[string]int)
+	for i := 0; i < len(endpoints); i++ {
+		stream, err := testServiceClient.FullDuplexCall(ctx)
+		if err != nil {
+			t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
+		}
+		p, ok := peer.FromContext(stream.Context())
+		if !ok {
+			t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
+		}
+		if p.Addr != nil {
+			gotAddrCount[p.Addr.String()]++
+		}
+	}
+	if diff := cmp.Diff(gotAddrCount, wantAddrCount); diff != "" {
+		t.Fatalf("addr count (-got:, +want): %v", diff)
+	}
+
+	// Shutdown the active address for endpoint[1] and endpoint[2]. This should
+	// result in their streams failing. Now the requests should roundrobin b/w
+	// endpoint[1] and endpoint[2].
+	backends[2].Stop()
+	backends[4].Stop()
+	index = 0
+	indexes = []uint32{
+		0, 1, 2, 2, 1, 0,
+	}
+	wantAddrs = []resolver.Address{
+		endpoints[1].Addresses[1],
+		endpoints[2].Addresses[1],
+	}
+	if err := checkRoundRobinRPCs(ctx, testServiceClient, wantAddrs); err != nil {
+		t.Fatalf("error in expected round robin: %v", err)
+	}
+}
+
+// Test tests that the least request balancer properly surfaces resolver
+// errors.
+func (s) TestLeastRequestEndpoints_ResolverError(t *testing.T) {
+	const sc = `{"loadBalancingConfig": [{"least_request_experimental": {}}]}`
+	mr := manual.NewBuilderWithScheme("lr-e2e")
+	defer mr.Close()
+
+	cc, err := grpc.NewClient(
+		mr.Scheme()+":///",
+		grpc.WithResolvers(mr),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultServiceConfig(sc),
+	)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	// We need to pass an endpoint with a valid address to the resolver before
+	// reporting an error - otherwise endpointsharding does not report the
+	// error through.
+	lis, err := testutils.LocalTCPListener()
+	if err != nil {
+		t.Fatalf("net.Listen() failed: %v", err)
+	}
+	// Act like a server that closes the connection without sending a server
+	// preface.
+	go func() {
+		conn, err := lis.Accept()
+		if err != nil {
+			t.Errorf("Unexpected error when accepting a connection: %v", err)
+		}
+		conn.Close()
+	}()
+	mr.UpdateState(resolver.State{
+		Endpoints: []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: lis.Addr().String()}}}},
+	})
+	cc.Connect()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+
+	// Report an error through the resolver
+	resolverErr := fmt.Errorf("simulated resolver error")
+	mr.CC().ReportError(resolverErr)
+
+	// Ensure the client returns the expected resolver error.
+	testServiceClient := testgrpc.NewTestServiceClient(cc)
+	for ; ctx.Err() == nil; <-time.After(defaultTestShortTimeout) {
+		_, err = testServiceClient.EmptyCall(ctx, &testpb.Empty{})
+		if strings.Contains(err.Error(), resolverErr.Error()) {
+			break
+		}
+	}
+	if ctx.Err() != nil {
+		t.Fatalf("Timeout when waiting for RPCs to fail with error containing %s. Last error: %v", resolverErr, err)
+	}
+}

balancer/pickfirst/internal/internal.go

@@ -0,0 +1,37 @@
+/*
+ * 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 internal contains code internal to the pickfirst package.
+package internal
+
+import (
+	rand "math/rand/v2"
+	"time"
+)
+
+var (
+	// RandShuffle pseudo-randomizes the order of addresses.
+	RandShuffle = rand.Shuffle
+	// RandFloat64 returns, as a float64, a pseudo-random number in [0.0,1.0).
+	RandFloat64 = rand.Float64
+	// TimeAfterFunc allows mocking the timer for testing connection delay
+	// related functionality.
+	TimeAfterFunc = func(d time.Duration, f func()) func() {
+		timer := time.AfterFunc(d, f)
+		return func() { timer.Stop() }
+	}
+)

balancer/pickfirst/pickfirst.go

@@ -0,0 +1,921 @@
+/*
+ *
+ * Copyright 2017 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 pickfirst contains the pick_first load balancing policy which
+// is the universal leaf policy.
+package pickfirst
+
+import (
+	"cmp"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"math"
+	"net"
+	"net/netip"
+	"slices"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst/internal"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/experimental/balancer/weight"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal/envconfig"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/pretty"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+func init() {
+	balancer.Register(pickfirstBuilder{})
+}
+
+// Name is the name of the pick_first balancer.
+const Name = "pick_first"
+
+// enableHealthListenerKeyType is a unique key type used in resolver
+// attributes to indicate whether the health listener usage is enabled.
+type enableHealthListenerKeyType struct{}
+
+var logger = grpclog.Component("pick-first-leaf-lb")
+
+const (
+	// TODO: change to pick-first when this becomes the default pick_first policy.
+	logPrefix = "[pick-first-leaf-lb %p] "
+	// connectionDelayInterval is the time to wait for during the happy eyeballs
+	// pass before starting the next connection attempt.
+	connectionDelayInterval = 250 * time.Millisecond
+)
+
+type ipAddrFamily int
+
+const (
+	// ipAddrFamilyUnknown represents strings that can't be parsed as an IP
+	// address.
+	ipAddrFamilyUnknown ipAddrFamily = iota
+	ipAddrFamilyV4
+	ipAddrFamilyV6
+)
+
+type pickfirstBuilder struct{}
+
+func (pickfirstBuilder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer {
+	b := &pickfirstBalancer{
+		cc: cc,
+
+		subConns:              resolver.NewAddressMapV2[*scData](),
+		state:                 connectivity.Connecting,
+		cancelConnectionTimer: func() {},
+	}
+	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
+	return b
+}
+
+func (b pickfirstBuilder) Name() string {
+	return Name
+}
+
+func (pickfirstBuilder) ParseConfig(js json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	var cfg pfConfig
+	if err := json.Unmarshal(js, &cfg); err != nil {
+		return nil, fmt.Errorf("pickfirst: unable to unmarshal LB policy config: %s, error: %v", string(js), err)
+	}
+	return cfg, nil
+}
+
+// EnableHealthListener updates the state to configure pickfirst for using a
+// generic health listener.
+//
+// # Experimental
+//
+// Notice: This API is EXPERIMENTAL and may be changed or removed in a later
+// release.
+func EnableHealthListener(state resolver.State) resolver.State {
+	state.Attributes = state.Attributes.WithValue(enableHealthListenerKeyType{}, true)
+	return state
+}
+
+type pfConfig struct {
+	serviceconfig.LoadBalancingConfig `json:"-"`
+
+	// If set to true, instructs the LB policy to shuffle the order of the list
+	// of endpoints received from the name resolver before attempting to
+	// connect to them.
+	ShuffleAddressList bool `json:"shuffleAddressList"`
+}
+
+// scData keeps track of the current state of the subConn.
+// It is not safe for concurrent access.
+type scData struct {
+	// The following fields are initialized at build time and read-only after
+	// that.
+	subConn balancer.SubConn
+	addr    resolver.Address
+
+	rawConnectivityState connectivity.State
+	// The effective connectivity state based on raw connectivity, health state
+	// and after following sticky TransientFailure behaviour defined in A62.
+	effectiveState              connectivity.State
+	lastErr                     error
+	connectionFailedInFirstPass bool
+}
+
+func (b *pickfirstBalancer) newSCData(addr resolver.Address) (*scData, error) {
+	sd := &scData{
+		rawConnectivityState: connectivity.Idle,
+		effectiveState:       connectivity.Idle,
+		addr:                 addr,
+	}
+	sc, err := b.cc.NewSubConn([]resolver.Address{addr}, balancer.NewSubConnOptions{
+		StateListener: func(state balancer.SubConnState) {
+			b.updateSubConnState(sd, state)
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	sd.subConn = sc
+	return sd, nil
+}
+
+type pickfirstBalancer struct {
+	// The following fields are initialized at build time and read-only after
+	// that and therefore do not need to be guarded by a mutex.
+	logger *internalgrpclog.PrefixLogger
+	cc     balancer.ClientConn
+
+	// The mutex is used to ensure synchronization of updates triggered
+	// from the idle picker and the already serialized resolver,
+	// SubConn state updates.
+	mu sync.Mutex
+	// State reported to the channel based on SubConn states and resolver
+	// updates.
+	state connectivity.State
+	// scData for active subonns mapped by address.
+	subConns              *resolver.AddressMapV2[*scData]
+	addressList           addressList
+	firstPass             bool
+	numTF                 int
+	cancelConnectionTimer func()
+	healthCheckingEnabled bool
+}
+
+// ResolverError is called by the ClientConn when the name resolver produces
+// an error or when pickfirst determined the resolver update to be invalid.
+func (b *pickfirstBalancer) ResolverError(err error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.resolverErrorLocked(err)
+}
+
+func (b *pickfirstBalancer) resolverErrorLocked(err error) {
+	if b.logger.V(2) {
+		b.logger.Infof("Received error from the name resolver: %v", err)
+	}
+
+	// The picker will not change since the balancer does not currently
+	// report an error. If the balancer hasn't received a single good resolver
+	// update yet, transition to TRANSIENT_FAILURE.
+	if b.state != connectivity.TransientFailure && b.addressList.size() > 0 {
+		if b.logger.V(2) {
+			b.logger.Infof("Ignoring resolver error because balancer is using a previous good update.")
+		}
+		return
+	}
+
+	b.updateBalancerState(balancer.State{
+		ConnectivityState: connectivity.TransientFailure,
+		Picker:            &picker{err: fmt.Errorf("name resolver error: %v", err)},
+	})
+}
+
+func (b *pickfirstBalancer) UpdateClientConnState(state balancer.ClientConnState) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.cancelConnectionTimer()
+	if len(state.ResolverState.Addresses) == 0 && len(state.ResolverState.Endpoints) == 0 {
+		// Cleanup state pertaining to the previous resolver state.
+		// Treat an empty address list like an error by calling b.ResolverError.
+		b.closeSubConnsLocked()
+		b.addressList.updateAddrs(nil)
+		b.resolverErrorLocked(errors.New("produced zero addresses"))
+		return balancer.ErrBadResolverState
+	}
+	b.healthCheckingEnabled = state.ResolverState.Attributes.Value(enableHealthListenerKeyType{}) != nil
+	cfg, ok := state.BalancerConfig.(pfConfig)
+	if state.BalancerConfig != nil && !ok {
+		return fmt.Errorf("pickfirst: received illegal BalancerConfig (type %T): %v: %w", state.BalancerConfig, state.BalancerConfig, balancer.ErrBadResolverState)
+	}
+
+	if b.logger.V(2) {
+		b.logger.Infof("Received new config %s, resolver state %s", pretty.ToJSON(cfg), pretty.ToJSON(state.ResolverState))
+	}
+
+	var newAddrs []resolver.Address
+	if endpoints := state.ResolverState.Endpoints; len(endpoints) != 0 {
+		// Perform the optional shuffling described in gRFC A62. The shuffling
+		// will change the order of endpoints but not touch the order of the
+		// addresses within each endpoint. - A61
+		if cfg.ShuffleAddressList {
+			if envconfig.PickFirstWeightedShuffling {
+				type weightedEndpoint struct {
+					endpoint resolver.Endpoint
+					weight   float64
+				}
+
+				// For each endpoint, compute a key as described in A113 and
+				// https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf:
+				var weightedEndpoints []weightedEndpoint
+				for _, endpoint := range endpoints {
+					u := internal.RandFloat64() // Random number in [0.0, 1.0)
+					weight := weightAttribute(endpoint)
+					weightedEndpoints = append(weightedEndpoints, weightedEndpoint{
+						endpoint: endpoint,
+						weight:   math.Pow(u, 1.0/float64(weight)),
+					})
+				}
+				// Sort endpoints by key in descending order and reconstruct the
+				// endpoints slice.
+				slices.SortFunc(weightedEndpoints, func(a, b weightedEndpoint) int {
+					return cmp.Compare(b.weight, a.weight)
+				})
+
+				// Here, and in the "else" block below, we clone the endpoints
+				// slice to avoid mutating the resolver state. Doing the latter
+				// would lead to data races if the caller is accessing the same
+				// slice concurrently.
+				sortedEndpoints := make([]resolver.Endpoint, len(endpoints))
+				for i, we := range weightedEndpoints {
+					sortedEndpoints[i] = we.endpoint
+				}
+				endpoints = sortedEndpoints
+			} else {
+				endpoints = slices.Clone(endpoints)
+				internal.RandShuffle(len(endpoints), func(i, j int) { endpoints[i], endpoints[j] = endpoints[j], endpoints[i] })
+			}
+		}
+
+		// "Flatten the list by concatenating the ordered list of addresses for
+		// each of the endpoints, in order." - A61
+		for _, endpoint := range endpoints {
+			newAddrs = append(newAddrs, endpoint.Addresses...)
+		}
+	} else {
+		// Endpoints not set, process addresses until we migrate resolver
+		// emissions fully to Endpoints. The top channel does wrap emitted
+		// addresses with endpoints, however some balancers such as weighted
+		// target do not forward the corresponding correct endpoints down/split
+		// endpoints properly. Once all balancers correctly forward endpoints
+		// down, can delete this else conditional.
+		newAddrs = state.ResolverState.Addresses
+		if cfg.ShuffleAddressList {
+			newAddrs = append([]resolver.Address{}, newAddrs...)
+			internal.RandShuffle(len(newAddrs), func(i, j int) { newAddrs[i], newAddrs[j] = newAddrs[j], newAddrs[i] })
+		}
+	}
+
+	// If an address appears in multiple endpoints or in the same endpoint
+	// multiple times, we keep it only once. We will create only one SubConn
+	// for the address because an AddressMap is used to store SubConns.
+	// Not de-duplicating would result in attempting to connect to the same
+	// SubConn multiple times in the same pass. We don't want this.
+	newAddrs = deDupAddresses(newAddrs)
+	newAddrs = interleaveAddresses(newAddrs)
+
+	prevAddr := b.addressList.currentAddress()
+	prevSCData, found := b.subConns.Get(prevAddr)
+	prevAddrsCount := b.addressList.size()
+	isPrevRawConnectivityStateReady := found && prevSCData.rawConnectivityState == connectivity.Ready
+	b.addressList.updateAddrs(newAddrs)
+
+	// If the previous ready SubConn exists in new address list,
+	// keep this connection and don't create new SubConns.
+	if isPrevRawConnectivityStateReady && b.addressList.seekTo(prevAddr) {
+		return nil
+	}
+
+	b.reconcileSubConnsLocked(newAddrs)
+	// If it's the first resolver update or the balancer was already READY
+	// (but the new address list does not contain the ready SubConn) or
+	// CONNECTING, enter CONNECTING.
+	// We may be in TRANSIENT_FAILURE due to a previous empty address list,
+	// we should still enter CONNECTING because the sticky TF behaviour
+	//  mentioned in A62 applies only when the TRANSIENT_FAILURE is reported
+	// due to connectivity failures.
+	if isPrevRawConnectivityStateReady || b.state == connectivity.Connecting || prevAddrsCount == 0 {
+		// Start connection attempt at first address.
+		b.forceUpdateConcludedStateLocked(balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            &picker{err: balancer.ErrNoSubConnAvailable},
+		})
+		b.startFirstPassLocked()
+	} else if b.state == connectivity.TransientFailure {
+		// If we're in TRANSIENT_FAILURE, we stay in TRANSIENT_FAILURE until
+		// we're READY. See A62.
+		b.startFirstPassLocked()
+	}
+	return nil
+}
+
+// UpdateSubConnState is unused as a StateListener is always registered when
+// creating SubConns.
+func (b *pickfirstBalancer) UpdateSubConnState(subConn balancer.SubConn, state balancer.SubConnState) {
+	b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", subConn, state)
+}
+
+func (b *pickfirstBalancer) Close() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.closeSubConnsLocked()
+	b.cancelConnectionTimer()
+	b.state = connectivity.Shutdown
+}
+
+// ExitIdle moves the balancer out of idle state. It can be called concurrently
+// by the idlePicker and clientConn so access to variables should be
+// synchronized.
+func (b *pickfirstBalancer) ExitIdle() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.state == connectivity.Idle {
+		// Move the balancer into CONNECTING state immediately. This is done to
+		// avoid staying in IDLE if a resolver update arrives before the first
+		// SubConn reports CONNECTING.
+		b.updateBalancerState(balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            &picker{err: balancer.ErrNoSubConnAvailable},
+		})
+		b.startFirstPassLocked()
+	}
+}
+
+func (b *pickfirstBalancer) startFirstPassLocked() {
+	b.firstPass = true
+	b.numTF = 0
+	// Reset the connection attempt record for existing SubConns.
+	for _, sd := range b.subConns.All() {
+		sd.connectionFailedInFirstPass = false
+	}
+	b.requestConnectionLocked()
+}
+
+func (b *pickfirstBalancer) closeSubConnsLocked() {
+	for _, sd := range b.subConns.All() {
+		sd.subConn.Shutdown()
+	}
+	b.subConns = resolver.NewAddressMapV2[*scData]()
+}
+
+// deDupAddresses ensures that each address appears only once in the slice.
+func deDupAddresses(addrs []resolver.Address) []resolver.Address {
+	seenAddrs := resolver.NewAddressMapV2[bool]()
+	retAddrs := []resolver.Address{}
+
+	for _, addr := range addrs {
+		if _, ok := seenAddrs.Get(addr); ok {
+			continue
+		}
+		seenAddrs.Set(addr, true)
+		retAddrs = append(retAddrs, addr)
+	}
+	return retAddrs
+}
+
+// interleaveAddresses interleaves addresses of both families (IPv4 and IPv6)
+// as per RFC-8305 section 4.
+// Whichever address family is first in the list is followed by an address of
+// the other address family; that is, if the first address in the list is IPv6,
+// then the first IPv4 address should be moved up in the list to be second in
+// the list. It doesn't support configuring "First Address Family Count", i.e.
+// there will always be a single member of the first address family at the
+// beginning of the interleaved list.
+// Addresses that are neither IPv4 nor IPv6 are treated as part of a third
+// "unknown" family for interleaving.
+// See: https://datatracker.ietf.org/doc/html/rfc8305#autoid-6
+func interleaveAddresses(addrs []resolver.Address) []resolver.Address {
+	familyAddrsMap := map[ipAddrFamily][]resolver.Address{}
+	interleavingOrder := []ipAddrFamily{}
+	for _, addr := range addrs {
+		family := addressFamily(addr.Addr)
+		if _, found := familyAddrsMap[family]; !found {
+			interleavingOrder = append(interleavingOrder, family)
+		}
+		familyAddrsMap[family] = append(familyAddrsMap[family], addr)
+	}
+
+	interleavedAddrs := make([]resolver.Address, 0, len(addrs))
+
+	for curFamilyIdx := 0; len(interleavedAddrs) < len(addrs); curFamilyIdx = (curFamilyIdx + 1) % len(interleavingOrder) {
+		// Some IP types may have fewer addresses than others, so we look for
+		// the next type that has a remaining member to add to the interleaved
+		// list.
+		family := interleavingOrder[curFamilyIdx]
+		remainingMembers := familyAddrsMap[family]
+		if len(remainingMembers) > 0 {
+			interleavedAddrs = append(interleavedAddrs, remainingMembers[0])
+			familyAddrsMap[family] = remainingMembers[1:]
+		}
+	}
+
+	return interleavedAddrs
+}
+
+// addressFamily returns the ipAddrFamily after parsing the address string.
+// If the address isn't of the format "ip-address:port", it returns
+// ipAddrFamilyUnknown. The address may be valid even if it's not an IP when
+// using a resolver like passthrough where the address may be a hostname in
+// some format that the dialer can resolve.
+func addressFamily(address string) ipAddrFamily {
+	// Parse the IP after removing the port.
+	host, _, err := net.SplitHostPort(address)
+	if err != nil {
+		return ipAddrFamilyUnknown
+	}
+	ip, err := netip.ParseAddr(host)
+	if err != nil {
+		return ipAddrFamilyUnknown
+	}
+	switch {
+	case ip.Is4() || ip.Is4In6():
+		return ipAddrFamilyV4
+	case ip.Is6():
+		return ipAddrFamilyV6
+	default:
+		return ipAddrFamilyUnknown
+	}
+}
+
+// reconcileSubConnsLocked updates the active subchannels based on a new address
+// list from the resolver. It does this by:
+//   - closing subchannels: any existing subchannels associated with addresses
+//     that are no longer in the updated list are shut down.
+//   - removing subchannels: entries for these closed subchannels are removed
+//     from the subchannel map.
+//
+// This ensures that the subchannel map accurately reflects the current set of
+// addresses received from the name resolver.
+func (b *pickfirstBalancer) reconcileSubConnsLocked(newAddrs []resolver.Address) {
+	newAddrsMap := resolver.NewAddressMapV2[bool]()
+	for _, addr := range newAddrs {
+		newAddrsMap.Set(addr, true)
+	}
+
+	for oldAddr := range b.subConns.All() {
+		if _, ok := newAddrsMap.Get(oldAddr); ok {
+			continue
+		}
+		val, _ := b.subConns.Get(oldAddr)
+		val.subConn.Shutdown()
+		b.subConns.Delete(oldAddr)
+	}
+}
+
+// shutdownRemainingLocked shuts down remaining subConns. Called when a subConn
+// becomes ready, which means that all other subConn must be shutdown.
+func (b *pickfirstBalancer) shutdownRemainingLocked(selected *scData) {
+	b.cancelConnectionTimer()
+	for _, sd := range b.subConns.All() {
+		if sd.subConn != selected.subConn {
+			sd.subConn.Shutdown()
+		}
+	}
+	b.subConns = resolver.NewAddressMapV2[*scData]()
+	b.subConns.Set(selected.addr, selected)
+}
+
+// requestConnectionLocked starts connecting on the subchannel corresponding to
+// the current address. If no subchannel exists, one is created. If the current
+// subchannel is in TransientFailure, a connection to the next address is
+// attempted until a subchannel is found.
+func (b *pickfirstBalancer) requestConnectionLocked() {
+	if !b.addressList.isValid() {
+		return
+	}
+	var lastErr error
+	for valid := true; valid; valid = b.addressList.increment() {
+		curAddr := b.addressList.currentAddress()
+		sd, ok := b.subConns.Get(curAddr)
+		if !ok {
+			var err error
+			// We want to assign the new scData to sd from the outer scope,
+			// hence we can't use := below.
+			sd, err = b.newSCData(curAddr)
+			if err != nil {
+				// This should never happen, unless the clientConn is being shut
+				// down.
+				if b.logger.V(2) {
+					b.logger.Infof("Failed to create a subConn for address %v: %v", curAddr.String(), err)
+				}
+				// Do nothing, the LB policy will be closed soon.
+				return
+			}
+			b.subConns.Set(curAddr, sd)
+		}
+
+		switch sd.rawConnectivityState {
+		case connectivity.Idle:
+			sd.subConn.Connect()
+			b.scheduleNextConnectionLocked()
+			return
+		case connectivity.TransientFailure:
+			// The SubConn is being re-used and failed during a previous pass
+			// over the addressList. It has not completed backoff yet.
+			// Mark it as having failed and try the next address.
+			sd.connectionFailedInFirstPass = true
+			lastErr = sd.lastErr
+			continue
+		case connectivity.Connecting:
+			// Wait for the connection attempt to complete or the timer to fire
+			// before attempting the next address.
+			b.scheduleNextConnectionLocked()
+			return
+		default:
+			b.logger.Errorf("SubConn with unexpected state %v present in SubConns map.", sd.rawConnectivityState)
+			return
+
+		}
+	}
+
+	// All the remaining addresses in the list are in TRANSIENT_FAILURE, end the
+	// first pass if possible.
+	b.endFirstPassIfPossibleLocked(lastErr)
+}
+
+func (b *pickfirstBalancer) scheduleNextConnectionLocked() {
+	b.cancelConnectionTimer()
+	if !b.addressList.hasNext() {
+		return
+	}
+	curAddr := b.addressList.currentAddress()
+	cancelled := false // Access to this is protected by the balancer's mutex.
+	closeFn := internal.TimeAfterFunc(connectionDelayInterval, func() {
+		b.mu.Lock()
+		defer b.mu.Unlock()
+		// If the scheduled task is cancelled while acquiring the mutex, return.
+		if cancelled {
+			return
+		}
+		if b.logger.V(2) {
+			b.logger.Infof("Happy Eyeballs timer expired while waiting for connection to %q.", curAddr.Addr)
+		}
+		if b.addressList.increment() {
+			b.requestConnectionLocked()
+		}
+	})
+	// Access to the cancellation callback held by the balancer is guarded by
+	// the balancer's mutex, so it's safe to set the boolean from the callback.
+	b.cancelConnectionTimer = sync.OnceFunc(func() {
+		cancelled = true
+		closeFn()
+	})
+}
+
+func (b *pickfirstBalancer) updateSubConnState(sd *scData, newState balancer.SubConnState) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	oldState := sd.rawConnectivityState
+	sd.rawConnectivityState = newState.ConnectivityState
+	// Previously relevant SubConns can still callback with state updates.
+	// To prevent pickers from returning these obsolete SubConns, this logic
+	// is included to check if the current list of active SubConns includes this
+	// SubConn.
+	if !b.isActiveSCData(sd) {
+		return
+	}
+	if newState.ConnectivityState == connectivity.Shutdown {
+		sd.effectiveState = connectivity.Shutdown
+		return
+	}
+
+	if newState.ConnectivityState == connectivity.TransientFailure {
+		sd.connectionFailedInFirstPass = true
+	}
+
+	if newState.ConnectivityState == connectivity.Ready {
+		b.shutdownRemainingLocked(sd)
+		if !b.addressList.seekTo(sd.addr) {
+			// This should not fail as we should have only one SubConn after
+			// entering READY. The SubConn should be present in the addressList.
+			b.logger.Errorf("Address %q not found address list in %v", sd.addr, b.addressList.addresses)
+			return
+		}
+		if !b.healthCheckingEnabled {
+			if b.logger.V(2) {
+				b.logger.Infof("SubConn %p reported connectivity state READY and the health listener is disabled. Transitioning SubConn to READY.", sd.subConn)
+			}
+
+			sd.effectiveState = connectivity.Ready
+			b.updateBalancerState(balancer.State{
+				ConnectivityState: connectivity.Ready,
+				Picker:            &picker{result: balancer.PickResult{SubConn: sd.subConn}},
+			})
+			return
+		}
+		if b.logger.V(2) {
+			b.logger.Infof("SubConn %p reported connectivity state READY. Registering health listener.", sd.subConn)
+		}
+		// Send a CONNECTING update to take the SubConn out of sticky-TF if
+		// required.
+		sd.effectiveState = connectivity.Connecting
+		b.updateBalancerState(balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            &picker{err: balancer.ErrNoSubConnAvailable},
+		})
+		sd.subConn.RegisterHealthListener(func(scs balancer.SubConnState) {
+			b.updateSubConnHealthState(sd, scs)
+		})
+		return
+	}
+
+	// If the LB policy is READY, and it receives a subchannel state change,
+	// it means that the READY subchannel has failed.
+	// A SubConn can also transition from CONNECTING directly to IDLE when
+	// a transport is successfully created, but the connection fails
+	// before the SubConn can send the notification for READY. We treat
+	// this as a successful connection and transition to IDLE.
+	// TODO: https://github.com/grpc/grpc-go/issues/7862 - Remove the second
+	// part of the if condition below once the issue is fixed.
+	if oldState == connectivity.Ready || (oldState == connectivity.Connecting && newState.ConnectivityState == connectivity.Idle) {
+		// Once a transport fails, the balancer enters IDLE and starts from
+		// the first address when the picker is used.
+		b.shutdownRemainingLocked(sd)
+		sd.effectiveState = newState.ConnectivityState
+		b.addressList.reset()
+		b.updateBalancerState(balancer.State{
+			ConnectivityState: connectivity.Idle,
+			Picker:            &idlePicker{exitIdle: sync.OnceFunc(b.ExitIdle)},
+		})
+		return
+	}
+
+	if b.firstPass {
+		switch newState.ConnectivityState {
+		case connectivity.Connecting:
+			// The effective state can be in either IDLE, CONNECTING or
+			// TRANSIENT_FAILURE. If it's  TRANSIENT_FAILURE, stay in
+			// TRANSIENT_FAILURE until it's READY. See A62.
+			if sd.effectiveState != connectivity.TransientFailure {
+				sd.effectiveState = connectivity.Connecting
+				b.updateBalancerState(balancer.State{
+					ConnectivityState: connectivity.Connecting,
+					Picker:            &picker{err: balancer.ErrNoSubConnAvailable},
+				})
+			}
+		case connectivity.TransientFailure:
+			sd.lastErr = newState.ConnectionError
+			sd.effectiveState = connectivity.TransientFailure
+			// Since we're re-using common SubConns while handling resolver
+			// updates, we could receive an out of turn TRANSIENT_FAILURE from
+			// a pass over the previous address list. Happy Eyeballs will also
+			// cause out of order updates to arrive.
+
+			if curAddr := b.addressList.currentAddress(); equalAddressIgnoringBalAttributes(&curAddr, &sd.addr) {
+				b.cancelConnectionTimer()
+				if b.addressList.increment() {
+					b.requestConnectionLocked()
+					return
+				}
+			}
+
+			// End the first pass if we've seen a TRANSIENT_FAILURE from all
+			// SubConns once.
+			b.endFirstPassIfPossibleLocked(newState.ConnectionError)
+		}
+		return
+	}
+
+	// We have finished the first pass, keep re-connecting failing SubConns.
+	switch newState.ConnectivityState {
+	case connectivity.TransientFailure:
+		b.numTF = (b.numTF + 1) % b.subConns.Len()
+		sd.lastErr = newState.ConnectionError
+		if b.numTF%b.subConns.Len() == 0 {
+			b.updateBalancerState(balancer.State{
+				ConnectivityState: connectivity.TransientFailure,
+				Picker:            &picker{err: newState.ConnectionError},
+			})
+		}
+		// We don't need to request re-resolution since the SubConn already
+		// does that before reporting TRANSIENT_FAILURE.
+		// TODO: #7534 - Move re-resolution requests from SubConn into
+		// pick_first.
+	case connectivity.Idle:
+		sd.subConn.Connect()
+	}
+}
+
+// endFirstPassIfPossibleLocked ends the first happy-eyeballs pass if all the
+// addresses are tried and their SubConns have reported a failure.
+func (b *pickfirstBalancer) endFirstPassIfPossibleLocked(lastErr error) {
+	// An optimization to avoid iterating over the entire SubConn map.
+	if b.addressList.isValid() {
+		return
+	}
+	// Connect() has been called on all the SubConns. The first pass can be
+	// ended if all the SubConns have reported a failure.
+	for _, sd := range b.subConns.All() {
+		if !sd.connectionFailedInFirstPass {
+			return
+		}
+	}
+	b.firstPass = false
+	b.updateBalancerState(balancer.State{
+		ConnectivityState: connectivity.TransientFailure,
+		Picker:            &picker{err: lastErr},
+	})
+	// Start re-connecting all the SubConns that are already in IDLE.
+	for _, sd := range b.subConns.All() {
+		if sd.rawConnectivityState == connectivity.Idle {
+			sd.subConn.Connect()
+		}
+	}
+}
+
+func (b *pickfirstBalancer) isActiveSCData(sd *scData) bool {
+	activeSD, found := b.subConns.Get(sd.addr)
+	return found && activeSD == sd
+}
+
+func (b *pickfirstBalancer) updateSubConnHealthState(sd *scData, state balancer.SubConnState) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	// Previously relevant SubConns can still callback with state updates.
+	// To prevent pickers from returning these obsolete SubConns, this logic
+	// is included to check if the current list of active SubConns includes
+	// this SubConn.
+	if !b.isActiveSCData(sd) {
+		return
+	}
+	sd.effectiveState = state.ConnectivityState
+	switch state.ConnectivityState {
+	case connectivity.Ready:
+		b.updateBalancerState(balancer.State{
+			ConnectivityState: connectivity.Ready,
+			Picker:            &picker{result: balancer.PickResult{SubConn: sd.subConn}},
+		})
+	case connectivity.TransientFailure:
+		b.updateBalancerState(balancer.State{
+			ConnectivityState: connectivity.TransientFailure,
+			Picker:            &picker{err: fmt.Errorf("pickfirst: health check failure: %v", state.ConnectionError)},
+		})
+	case connectivity.Connecting:
+		b.updateBalancerState(balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            &picker{err: balancer.ErrNoSubConnAvailable},
+		})
+	default:
+		b.logger.Errorf("Got unexpected health update for SubConn %p: %v", state)
+	}
+}
+
+// updateBalancerState stores the state reported to the channel and calls
+// ClientConn.UpdateState(). As an optimization, it avoids sending duplicate
+// updates to the channel.
+func (b *pickfirstBalancer) updateBalancerState(newState balancer.State) {
+	// In case of TransientFailures allow the picker to be updated to update
+	// the connectivity error, in all other cases don't send duplicate state
+	// updates.
+	if newState.ConnectivityState == b.state && b.state != connectivity.TransientFailure {
+		return
+	}
+	b.forceUpdateConcludedStateLocked(newState)
+}
+
+// forceUpdateConcludedStateLocked stores the state reported to the channel and
+// calls ClientConn.UpdateState().
+// A separate function is defined to force update the ClientConn state since the
+// channel doesn't correctly assume that LB policies start in CONNECTING and
+// relies on LB policy to send an initial CONNECTING update.
+func (b *pickfirstBalancer) forceUpdateConcludedStateLocked(newState balancer.State) {
+	b.state = newState.ConnectivityState
+	b.cc.UpdateState(newState)
+}
+
+type picker struct {
+	result balancer.PickResult
+	err    error
+}
+
+func (p *picker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	return p.result, p.err
+}
+
+// idlePicker is used when the SubConn is IDLE and kicks the SubConn into
+// CONNECTING when Pick is called.
+type idlePicker struct {
+	exitIdle func()
+}
+
+func (i *idlePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	i.exitIdle()
+	return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+}
+
+// addressList manages sequentially iterating over addresses present in a list
+// of endpoints. It provides a 1 dimensional view of the addresses present in
+// the endpoints.
+// This type is not safe for concurrent access.
+type addressList struct {
+	addresses []resolver.Address
+	idx       int
+}
+
+func (al *addressList) isValid() bool {
+	return al.idx < len(al.addresses)
+}
+
+func (al *addressList) size() int {
+	return len(al.addresses)
+}
+
+// increment moves to the next index in the address list.
+// This method returns false if it went off the list, true otherwise.
+func (al *addressList) increment() bool {
+	if !al.isValid() {
+		return false
+	}
+	al.idx++
+	return al.idx < len(al.addresses)
+}
+
+// currentAddress returns the current address pointed to in the addressList.
+// If the list is in an invalid state, it returns an empty address instead.
+func (al *addressList) currentAddress() resolver.Address {
+	if !al.isValid() {
+		return resolver.Address{}
+	}
+	return al.addresses[al.idx]
+}
+
+func (al *addressList) reset() {
+	al.idx = 0
+}
+
+func (al *addressList) updateAddrs(addrs []resolver.Address) {
+	al.addresses = addrs
+	al.reset()
+}
+
+// seekTo returns false if the needle was not found and the current index was
+// left unchanged.
+func (al *addressList) seekTo(needle resolver.Address) bool {
+	for ai, addr := range al.addresses {
+		if !equalAddressIgnoringBalAttributes(&addr, &needle) {
+			continue
+		}
+		al.idx = ai
+		return true
+	}
+	return false
+}
+
+// hasNext returns whether incrementing the addressList will result in moving
+// past the end of the list. If the list has already moved past the end, it
+// returns false.
+func (al *addressList) hasNext() bool {
+	if !al.isValid() {
+		return false
+	}
+	return al.idx+1 < len(al.addresses)
+}
+
+// equalAddressIgnoringBalAttributes returns true is a and b are considered
+// equal. This is different from the Equal method on the resolver.Address type
+// which considers all fields to determine equality. Here, we only consider
+// fields that are meaningful to the SubConn.
+func equalAddressIgnoringBalAttributes(a, b *resolver.Address) bool {
+	return a.Addr == b.Addr && a.ServerName == b.ServerName &&
+		a.Attributes.Equal(b.Attributes)
+}
+
+// weightAttribute is a convenience function which returns the value of the
+// weight endpoint Attribute.
+//
+// When used in the xDS context, the weight attribute is guaranteed to be
+// non-zero. But, when used in a non-xDS context, the weight attribute could be
+// unset. A Default of 1 is used in the latter case.
+func weightAttribute(e resolver.Endpoint) uint32 {
+	w := weight.FromEndpoint(e).Weight
+	if w == 0 {
+		return 1
+	}
+	return w
+}

balancer/pickfirst/pickfirst_test.go

@@ -0,0 +1,376 @@
+/*
+ *
+ * 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 pickfirst
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/resolver"
+)
+
+const (
+	// Default timeout for tests in this package.
+	defaultTestTimeout = 10 * time.Second
+	// Default short timeout, to be used when waiting for events which are not
+	// expected to happen.
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// TestPickFirst_InitialResolverError sends a resolver error to the balancer
+// before a valid resolver update. It verifies that the clientconn state is
+// updated to TRANSIENT_FAILURE.
+func (s) TestPickFirst_InitialResolverError(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	cc := testutils.NewBalancerClientConn(t)
+	bal := balancer.Get(Name).Build(cc, balancer.BuildOptions{})
+	defer bal.Close()
+	bal.ResolverError(errors.New("resolution failed: test error"))
+
+	if err := cc.WaitForConnectivityState(ctx, connectivity.TransientFailure); err != nil {
+		t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.TransientFailure, err)
+	}
+
+	// After sending a valid update, the LB policy should report CONNECTING.
+	ccState := balancer.ClientConnState{
+		ResolverState: resolver.State{
+			Endpoints: []resolver.Endpoint{
+				{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}},
+				{Addresses: []resolver.Address{{Addr: "2.2.2.2:2"}}},
+			},
+		},
+	}
+	if err := bal.UpdateClientConnState(ccState); err != nil {
+		t.Fatalf("UpdateClientConnState(%v) returned error: %v", ccState, err)
+	}
+
+	if err := cc.WaitForConnectivityState(ctx, connectivity.Connecting); err != nil {
+		t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.Connecting, err)
+	}
+}
+
+// TestPickFirst_ResolverErrorinTF sends a resolver error to the balancer
+// before when it's attempting to connect to a SubConn TRANSIENT_FAILURE. It
+// verifies that the picker is updated and the SubConn is not closed.
+func (s) TestPickFirst_ResolverErrorinTF(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	cc := testutils.NewBalancerClientConn(t)
+	bal := balancer.Get(Name).Build(cc, balancer.BuildOptions{})
+	defer bal.Close()
+
+	// After sending a valid update, the LB policy should report CONNECTING.
+	ccState := balancer.ClientConnState{
+		ResolverState: resolver.State{
+			Endpoints: []resolver.Endpoint{
+				{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}},
+			},
+		},
+	}
+
+	if err := bal.UpdateClientConnState(ccState); err != nil {
+		t.Fatalf("UpdateClientConnState(%v) returned error: %v", ccState, err)
+	}
+
+	sc1 := <-cc.NewSubConnCh
+	if err := cc.WaitForConnectivityState(ctx, connectivity.Connecting); err != nil {
+		t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.Connecting, err)
+	}
+
+	scErr := fmt.Errorf("test error: connection refused")
+	sc1.UpdateState(balancer.SubConnState{
+		ConnectivityState: connectivity.TransientFailure,
+		ConnectionError:   scErr,
+	})
+
+	if err := cc.WaitForPickerWithErr(ctx, scErr); err != nil {
+		t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", scErr, err)
+	}
+
+	bal.ResolverError(errors.New("resolution failed: test error"))
+	if err := cc.WaitForErrPicker(ctx); err != nil {
+		t.Fatalf("cc.WaitForPickerWithErr() returned error: %v", err)
+	}
+
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case sc := <-cc.ShutdownSubConnCh:
+		t.Fatalf("Unexpected SubConn shutdown: %v", sc)
+	}
+}
+
+// TestAddressList_Iteration verifies the behaviour of the addressList while
+// iterating through the entries.
+func (s) TestAddressList_Iteration(t *testing.T) {
+	addrs := []resolver.Address{
+		{
+			Addr:               "192.168.1.1",
+			ServerName:         "test-host-1",
+			Attributes:         attributes.New("key-1", "val-1"),
+			BalancerAttributes: attributes.New("bal-key-1", "bal-val-1"),
+		},
+		{
+			Addr:               "192.168.1.2",
+			ServerName:         "test-host-2",
+			Attributes:         attributes.New("key-2", "val-2"),
+			BalancerAttributes: attributes.New("bal-key-2", "bal-val-2"),
+		},
+		{
+			Addr:               "192.168.1.3",
+			ServerName:         "test-host-3",
+			Attributes:         attributes.New("key-3", "val-3"),
+			BalancerAttributes: attributes.New("bal-key-3", "bal-val-3"),
+		},
+	}
+
+	addressList := addressList{}
+	addressList.updateAddrs(addrs)
+
+	for i := 0; i < len(addrs); i++ {
+		if got, want := addressList.isValid(), true; got != want {
+			t.Fatalf("addressList.isValid() = %t, want %t", got, want)
+		}
+		if got, want := addressList.currentAddress(), addrs[i]; !want.Equal(got) {
+			t.Errorf("addressList.currentAddress() = %v, want %v", got, want)
+		}
+		if got, want := addressList.increment(), i+1 < len(addrs); got != want {
+			t.Fatalf("addressList.increment() = %t, want %t", got, want)
+		}
+	}
+
+	if got, want := addressList.isValid(), false; got != want {
+		t.Fatalf("addressList.isValid() = %t, want %t", got, want)
+	}
+
+	// increment an invalid address list.
+	if got, want := addressList.increment(), false; got != want {
+		t.Errorf("addressList.increment() = %t, want %t", got, want)
+	}
+
+	if got, want := addressList.isValid(), false; got != want {
+		t.Errorf("addressList.isValid() = %t, want %t", got, want)
+	}
+
+	addressList.reset()
+	for i := 0; i < len(addrs); i++ {
+		if got, want := addressList.isValid(), true; got != want {
+			t.Fatalf("addressList.isValid() = %t, want %t", got, want)
+		}
+		if got, want := addressList.currentAddress(), addrs[i]; !want.Equal(got) {
+			t.Errorf("addressList.currentAddress() = %v, want %v", got, want)
+		}
+		if got, want := addressList.increment(), i+1 < len(addrs); got != want {
+			t.Fatalf("addressList.increment() = %t, want %t", got, want)
+		}
+	}
+}
+
+// TestAddressList_SeekTo verifies the behaviour of addressList.seekTo.
+func (s) TestAddressList_SeekTo(t *testing.T) {
+	addrs := []resolver.Address{
+		{
+			Addr:               "192.168.1.1",
+			ServerName:         "test-host-1",
+			Attributes:         attributes.New("key-1", "val-1"),
+			BalancerAttributes: attributes.New("bal-key-1", "bal-val-1"),
+		},
+		{
+			Addr:               "192.168.1.2",
+			ServerName:         "test-host-2",
+			Attributes:         attributes.New("key-2", "val-2"),
+			BalancerAttributes: attributes.New("bal-key-2", "bal-val-2"),
+		},
+		{
+			Addr:               "192.168.1.3",
+			ServerName:         "test-host-3",
+			Attributes:         attributes.New("key-3", "val-3"),
+			BalancerAttributes: attributes.New("bal-key-3", "bal-val-3"),
+		},
+	}
+
+	addressList := addressList{}
+	addressList.updateAddrs(addrs)
+
+	// Try finding an address in the list.
+	key := resolver.Address{
+		Addr:               "192.168.1.2",
+		ServerName:         "test-host-2",
+		Attributes:         attributes.New("key-2", "val-2"),
+		BalancerAttributes: attributes.New("ignored", "bal-val-2"),
+	}
+
+	if got, want := addressList.seekTo(key), true; got != want {
+		t.Errorf("addressList.seekTo(%v) = %t, want %t", key, got, want)
+	}
+
+	// It should be possible to increment once more now that the pointer has advanced.
+	if got, want := addressList.increment(), true; got != want {
+		t.Errorf("addressList.increment() = %t, want %t", got, want)
+	}
+
+	if got, want := addressList.increment(), false; got != want {
+		t.Errorf("addressList.increment() = %t, want %t", got, want)
+	}
+
+	// Seek to the key again, it is behind the pointer now.
+	if got, want := addressList.seekTo(key), true; got != want {
+		t.Errorf("addressList.seekTo(%v) = %t, want %t", key, got, want)
+	}
+
+	// Seek to a key not in the list.
+	key = resolver.Address{
+		Addr:               "192.168.1.5",
+		ServerName:         "test-host-5",
+		Attributes:         attributes.New("key-5", "val-5"),
+		BalancerAttributes: attributes.New("ignored", "bal-val-5"),
+	}
+
+	if got, want := addressList.seekTo(key), false; got != want {
+		t.Errorf("addressList.seekTo(%v) = %t, want %t", key, got, want)
+	}
+
+	// It should be possible to increment once more since the pointer has not advanced.
+	if got, want := addressList.increment(), true; got != want {
+		t.Errorf("addressList.increment() = %t, want %t", got, want)
+	}
+
+	if got, want := addressList.increment(), false; got != want {
+		t.Errorf("addressList.increment() = %t, want %t", got, want)
+	}
+}
+
+// TestPickFirstLeaf_TFPickerUpdate sends TRANSIENT_FAILURE SubConn state updates
+// for each SubConn managed by a pickfirst balancer. It verifies that the picker
+// is updated with the expected frequency.
+func (s) TestPickFirstLeaf_TFPickerUpdate(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	cc := testutils.NewBalancerClientConn(t)
+	bal := pickfirstBuilder{}.Build(cc, balancer.BuildOptions{})
+	defer bal.Close()
+	ccState := balancer.ClientConnState{
+		ResolverState: resolver.State{
+			Endpoints: []resolver.Endpoint{
+				{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}},
+				{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}}, // duplicate, should be ignored.
+				{Addresses: []resolver.Address{{Addr: "2.2.2.2:2"}}},
+				{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}}, // duplicate, should be ignored.
+			},
+		},
+	}
+	if err := bal.UpdateClientConnState(ccState); err != nil {
+		t.Fatalf("UpdateClientConnState(%v) returned error: %v", ccState, err)
+	}
+
+	// PF should report TRANSIENT_FAILURE only once all the sunbconns have failed
+	// once.
+	tfErr := fmt.Errorf("test err: connection refused")
+	sc1 := <-cc.NewSubConnCh
+	select {
+	case <-sc1.ConnectCh:
+	case <-ctx.Done():
+		t.Fatal("Context timed out waiting for Connect() to be called on sc1.")
+	}
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: tfErr})
+
+	// Move the subconn back to IDLE, it should not be re-connected until the
+	// first pass is complete.
+	shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer shortCancel()
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
+	select {
+	case <-sc1.ConnectCh:
+		t.Fatal("Connect() unexpectedly called on sc1.")
+	case <-shortCtx.Done():
+	}
+
+	if err := cc.WaitForPickerWithErr(ctx, balancer.ErrNoSubConnAvailable); err != nil {
+		t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", balancer.ErrNoSubConnAvailable, err)
+	}
+
+	sc2 := <-cc.NewSubConnCh
+	select {
+	case <-sc2.ConnectCh:
+	case <-ctx.Done():
+		t.Fatal("Context timed out waiting for Connect() to be called on sc2.")
+	}
+	sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: tfErr})
+
+	if err := cc.WaitForPickerWithErr(ctx, tfErr); err != nil {
+		t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", tfErr, err)
+	}
+
+	// Subsequent TRANSIENT_FAILUREs should be reported only after seeing "# of SubConns"
+	// TRANSIENT_FAILUREs.
+
+	// Both the subconns should be connected in parallel.
+	select {
+	case <-sc1.ConnectCh:
+	case <-ctx.Done():
+		t.Fatal("Context timed out waiting for Connect() to be called on sc1.")
+	}
+
+	shortCtx, shortCancel = context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer shortCancel()
+	select {
+	case <-sc2.ConnectCh:
+		t.Fatal("Connect() called on sc2 before it completed backing-off.")
+	case <-shortCtx.Done():
+	}
+
+	sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
+	select {
+	case <-sc2.ConnectCh:
+	case <-ctx.Done():
+		t.Fatal("Context timed out waiting for Connect() to be called on sc2.")
+	}
+
+	newTfErr := fmt.Errorf("test err: unreachable")
+	sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: newTfErr})
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case p := <-cc.NewPickerCh:
+		sc, err := p.Pick(balancer.PickInfo{})
+		t.Fatalf("Unexpected picker update: %v, %v", sc, err)
+	}
+
+	sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: newTfErr})
+	if err := cc.WaitForPickerWithErr(ctx, newTfErr); err != nil {
+		t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", newTfErr, err)
+	}
+}

balancer/randomsubsetting/randomsubsetting.go

@@ -0,0 +1,201 @@
+/*
+ *
+ * Copyright 2025 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 randomsubsetting implements the random_subsetting LB policy specified
+// here: https://github.com/grpc/proposal/blob/master/A68-random-subsetting.md
+//
+// To install the LB policy, import this package as:
+//
+//	import _ "google.golang.org/grpc/balancer/randomsubsetting"
+//
+// # Experimental
+//
+// Notice: This package is EXPERIMENTAL and may be changed or removed in a
+// later release.
+package randomsubsetting
+
+import (
+	"cmp"
+	"encoding/json"
+	"fmt"
+	"math/rand/v2"
+	"slices"
+
+	xxhash "github.com/cespare/xxhash/v2"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal/balancer/gracefulswitch"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// Name is the name of the random subsetting load balancer.
+const Name = "random_subsetting_experimental"
+
+var (
+	logger     = grpclog.Component(Name)
+	randUint64 = rand.Uint64
+)
+
+func prefixLogger(p *subsettingBalancer) *internalgrpclog.PrefixLogger {
+	return internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[random-subsetting-lb %p] ", p))
+}
+
+func init() {
+	balancer.Register(bb{})
+}
+
+type bb struct{}
+
+func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+	b := &subsettingBalancer{
+		Balancer:   gracefulswitch.NewBalancer(cc, bOpts),
+		hashSeed:   randUint64(),
+		hashDigest: xxhash.New(),
+	}
+	b.logger = prefixLogger(b)
+	b.logger.Infof("Created")
+	return b
+}
+
+type lbConfig struct {
+	serviceconfig.LoadBalancingConfig `json:"-"`
+
+	SubsetSize  uint32                         `json:"subsetSize,omitempty"`
+	ChildPolicy *iserviceconfig.BalancerConfig `json:"childPolicy,omitempty"`
+}
+
+func (bb) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	lbCfg := &lbConfig{}
+
+	// Ensure that the specified child policy is registered and validates its
+	// config, if present.
+	if err := json.Unmarshal(s, lbCfg); err != nil {
+		return nil, fmt.Errorf("randomsubsetting: json.Unmarshal failed for configuration: %s with error: %v", string(s), err)
+	}
+	if lbCfg.SubsetSize == 0 {
+		return nil, fmt.Errorf("randomsubsetting: SubsetSize must be greater than 0")
+	}
+	if lbCfg.ChildPolicy == nil {
+		return nil, fmt.Errorf("randomsubsetting: ChildPolicy must be specified")
+	}
+
+	return lbCfg, nil
+}
+
+func (bb) Name() string {
+	return Name
+}
+
+type subsettingBalancer struct {
+	*gracefulswitch.Balancer
+
+	logger     *internalgrpclog.PrefixLogger
+	cfg        *lbConfig
+	hashSeed   uint64
+	hashDigest *xxhash.Digest
+}
+
+func (b *subsettingBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
+	lbCfg, ok := s.BalancerConfig.(*lbConfig)
+	if !ok {
+		b.logger.Warningf("Received config with unexpected type %T: %v", s.BalancerConfig, s.BalancerConfig)
+		return balancer.ErrBadResolverState
+	}
+
+	// Build config for the gracefulswitch balancer. It is safe to ignore
+	// JSON marshaling errors here, since the config was already validated
+	// as part of ParseConfig().
+	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
+	}
+
+	// Empty endpoints cannot be hashed and are ignored by child policies.
+	// Remove them here so they do not count toward the requested subset size.
+	isEmpty := func(endpoint resolver.Endpoint) bool {
+		return len(endpoint.Addresses) == 0
+	}
+	if slices.ContainsFunc(endpoints, isEmpty) {
+		endpoints = slices.DeleteFunc(slices.Clone(endpoints), isEmpty)
+	}
+
+	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,
+		}
+	}
+
+	slices.SortFunc(hashedEndpoints, func(a, b endpointWithHash) int {
+		// Note: This uses the standard library cmp package, not the
+		// github.com/google/go-cmp/cmp package. The latter is intended for
+		// testing purposes only.
+		return cmp.Compare(a.hash, b.hash)
+	})
+
+	// Convert back to resolver.Endpoints
+	endpointSubset := make([]resolver.Endpoint, subsetSize)
+	for i, endpoint := range hashedEndpoints[:subsetSize] {
+		endpointSubset[i] = endpoint.ep
+	}
+
+	return endpointSubset
+}

balancer/randomsubsetting/randomsubsetting_test.go

@@ -0,0 +1,281 @@
+/*
+ *
+ * Copyright 2025 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 randomsubsetting
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	xxhash "github.com/cespare/xxhash/v2"
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/grpctest"
+	iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+
+	_ "google.golang.org/grpc/balancer/roundrobin" // For round_robin LB policy in tests
+)
+
+var defaultTestTimeout = 5 * time.Second
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestParseConfig(t *testing.T) {
+	parser := bb{}
+	tests := []struct {
+		name    string
+		input   string
+		wantCfg serviceconfig.LoadBalancingConfig
+		wantErr string
+	}{
+		{
+			name:    "invalid_json",
+			input:   "{{invalidjson{{",
+			wantErr: "json.Unmarshal failed for configuration",
+		},
+		{
+			name:    "empty_config",
+			input:   `{}`,
+			wantErr: "SubsetSize must be greater than 0",
+		},
+		{
+			name:    "subset_size_zero",
+			input:   `{ "subsetSize": 0 }`,
+			wantErr: "SubsetSize must be greater than 0",
+		},
+		{
+			name:    "child_policy_missing",
+			input:   `{ "subsetSize": 1 }`,
+			wantErr: "ChildPolicy must be specified",
+		},
+		{
+			name:    "child_policy_not_registered",
+			input:   `{ "subsetSize": 1 , "childPolicy": [{"unregistered_lb": {}}] }`,
+			wantErr: "no supported policies found",
+		},
+		{
+			name:  "success",
+			input: `{ "subsetSize": 3, "childPolicy": [{"round_robin": {}}]}`,
+			wantCfg: &lbConfig{
+				SubsetSize:  3,
+				ChildPolicy: &iserviceconfig.BalancerConfig{Name: "round_robin"},
+			},
+		},
+	}
+	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(test.wantCfg, gotCfg); diff != "" {
+				t.Fatalf("ParseConfig(%s) got unexpected output, diff (-want +got):\n%s", test.input, diff)
+			}
+		})
+	}
+}
+
+func makeEndpoints(n int) []resolver.Endpoint {
+	endpoints := make([]resolver.Endpoint, n)
+	for i := 0; i < n; i++ {
+		endpoints[i] = resolver.Endpoint{
+			Addresses: []resolver.Address{{Addr: fmt.Sprintf("endpoint-%d", i)}},
+		}
+	}
+	return endpoints
+}
+
+func (s) TestCalculateSubset_Simple(t *testing.T) {
+	tests := []struct {
+		name       string
+		endpoints  []resolver.Endpoint
+		subsetSize uint32
+		want       []resolver.Endpoint
+	}{
+		{
+			name:       "NoEndpoints",
+			endpoints:  []resolver.Endpoint{},
+			subsetSize: 3,
+			want:       []resolver.Endpoint{},
+		},
+		{
+			name:       "SubsetSizeLargerThanNumberOfEndpoints",
+			endpoints:  makeEndpoints(5),
+			subsetSize: 10,
+			want:       makeEndpoints(5),
+		},
+		{
+			name:       "SubsetSizeEqualToNumberOfEndpoints",
+			endpoints:  makeEndpoints(5),
+			subsetSize: 5,
+			want:       makeEndpoints(5),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b := &subsettingBalancer{
+				cfg:        &lbConfig{SubsetSize: tt.subsetSize},
+				hashSeed:   0,
+				hashDigest: xxhash.New(),
+			}
+			got := b.calculateSubset(tt.endpoints)
+			if diff := cmp.Diff(tt.want, got); diff != "" {
+				t.Errorf("calculateSubset() returned diff (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func (s) TestCalculateSubset_IgnoresEndpointsWithoutAddresses(t *testing.T) {
+	emptyEndpoint := resolver.Endpoint{Addresses: []resolver.Address{}}
+	endpoints := append([]resolver.Endpoint{emptyEndpoint}, makeEndpoints(3)...)
+	b := &subsettingBalancer{
+		cfg:        &lbConfig{SubsetSize: 2},
+		hashSeed:   0,
+		hashDigest: xxhash.New(),
+	}
+
+	got := b.calculateSubset(endpoints)
+	if len(got) != 2 {
+		t.Fatalf("calculateSubset() returned %d endpoints, want 2", len(got))
+	}
+	for _, endpoint := range got {
+		if len(endpoint.Addresses) == 0 {
+			t.Errorf("calculateSubset() returned an endpoint without addresses")
+		}
+	}
+}
+
+func (s) TestCalculateSubset_EndpointsRetainHashValues(t *testing.T) {
+	endpoints := makeEndpoints(10)
+	const subsetSize = 5
+	// The subset is deterministic based on the hash, so we can hardcode
+	// the expected output.
+	want := []resolver.Endpoint{
+		{Addresses: []resolver.Address{{Addr: "endpoint-6"}}},
+		{Addresses: []resolver.Address{{Addr: "endpoint-0"}}},
+		{Addresses: []resolver.Address{{Addr: "endpoint-1"}}},
+		{Addresses: []resolver.Address{{Addr: "endpoint-7"}}},
+		{Addresses: []resolver.Address{{Addr: "endpoint-3"}}},
+	}
+
+	b := &subsettingBalancer{
+		cfg:        &lbConfig{SubsetSize: subsetSize},
+		hashSeed:   0,
+		hashDigest: xxhash.New(),
+	}
+	for range 10 {
+		got := b.calculateSubset(endpoints)
+		if diff := cmp.Diff(want, got); diff != "" {
+			t.Fatalf("calculateSubset() returned diff (-want +got):\n%s", diff)
+		}
+	}
+}
+
+func (s) TestSubsettingBalancer_DeterministicSubset(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Register the stub balancer builder, which will be used as the child
+	// policy in the random_subsetting balancer.
+	updateCh := make(chan balancer.ClientConnState, 1)
+	stub.Register("stub-child-balancer", stub.BalancerFuncs{
+		UpdateClientConnState: func(_ *stub.BalancerData, s balancer.ClientConnState) error {
+			select {
+			case <-ctx.Done():
+			case updateCh <- s:
+			}
+			return nil
+		},
+	})
+
+	// Create a random_subsetting balancer.
+	tcc := testutils.NewBalancerClientConn(t)
+	rsb := balancer.Get(Name).Build(tcc, balancer.BuildOptions{})
+	defer rsb.Close()
+
+	// Prepare the configuration and resolver state to be passed to the
+	// random_subsetting balancer.
+	endpoints := makeEndpoints(10)
+	state := balancer.ClientConnState{
+		ResolverState: resolver.State{Endpoints: endpoints},
+		BalancerConfig: &lbConfig{
+			SubsetSize:  5,
+			ChildPolicy: &iserviceconfig.BalancerConfig{Name: "stub-child-balancer"},
+		},
+	}
+
+	// Send the resolver state to the random_subsetting balancer and verify that
+	// the child policy receives the expected number of endpoints.
+	if err := rsb.UpdateClientConnState(state); err != nil {
+		t.Fatalf("UpdateClientConnState failed: %v", err)
+	}
+
+	var wantEndpoints []resolver.Endpoint
+	select {
+	case s := <-updateCh:
+		if len(s.ResolverState.Endpoints) != 5 {
+			t.Fatalf("Child policy received %d endpoints, want 5", len(s.ResolverState.Endpoints))
+		}
+		// Store the subset for the next comparison.
+		wantEndpoints = s.ResolverState.Endpoints
+	case <-ctx.Done():
+		t.Fatal("Timed out waiting for child policy to receive an update")
+	}
+
+	// Call UpdateClientConnState again with the same configuration.
+	if err := rsb.UpdateClientConnState(state); err != nil {
+		t.Fatalf("Second UpdateClientConnState failed: %v", err)
+	}
+
+	// Verify that the child policy receives the same subset of endpoints.
+	select {
+	case s := <-updateCh:
+		if diff := cmp.Diff(wantEndpoints, s.ResolverState.Endpoints); diff != "" {
+			t.Fatalf("Child policy received a different subset of endpoints on second update, diff (-want +got):\n%s", diff)
+		}
+	case <-ctx.Done():
+		t.Fatal("Timed out waiting for second child policy update")
+	}
+}

balancer/ringhash/config.go

@@ -0,0 +1,77 @@
+/*
+ *
+ * Copyright 2021 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 ringhash
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"google.golang.org/grpc/internal/envconfig"
+	"google.golang.org/grpc/internal/metadata"
+	iringhash "google.golang.org/grpc/internal/ringhash"
+)
+
+const (
+	defaultMinSize         = 1024
+	defaultMaxSize         = 4096
+	ringHashSizeUpperBound = 8 * 1024 * 1024 // 8M
+)
+
+func parseConfig(c json.RawMessage) (*iringhash.LBConfig, error) {
+	var cfg iringhash.LBConfig
+	if err := json.Unmarshal(c, &cfg); err != nil {
+		return nil, err
+	}
+	if cfg.MinRingSize > ringHashSizeUpperBound {
+		return nil, fmt.Errorf("min_ring_size value of %d is greater than max supported value %d for this field", cfg.MinRingSize, ringHashSizeUpperBound)
+	}
+	if cfg.MaxRingSize > ringHashSizeUpperBound {
+		return nil, fmt.Errorf("max_ring_size value of %d is greater than max supported value %d for this field", cfg.MaxRingSize, ringHashSizeUpperBound)
+	}
+	if cfg.MinRingSize == 0 {
+		cfg.MinRingSize = defaultMinSize
+	}
+	if cfg.MaxRingSize == 0 {
+		cfg.MaxRingSize = defaultMaxSize
+	}
+	if cfg.MinRingSize > cfg.MaxRingSize {
+		return nil, fmt.Errorf("min %v is greater than max %v", cfg.MinRingSize, cfg.MaxRingSize)
+	}
+	if cfg.MinRingSize > envconfig.RingHashCap {
+		cfg.MinRingSize = envconfig.RingHashCap
+	}
+	if cfg.MaxRingSize > envconfig.RingHashCap {
+		cfg.MaxRingSize = envconfig.RingHashCap
+	}
+	if !envconfig.RingHashSetRequestHashKey {
+		cfg.RequestHashHeader = ""
+	}
+	if cfg.RequestHashHeader != "" {
+		cfg.RequestHashHeader = strings.ToLower(cfg.RequestHashHeader)
+		// See rules in https://github.com/grpc/proposal/blob/master/A76-ring-hash-improvements.md#explicitly-setting-the-request-hash-key
+		if err := metadata.ValidateKey(cfg.RequestHashHeader); err != nil {
+			return nil, fmt.Errorf("invalid requestHashHeader %q: %v", cfg.RequestHashHeader, err)
+		}
+		if strings.HasSuffix(cfg.RequestHashHeader, "-bin") {
+			return nil, fmt.Errorf("invalid requestHashHeader %q: key must not end with \"-bin\"", cfg.RequestHashHeader)
+		}
+	}
+	return &cfg, nil
+}

balancer/ringhash/config_test.go

@@ -0,0 +1,173 @@
+/*
+ *
+ * Copyright 2021 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 ringhash
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+
+	"google.golang.org/grpc/internal/envconfig"
+	iringhash "google.golang.org/grpc/internal/ringhash"
+	"google.golang.org/grpc/internal/testutils"
+)
+
+func (s) TestParseConfig(t *testing.T) {
+	tests := []struct {
+		name                string
+		js                  string
+		envConfigCap        uint64
+		requestHeaderEnvVar bool
+		want                *iringhash.LBConfig
+		wantErr             bool
+	}{
+		{
+			name:                "OK",
+			js:                  `{"minRingSize": 1, "maxRingSize": 2}`,
+			requestHeaderEnvVar: true,
+			want:                &iringhash.LBConfig{MinRingSize: 1, MaxRingSize: 2},
+		},
+		{
+			name:                "OK with default min",
+			js:                  `{"maxRingSize": 2000}`,
+			requestHeaderEnvVar: true,
+			want:                &iringhash.LBConfig{MinRingSize: defaultMinSize, MaxRingSize: 2000},
+		},
+		{
+			name:                "OK with default max",
+			js:                  `{"minRingSize": 2000}`,
+			requestHeaderEnvVar: true,
+			want:                &iringhash.LBConfig{MinRingSize: 2000, MaxRingSize: defaultMaxSize},
+		},
+		{
+			name:                "min greater than max",
+			js:                  `{"minRingSize": 10, "maxRingSize": 2}`,
+			requestHeaderEnvVar: true,
+			want:                nil,
+			wantErr:             true,
+		},
+		{
+			name:                "min greater than max greater than global limit",
+			js:                  `{"minRingSize": 6000, "maxRingSize": 5000}`,
+			requestHeaderEnvVar: true,
+			want:                nil,
+			wantErr:             true,
+		},
+		{
+			name:                "max greater than global limit",
+			js:                  `{"minRingSize": 1, "maxRingSize": 6000}`,
+			requestHeaderEnvVar: true,
+			want:                &iringhash.LBConfig{MinRingSize: 1, MaxRingSize: 4096},
+		},
+		{
+			name:                "min and max greater than global limit",
+			js:                  `{"minRingSize": 5000, "maxRingSize": 6000}`,
+			requestHeaderEnvVar: true,
+			want:                &iringhash.LBConfig{MinRingSize: 4096, MaxRingSize: 4096},
+		},
+		{
+			name:                "min and max less than raised global limit",
+			js:                  `{"minRingSize": 5000, "maxRingSize": 6000}`,
+			envConfigCap:        8000,
+			requestHeaderEnvVar: true,
+			want:                &iringhash.LBConfig{MinRingSize: 5000, MaxRingSize: 6000},
+		},
+		{
+			name:                "min and max greater than raised global limit",
+			js:                  `{"minRingSize": 10000, "maxRingSize": 10000}`,
+			envConfigCap:        8000,
+			requestHeaderEnvVar: true,
+			want:                &iringhash.LBConfig{MinRingSize: 8000, MaxRingSize: 8000},
+		},
+		{
+			name:                "min greater than upper bound",
+			js:                  `{"minRingSize": 8388610, "maxRingSize": 10}`,
+			requestHeaderEnvVar: true,
+			want:                nil,
+			wantErr:             true,
+		},
+		{
+			name:                "max greater than upper bound",
+			js:                  `{"minRingSize": 10, "maxRingSize": 8388610}`,
+			requestHeaderEnvVar: true,
+			want:                nil,
+			wantErr:             true,
+		},
+		{
+			name:                "request metadata key set",
+			js:                  `{"requestHashHeader": "x-foo"}`,
+			requestHeaderEnvVar: true,
+			want: &iringhash.LBConfig{
+				MinRingSize:       defaultMinSize,
+				MaxRingSize:       defaultMaxSize,
+				RequestHashHeader: "x-foo",
+			},
+		},
+		{
+			name:                "request metadata key set with uppercase letters",
+			js:                  `{"requestHashHeader": "x-FOO"}`,
+			requestHeaderEnvVar: true,
+			want: &iringhash.LBConfig{
+				MinRingSize:       defaultMinSize,
+				MaxRingSize:       defaultMaxSize,
+				RequestHashHeader: "x-foo",
+			},
+		},
+		{
+			name:                "invalid request hash header",
+			js:                  `{"requestHashHeader": "!invalid"}`,
+			requestHeaderEnvVar: true,
+			want:                nil,
+			wantErr:             true,
+		},
+		{
+			name:                "binary request hash header",
+			js:                  `{"requestHashHeader": "header-with-bin"}`,
+			requestHeaderEnvVar: true,
+			want:                nil,
+			wantErr:             true,
+		},
+		{
+			name:                "request hash header cleared when RingHashSetRequestHashKey env var is false",
+			js:                  `{"requestHashHeader": "x-foo"}`,
+			requestHeaderEnvVar: false,
+			want: &iringhash.LBConfig{
+				MinRingSize: defaultMinSize,
+				MaxRingSize: defaultMaxSize,
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.envConfigCap != 0 {
+				testutils.SetEnvConfig(t, &envconfig.RingHashCap, tt.envConfigCap)
+			}
+			testutils.SetEnvConfig(t, &envconfig.RingHashSetRequestHashKey, tt.requestHeaderEnvVar)
+			got, err := parseConfig(json.RawMessage(tt.js))
+			if (err != nil) != tt.wantErr {
+				t.Errorf("parseConfig() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if diff := cmp.Diff(got, tt.want); diff != "" {
+				t.Errorf("parseConfig() got unexpected output, diff (-got +want): %v", diff)
+			}
+		})
+	}
+}

balancer/ringhash/logging.go

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

balancer/ringhash/picker.go

@@ -0,0 +1,124 @@
+/*
+ *
+ * Copyright 2021 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 ringhash
+
+import (
+	"fmt"
+	"strings"
+
+	xxhash "github.com/cespare/xxhash/v2"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	iringhash "google.golang.org/grpc/internal/ringhash"
+	"google.golang.org/grpc/metadata"
+)
+
+type picker struct {
+	ring *ring
+
+	// endpointStates is a cache of endpoint states.
+	// The ringhash balancer stores endpoint states in a `resolver.EndpointMap`,
+	// with access guarded by `ringhashBalancer.mu`. The `endpointStates` cache
+	// in the picker helps avoid locking the ringhash balancer's mutex when
+	// reading the latest state at RPC time.
+	endpointStates map[string]endpointState // endpointState.hashKey -> endpointState
+
+	// requestHashHeader is the header key to look for the request hash. If it's
+	// empty, the request hash is expected to be set in the context via xDS.
+	// See gRFC A76.
+	requestHashHeader string
+
+	// hasEndpointInConnectingState is true if any of the endpoints is in
+	// CONNECTING.
+	hasEndpointInConnectingState bool
+
+	randUint64 func() uint64
+}
+
+func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	usingRandomHash := false
+	var requestHash uint64
+	if p.requestHashHeader == "" {
+		var ok bool
+		if requestHash, ok = iringhash.XDSRequestHash(info.Ctx); !ok {
+			return balancer.PickResult{}, fmt.Errorf("ringhash: expected xDS config selector to set the request hash")
+		}
+	} else {
+		md, ok := metadata.FromOutgoingContext(info.Ctx)
+		if !ok || len(md.Get(p.requestHashHeader)) == 0 {
+			requestHash = p.randUint64()
+			usingRandomHash = true
+		} else {
+			values := strings.Join(md.Get(p.requestHashHeader), ",")
+			requestHash = xxhash.Sum64String(values)
+		}
+	}
+
+	e := p.ring.pick(requestHash)
+	ringSize := len(p.ring.items)
+	if !usingRandomHash {
+		// Per gRFC A61, because of sticky-TF with PickFirst's auto reconnect on TF,
+		// we ignore all TF subchannels and find the first ring entry in READY,
+		// CONNECTING or IDLE.  If that entry is in IDLE, we need to initiate a
+		// connection. The idlePicker returned by the LazyLB or the new Pickfirst
+		// should do this automatically.
+		for i := 0; i < ringSize; i++ {
+			index := (e.idx + i) % ringSize
+			es := p.endpointState(p.ring.items[index])
+			switch es.state.ConnectivityState {
+			case connectivity.Ready, connectivity.Connecting, connectivity.Idle:
+				return es.state.Picker.Pick(info)
+			case connectivity.TransientFailure:
+			default:
+				panic(fmt.Sprintf("Found child balancer in unknown state: %v", es.state.ConnectivityState))
+			}
+		}
+	} else {
+		// If the picker has generated a random hash, it will walk the ring from
+		// this hash, and pick the first READY endpoint. If no endpoint is
+		// currently in CONNECTING state, it will trigger a connection attempt
+		// on at most one endpoint that is in IDLE state along the way. - A76
+		requestedConnection := p.hasEndpointInConnectingState
+		for i := 0; i < ringSize; i++ {
+			index := (e.idx + i) % ringSize
+			es := p.endpointState(p.ring.items[index])
+			if es.state.ConnectivityState == connectivity.Ready {
+				return es.state.Picker.Pick(info)
+			}
+			if !requestedConnection && es.state.ConnectivityState == connectivity.Idle {
+				requestedConnection = true
+				// If the SubChannel is in idle state, initiate a connection but
+				// continue to check other pickers to see if there is one in
+				// ready state.
+				es.exitIdle()
+			}
+		}
+		if requestedConnection {
+			return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+		}
+	}
+
+	// All children are in transient failure. Return the first failure.
+	return p.endpointState(e).state.Picker.Pick(info)
+}
+
+func (p *picker) endpointState(e *ringEntry) endpointState {
+	return p.endpointStates[e.hashKey]
+}

balancer/ringhash/picker_test.go

@@ -0,0 +1,301 @@
+/*
+ *
+ * Copyright 2021 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 ringhash
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	iringhash "google.golang.org/grpc/internal/ringhash"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/metadata"
+)
+
+var (
+	testSubConns []*testutils.TestSubConn
+	errPicker    = errors.New("picker in TransientFailure")
+)
+
+func init() {
+	for i := 0; i < 8; i++ {
+		testSubConns = append(testSubConns, testutils.NewTestSubConn(fmt.Sprint(i)))
+	}
+}
+
+// fakeChildPicker is used to mock pickers from child pickfirst balancers.
+type fakeChildPicker struct {
+	connectivityState connectivity.State
+	subConn           *testutils.TestSubConn
+	tfError           error
+}
+
+func (p *fakeChildPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	switch p.connectivityState {
+	case connectivity.Idle:
+		p.subConn.Connect()
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+	case connectivity.Connecting:
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+	case connectivity.Ready:
+		return balancer.PickResult{SubConn: p.subConn}, nil
+	default:
+		return balancer.PickResult{}, p.tfError
+	}
+}
+
+func testRingAndEndpointStates(states []connectivity.State) (*ring, map[string]endpointState) {
+	var items []*ringEntry
+	epStates := map[string]endpointState{}
+	for i, st := range states {
+		testSC := testSubConns[i]
+		items = append(items, &ringEntry{
+			idx:     i,
+			hash:    math.MaxUint64 / uint64(len(states)) * uint64(i),
+			hashKey: testSC.String(),
+		})
+		epState := endpointState{
+			state: balancer.State{
+				ConnectivityState: st,
+				Picker: &fakeChildPicker{
+					connectivityState: st,
+					tfError:           fmt.Errorf("%d: %w", i, errPicker),
+					subConn:           testSC,
+				},
+			},
+			exitIdle: func() { testSC.Connect() },
+		}
+		epStates[testSC.String()] = epState
+	}
+	return &ring{items: items}, epStates
+}
+
+func (s) TestPickerPickFirstTwo(t *testing.T) {
+	tests := []struct {
+		name               string
+		connectivityStates []connectivity.State
+		wantSC             balancer.SubConn
+		wantErr            error
+		wantSCToConnect    balancer.SubConn
+	}{
+		{
+			name:               "picked is Ready",
+			connectivityStates: []connectivity.State{connectivity.Ready, connectivity.Idle},
+			wantSC:             testSubConns[0],
+		},
+		{
+			name:               "picked is connecting, queue",
+			connectivityStates: []connectivity.State{connectivity.Connecting, connectivity.Idle},
+			wantErr:            balancer.ErrNoSubConnAvailable,
+		},
+		{
+			name:               "picked is Idle, connect and queue",
+			connectivityStates: []connectivity.State{connectivity.Idle, connectivity.Idle},
+			wantErr:            balancer.ErrNoSubConnAvailable,
+			wantSCToConnect:    testSubConns[0],
+		},
+		{
+			name:               "picked is TransientFailure, next is ready, return",
+			connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.Ready},
+			wantSC:             testSubConns[1],
+		},
+		{
+			name:               "picked is TransientFailure, next is connecting, queue",
+			connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.Connecting},
+			wantErr:            balancer.ErrNoSubConnAvailable,
+		},
+		{
+			name:               "picked is TransientFailure, next is Idle, connect and queue",
+			connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.Idle},
+			wantErr:            balancer.ErrNoSubConnAvailable,
+			wantSCToConnect:    testSubConns[1],
+		},
+		{
+			name:               "all are in TransientFailure, return picked failure",
+			connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.TransientFailure},
+			wantErr:            errPicker,
+		},
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ring, epStates := testRingAndEndpointStates(tt.connectivityStates)
+			p := &picker{
+				ring:           ring,
+				endpointStates: epStates,
+			}
+			got, err := p.Pick(balancer.PickInfo{
+				Ctx: iringhash.SetXDSRequestHash(ctx, 0), // always pick the first endpoint on the ring.
+			})
+			if (err != nil || tt.wantErr != nil) && !errors.Is(err, tt.wantErr) {
+				t.Errorf("Pick() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got.SubConn != tt.wantSC {
+				t.Errorf("Pick() got = %v, want picked SubConn: %v", got, tt.wantSC)
+			}
+			if sc := tt.wantSCToConnect; sc != nil {
+				select {
+				case <-sc.(*testutils.TestSubConn).ConnectCh:
+				case <-time.After(defaultTestShortTimeout):
+					t.Errorf("timeout waiting for Connect() from SubConn %v", sc)
+				}
+			}
+		})
+	}
+}
+
+func (s) TestPickerNoRequestHash(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	ring, epStates := testRingAndEndpointStates([]connectivity.State{connectivity.Ready})
+	p := &picker{
+		ring:           ring,
+		endpointStates: epStates,
+	}
+	if _, err := p.Pick(balancer.PickInfo{Ctx: ctx}); err == nil {
+		t.Errorf("Pick() should have failed with no request hash")
+	}
+}
+
+func (s) TestPickerRequestHashKey(t *testing.T) {
+	tests := []struct {
+		name         string
+		headerValues []string
+		expectedPick int
+	}{
+		{
+			name:         "header not set",
+			expectedPick: 0, // Random hash set to 0, which is within (MaxUint64 / 3 * 2, 0]
+		},
+		{
+			name:         "header empty",
+			headerValues: []string{""},
+			expectedPick: 0, // xxhash.Sum64String("value1,value2") is within (MaxUint64 / 3 * 2, 0]
+		},
+		{
+			name:         "header set to one value",
+			headerValues: []string{"some-value"},
+			expectedPick: 1, // xxhash.Sum64String("some-value") is within (0, MaxUint64 / 3]
+		},
+		{
+			name:         "header set to multiple values",
+			headerValues: []string{"value1", "value2"},
+			expectedPick: 2, // xxhash.Sum64String("value1,value2") is within (MaxUint64 / 3, MaxUint64 / 3 * 2]
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+			defer cancel()
+
+			ring, epStates := testRingAndEndpointStates(
+				[]connectivity.State{
+					connectivity.Ready,
+					connectivity.Ready,
+					connectivity.Ready,
+				})
+			headerName := "some-header"
+			p := &picker{
+				ring:              ring,
+				endpointStates:    epStates,
+				requestHashHeader: headerName,
+				randUint64:        func() uint64 { return 0 },
+			}
+			for _, v := range tt.headerValues {
+				ctx = metadata.AppendToOutgoingContext(ctx, headerName, v)
+			}
+			if res, err := p.Pick(balancer.PickInfo{Ctx: ctx}); err != nil {
+				t.Errorf("Pick() failed: %v", err)
+			} else if res.SubConn != testSubConns[tt.expectedPick] {
+				t.Errorf("Pick() got = %v, want SubConn: %v", res.SubConn, testSubConns[tt.expectedPick])
+			}
+		})
+	}
+}
+
+func (s) TestPickerRandomHash(t *testing.T) {
+	tests := []struct {
+		name                         string
+		hash                         uint64
+		connectivityStates           []connectivity.State
+		wantSC                       balancer.SubConn
+		wantErr                      error
+		wantSCToConnect              balancer.SubConn
+		hasEndpointInConnectingState bool
+	}{
+		{
+			name:               "header not set, picked is Ready",
+			connectivityStates: []connectivity.State{connectivity.Ready, connectivity.Idle},
+			wantSC:             testSubConns[0],
+		},
+		{
+			name:               "header not set, picked is Idle, another is Ready. Connect and pick Ready",
+			connectivityStates: []connectivity.State{connectivity.Idle, connectivity.Ready},
+			wantSC:             testSubConns[1],
+			wantSCToConnect:    testSubConns[0],
+		},
+		{
+			name:                         "header not set, picked is Idle, there is at least one Connecting",
+			connectivityStates:           []connectivity.State{connectivity.Connecting, connectivity.Idle},
+			wantErr:                      balancer.ErrNoSubConnAvailable,
+			hasEndpointInConnectingState: true,
+		},
+		{
+			name:               "header not set, all Idle or TransientFailure, connect",
+			connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.Idle},
+			wantErr:            balancer.ErrNoSubConnAvailable,
+			wantSCToConnect:    testSubConns[1],
+		},
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ring, epStates := testRingAndEndpointStates(tt.connectivityStates)
+			p := &picker{
+				ring:                         ring,
+				endpointStates:               epStates,
+				requestHashHeader:            "some-header",
+				hasEndpointInConnectingState: tt.hasEndpointInConnectingState,
+				randUint64:                   func() uint64 { return 0 }, // always return the first endpoint on the ring.
+			}
+			if got, err := p.Pick(balancer.PickInfo{Ctx: ctx}); err != tt.wantErr {
+				t.Errorf("Pick() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			} else if got.SubConn != tt.wantSC {
+				t.Errorf("Pick() got = %v, want picked SubConn: %v", got, tt.wantSC)
+			}
+			if sc := tt.wantSCToConnect; sc != nil {
+				select {
+				case <-sc.(*testutils.TestSubConn).ConnectCh:
+				case <-time.After(defaultTestShortTimeout):
+					t.Errorf("timeout waiting for Connect() from SubConn %v", sc)
+				}
+			}
+		})
+	}
+}

balancer/ringhash/ring.go

@@ -0,0 +1,189 @@
+/*
+ *
+ * Copyright 2021 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 ringhash
+
+import (
+	"math"
+	"sort"
+	"strconv"
+
+	xxhash "github.com/cespare/xxhash/v2"
+	"google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/resolver"
+)
+
+type ring struct {
+	items []*ringEntry
+}
+
+type endpointInfo struct {
+	hashKey        string
+	scaledWeight   float64
+	originalWeight uint32
+}
+
+type ringEntry struct {
+	idx     int
+	hash    uint64
+	hashKey string
+	weight  uint32
+}
+
+// newRing creates a ring from the endpoints stored in the EndpointMap. The ring
+// size is limited by the passed in max/min.
+//
+// ring entries will be created for each endpoint, and endpoints with high
+// weight (specified by the endpoint) may have multiple entries.
+//
+// For example, for endpoints with weights {a:3, b:3, c:4}, a generated ring of
+// size 10 could be:
+// - {idx:0 hash:3689675255460411075  b}
+// - {idx:1 hash:4262906501694543955  c}
+// - {idx:2 hash:5712155492001633497  c}
+// - {idx:3 hash:8050519350657643659  b}
+// - {idx:4 hash:8723022065838381142  b}
+// - {idx:5 hash:11532782514799973195 a}
+// - {idx:6 hash:13157034721563383607 c}
+// - {idx:7 hash:14468677667651225770 c}
+// - {idx:8 hash:17336016884672388720 a}
+// - {idx:9 hash:18151002094784932496 a}
+//
+// To pick from a ring, a binary search will be done for the given target hash,
+// and first item with hash >= given hash will be returned.
+//
+// Must be called with a non-empty endpoints map.
+func newRing(endpoints *resolver.EndpointMap[*endpointState], minRingSize, maxRingSize uint64, logger *grpclog.PrefixLogger) *ring {
+	if logger.V(2) {
+		logger.Infof("newRing: number of endpoints is %d, minRingSize is %d, maxRingSize is %d", endpoints.Len(), minRingSize, maxRingSize)
+	}
+
+	// https://github.com/envoyproxy/envoy/blob/765c970f06a4c962961a0e03a467e165b276d50f/source/common/upstream/ring_hash_lb.cc#L114
+	normalizedWeights, minWeight := normalizeWeights(endpoints)
+	if logger.V(2) {
+		logger.Infof("newRing: normalized endpoint weights is %v", normalizedWeights)
+	}
+
+	// Normalized weights for {3,3,4} is {0.3,0.3,0.4}.
+
+	// Scale up the size of the ring such that the least-weighted host gets a
+	// whole number of hashes on the ring.
+	//
+	// Note that size is limited by the input max/min.
+	scale := math.Min(math.Ceil(minWeight*float64(minRingSize))/minWeight, float64(maxRingSize))
+	ringSize := math.Ceil(scale)
+	items := make([]*ringEntry, 0, int(ringSize))
+	if logger.V(2) {
+		logger.Infof("newRing: creating new ring of size %v", ringSize)
+	}
+
+	// For each entry, scale*weight nodes are generated in the ring.
+	//
+	// Not all of these are whole numbers. E.g. for weights {a:3,b:3,c:4}, if
+	// ring size is 7, scale is 6.66. The numbers of nodes will be
+	// {a,a,b,b,c,c,c}.
+	//
+	// A hash is generated for each item, and later the results will be sorted
+	// based on the hash.
+	var currentHashes, targetHashes float64
+	for _, epInfo := range normalizedWeights {
+		targetHashes += scale * epInfo.scaledWeight
+		// This index ensures that ring entries corresponding to the same
+		// endpoint hash to different values. And since this index is
+		// per-endpoint, these entries hash to the same value across address
+		// updates.
+		idx := 0
+		for currentHashes < targetHashes {
+			h := xxhash.Sum64String(epInfo.hashKey + "_" + strconv.Itoa(idx))
+			items = append(items, &ringEntry{hash: h, hashKey: epInfo.hashKey, weight: epInfo.originalWeight})
+			idx++
+			currentHashes++
+		}
+	}
+
+	// Sort items based on hash, to prepare for binary search.
+	sort.Slice(items, func(i, j int) bool { return items[i].hash < items[j].hash })
+	for i, ii := range items {
+		ii.idx = i
+	}
+	return &ring{items: items}
+}
+
+// normalizeWeights calculates the normalized weights for each endpoint in the
+// given endpoints map. It returns a slice of endpointWithState structs, where
+// each struct contains the picker for an endpoint and its corresponding weight.
+// The function also returns the minimum weight among all endpoints.
+//
+// The normalized weight of each endpoint is calculated by dividing its weight
+// attribute by the sum of all endpoint weights. If the weight attribute is not
+// found on the endpoint, a default weight of 1 is used.
+//
+// The endpoints are sorted in ascending order to ensure consistent results.
+//
+// Must be called with a non-empty endpoints map.
+func normalizeWeights(endpoints *resolver.EndpointMap[*endpointState]) ([]endpointInfo, float64) {
+	var weightSum uint32
+	// Since attributes are explicitly ignored in the EndpointMap key, we need
+	// to iterate over the values to get the weights.
+	endpointVals := endpoints.Values()
+	for _, epState := range endpointVals {
+		weightSum += epState.weight
+	}
+	ret := make([]endpointInfo, 0, endpoints.Len())
+	min := 1.0
+	for _, epState := range endpointVals {
+		// (*endpointState).weight is set to 1 if the weight attribute is not
+		// found on the endpoint. And since this function is guaranteed to be
+		// called with a non-empty endpoints map, weightSum is guaranteed to be
+		// non-zero. So, we need not worry about divide by zero error here.
+		nw := float64(epState.weight) / float64(weightSum)
+		ret = append(ret, endpointInfo{
+			hashKey:        epState.hashKey,
+			scaledWeight:   nw,
+			originalWeight: epState.weight,
+		})
+		min = math.Min(min, nw)
+	}
+	// Sort the endpoints to return consistent results.
+	//
+	// Note: this might not be necessary, but this makes sure the ring is
+	// consistent as long as the endpoints are the same, for example, in cases
+	// where an endpoint is added and then removed, the RPCs will still pick the
+	// same old endpoint.
+	sort.Slice(ret, func(i, j int) bool {
+		return ret[i].hashKey < ret[j].hashKey
+	})
+	return ret, min
+}
+
+// pick does a binary search. It returns the item with smallest index i that
+// r.items[i].hash >= h.
+func (r *ring) pick(h uint64) *ringEntry {
+	i := sort.Search(len(r.items), func(i int) bool { return r.items[i].hash >= h })
+	if i == len(r.items) {
+		// If not found, and h is greater than the largest hash, return the
+		// first item.
+		i = 0
+	}
+	return r.items[i]
+}
+
+// next returns the next entry.
+func (r *ring) next(e *ringEntry) *ringEntry {
+	return r.items[(e.idx+1)%len(r.items)]
+}

balancer/ringhash/ring_test.go

@@ -0,0 +1,112 @@
+/*
+ *
+ * Copyright 2021 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 ringhash
+
+import (
+	"fmt"
+	"math"
+	"testing"
+
+	xxhash "github.com/cespare/xxhash/v2"
+	"google.golang.org/grpc/experimental/balancer/weight"
+	"google.golang.org/grpc/resolver"
+)
+
+var testEndpoints []resolver.Endpoint
+var testEndpointStateMap *resolver.EndpointMap[*endpointState]
+
+func init() {
+	testEndpoints = []resolver.Endpoint{
+		testEndpoint("a", 3),
+		testEndpoint("b", 3),
+		testEndpoint("c", 4),
+	}
+	testEndpointStateMap = resolver.NewEndpointMap[*endpointState]()
+	testEndpointStateMap.Set(testEndpoints[0], &endpointState{hashKey: "a", weight: 3})
+	testEndpointStateMap.Set(testEndpoints[1], &endpointState{hashKey: "b", weight: 3})
+	testEndpointStateMap.Set(testEndpoints[2], &endpointState{hashKey: "c", weight: 4})
+}
+
+func testEndpoint(addr string, endpointWeight uint32) resolver.Endpoint {
+	ep := resolver.Endpoint{Addresses: []resolver.Address{{Addr: addr}}}
+	return weight.Set(ep, weight.EndpointInfo{Weight: endpointWeight})
+}
+
+func (s) TestRingNew(t *testing.T) {
+	var totalWeight float64 = 10
+	for _, min := range []uint64{3, 4, 6, 8} {
+		for _, max := range []uint64{20, 8} {
+			t.Run(fmt.Sprintf("size-min-%v-max-%v", min, max), func(t *testing.T) {
+				r := newRing(testEndpointStateMap, min, max, nil)
+				totalCount := len(r.items)
+				if totalCount < int(min) || totalCount > int(max) {
+					t.Fatalf("unexpected size %v, want min %v, max %v", totalCount, min, max)
+				}
+				for _, e := range testEndpoints {
+					var count int
+					for _, ii := range r.items {
+						if ii.hashKey == hashKey(e) {
+							count++
+						}
+					}
+					got := float64(count) / float64(totalCount)
+					want := float64(getWeightAttribute(e)) / totalWeight
+					if !equalApproximately(got, want) {
+						t.Fatalf("unexpected item weight in ring: %v != %v", got, want)
+					}
+				}
+			})
+		}
+	}
+}
+
+func equalApproximately(x, y float64) bool {
+	delta := math.Abs(x - y)
+	mean := math.Abs(x+y) / 2.0
+	return delta/mean < 0.25
+}
+
+func (s) TestRingPick(t *testing.T) {
+	r := newRing(testEndpointStateMap, 10, 20, nil)
+	for _, h := range []uint64{xxhash.Sum64String("1"), xxhash.Sum64String("2"), xxhash.Sum64String("3"), xxhash.Sum64String("4")} {
+		t.Run(fmt.Sprintf("picking-hash-%v", h), func(t *testing.T) {
+			e := r.pick(h)
+			var low uint64
+			if e.idx > 0 {
+				low = r.items[e.idx-1].hash
+			}
+			high := e.hash
+			// h should be in [low, high).
+			if h < low || h >= high {
+				t.Fatalf("unexpected item picked, hash: %v, low: %v, high: %v", h, low, high)
+			}
+		})
+	}
+}
+
+func (s) TestRingNext(t *testing.T) {
+	r := newRing(testEndpointStateMap, 10, 20, nil)
+
+	for _, e := range r.items {
+		ne := r.next(e)
+		if ne.idx != (e.idx+1)%len(r.items) {
+			t.Fatalf("next(%+v) returned unexpected %+v", e, ne)
+		}
+	}
+}

balancer/ringhash/ringhash.go

@@ -0,0 +1,410 @@
+/*
+ *
+ * Copyright 2021 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 ringhash implements the ringhash balancer. See the following
+// gRFCs for details:
+// - https://github.com/grpc/proposal/blob/master/A42-xds-ring-hash-lb-policy.md
+// - https://github.com/grpc/proposal/blob/master/A61-IPv4-IPv6-dualstack-backends.md#ring-hash
+// - https://github.com/grpc/proposal/blob/master/A76-ring-hash-improvements.md
+//
+// # Experimental
+//
+// Notice: This package is EXPERIMENTAL and may be changed or removed in a
+// later release.
+package ringhash
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"math/rand/v2"
+	"sort"
+	"sync"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/lazy"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/experimental/balancer/weight"
+	"google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/pretty"
+	iringhash "google.golang.org/grpc/internal/ringhash"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/ringhash"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// Name is the name of the ring_hash balancer.
+const Name = "ring_hash_experimental"
+
+func lazyPickFirstBuilder(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	return lazy.NewBalancer(cc, opts, balancer.Get(pickfirst.Name).Build)
+}
+
+func init() {
+	balancer.Register(bb{})
+}
+
+type bb struct{}
+
+func (bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	b := &ringhashBalancer{
+		ClientConn:     cc,
+		endpointStates: resolver.NewEndpointMap[*endpointState](),
+	}
+	esOpts := endpointsharding.Options{DisableAutoReconnect: true}
+	b.child = endpointsharding.NewBalancer(b, opts, lazyPickFirstBuilder, esOpts)
+	b.logger = prefixLogger(b)
+	b.logger.Infof("Created")
+	return b
+}
+
+func (bb) Name() string {
+	return Name
+}
+
+func (bb) ParseConfig(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	return parseConfig(c)
+}
+
+type ringhashBalancer struct {
+	// The following fields are initialized at build time and read-only after
+	// that and therefore do not need to be guarded by a mutex.
+
+	// ClientConn is embedded to intercept UpdateState calls from the child
+	// endpointsharding balancer.
+	balancer.ClientConn
+	logger *grpclog.PrefixLogger
+	child  balancer.Balancer
+
+	mu                   sync.Mutex
+	config               *iringhash.LBConfig
+	inhibitChildUpdates  bool
+	shouldRegenerateRing bool
+	endpointStates       *resolver.EndpointMap[*endpointState]
+
+	// ring is always in sync with endpoints. When endpoints change, a new ring
+	// is generated. Note that address weights updates also regenerates the
+	// ring.
+	ring *ring
+}
+
+// hashKey returns the hash key to use for an endpoint. Per gRFC A61, each entry
+// in the ring is a hash of the endpoint's hash key concatenated with a
+// per-entry unique suffix.
+func hashKey(endpoint resolver.Endpoint) string {
+	if hk := ringhash.HashKey(endpoint); hk != "" {
+		return hk
+	}
+	// If no hash key is set, use the endpoint's first address as the hash key.
+	// This is the default behavior when no hash key is set.
+	return endpoint.Addresses[0].Addr
+}
+
+// UpdateState intercepts child balancer state updates. It updates the
+// per-endpoint state stored in the ring, and also the aggregated state based on
+// the child picker. It also reconciles the endpoint list. It sets
+// `b.shouldRegenerateRing` to true if the new endpoint list is different from
+// the previous, i.e. any of the following is true:
+// - an endpoint was added
+// - an endpoint was removed
+// - an endpoint's weight was updated
+// - the first addresses of the endpoint has changed
+func (b *ringhashBalancer) UpdateState(state balancer.State) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	childStates := endpointsharding.ChildStatesFromPicker(state.Picker)
+	// endpointsSet is the set converted from endpoints, used for quick lookup.
+	endpointsSet := resolver.NewEndpointMap[bool]()
+
+	for _, childState := range childStates {
+		endpoint := childState.Endpoint
+		endpointsSet.Set(endpoint, true)
+		newWeight := getWeightAttribute(endpoint)
+		hk := hashKey(endpoint)
+		es, ok := b.endpointStates.Get(endpoint)
+		if !ok {
+			es := &endpointState{
+				hashKey:  hk,
+				weight:   newWeight,
+				state:    childState.State,
+				exitIdle: childState.ExitIdle,
+			}
+			b.endpointStates.Set(endpoint, es)
+			b.shouldRegenerateRing = true
+		} else {
+			// We have seen this endpoint before and created a `endpointState`
+			// object for it. If the weight or the hash key of the endpoint has
+			// changed, update the endpoint state map with the new weight or
+			// hash key. This will be used when a new ring is created.
+			if oldWeight := es.weight; oldWeight != newWeight {
+				b.shouldRegenerateRing = true
+				es.weight = newWeight
+			}
+			if es.hashKey != hk {
+				b.shouldRegenerateRing = true
+				es.hashKey = hk
+			}
+			es.state = childState.State
+		}
+	}
+
+	for endpoint := range b.endpointStates.All() {
+		if _, ok := endpointsSet.Get(endpoint); ok {
+			continue
+		}
+		// endpoint was removed by resolver.
+		b.endpointStates.Delete(endpoint)
+		b.shouldRegenerateRing = true
+	}
+
+	b.updatePickerLocked()
+}
+
+func (b *ringhashBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	if b.logger.V(2) {
+		b.logger.Infof("Received update from resolver, balancer config: %+v", pretty.ToJSON(ccs.BalancerConfig))
+	}
+
+	newConfig, ok := ccs.BalancerConfig.(*iringhash.LBConfig)
+	if !ok {
+		return fmt.Errorf("unexpected balancer config with type: %T", ccs.BalancerConfig)
+	}
+
+	b.mu.Lock()
+	b.inhibitChildUpdates = true
+	b.mu.Unlock()
+
+	defer func() {
+		b.mu.Lock()
+		b.inhibitChildUpdates = false
+		b.updatePickerLocked()
+		b.mu.Unlock()
+	}()
+
+	if err := b.child.UpdateClientConnState(balancer.ClientConnState{
+		// Make pickfirst children use health listeners for outlier detection
+		// and health checking to work.
+		ResolverState: pickfirst.EnableHealthListener(ccs.ResolverState),
+	}); err != nil {
+		return err
+	}
+
+	b.mu.Lock()
+	// Ring updates can happen due to the following:
+	// 1. Addition or deletion of endpoints: The synchronous picker update from
+	//    the child endpointsharding balancer would contain the list of updated
+	//    endpoints.  Updates triggered by the child after handling the
+	//    `UpdateClientConnState` call will not change the endpoint list.
+	// 2. Change in the `LoadBalancerConfig`: Ring config such as max/min ring
+	//    size.
+	// To avoid extra ring updates, a boolean is used to track the need for a
+	// ring update and the update is done only once at the end.
+	//
+	// If the ring configuration has changed, we need to regenerate the ring
+	// while sending a new picker.
+	if b.config == nil || b.config.MinRingSize != newConfig.MinRingSize || b.config.MaxRingSize != newConfig.MaxRingSize {
+		b.shouldRegenerateRing = true
+	}
+	b.config = newConfig
+	b.mu.Unlock()
+	return nil
+}
+
+func (b *ringhashBalancer) ResolverError(err error) {
+	b.child.ResolverError(err)
+}
+
+func (b *ringhashBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+func (b *ringhashBalancer) updatePickerLocked() {
+	state := b.aggregatedStateLocked()
+	// Start connecting to new endpoints if necessary.
+	if state == connectivity.Connecting || state == connectivity.TransientFailure {
+		// When overall state is TransientFailure, we need to make sure at least
+		// one endpoint is attempting to connect, otherwise this balancer may
+		// never get picks if the parent is priority.
+		//
+		// Because we report Connecting as the overall state when only one
+		// endpoint is in TransientFailure, we do the same check for Connecting
+		// here.
+		//
+		// Note that this check also covers deleting endpoints. E.g. if the
+		// endpoint attempting to connect is deleted, and the overall state is
+		// TF. Since there must be at least one endpoint attempting to connect,
+		// we need to trigger one.
+		//
+		// After calling `ExitIdle` on a child balancer, the child will send a
+		// picker update asynchronously. A race condition may occur if another
+		// picker update from endpointsharding arrives before the child's
+		// picker update. The received picker may trigger a re-execution of the
+		// loop below to find an idle child. Since map iteration order is
+		// non-deterministic, the list of `endpointState`s must be sorted to
+		// ensure `ExitIdle` is called on the same child, preventing unnecessary
+		// connections.
+		endpointStates := make([]*endpointState, 0, b.endpointStates.Len())
+		for _, s := range b.endpointStates.All() {
+			endpointStates = append(endpointStates, s)
+		}
+		sort.Slice(endpointStates, func(i, j int) bool {
+			return endpointStates[i].hashKey < endpointStates[j].hashKey
+		})
+
+		// Store the function to ExitIdle on the first IDLE endpoint.
+		var exitIdle func()
+		for _, es := range endpointStates {
+			connState := es.state.ConnectivityState
+			if connState == connectivity.Connecting {
+				exitIdle = nil
+				break
+			}
+			if exitIdle == nil && connState == connectivity.Idle {
+				exitIdle = es.exitIdle
+			}
+		}
+		if exitIdle != nil {
+			exitIdle()
+		}
+	}
+
+	if b.inhibitChildUpdates {
+		return
+	}
+
+	// Update the channel.
+	if b.endpointStates.Len() > 0 && b.shouldRegenerateRing {
+		// with a non-empty list of endpoints.
+		b.ring = newRing(b.endpointStates, b.config.MinRingSize, b.config.MaxRingSize, b.logger)
+	}
+	b.shouldRegenerateRing = false
+	var newPicker balancer.Picker
+	if b.endpointStates.Len() == 0 {
+		newPicker = base.NewErrPicker(errors.New("produced zero addresses"))
+	} else {
+		newPicker = b.newPickerLocked()
+	}
+	b.ClientConn.UpdateState(balancer.State{
+		ConnectivityState: state,
+		Picker:            newPicker,
+	})
+}
+
+func (b *ringhashBalancer) Close() {
+	b.logger.Infof("Shutdown")
+	b.child.Close()
+}
+
+func (b *ringhashBalancer) ExitIdle() {
+	// ExitIdle implementation is a no-op because connections are either
+	// triggers from picks or from child balancer state changes.
+}
+
+// newPickerLocked generates a picker. The picker copies the endpoint states
+// over to avoid locking the mutex at RPC time. The picker should be
+// re-generated every time an endpoint state is updated.
+func (b *ringhashBalancer) newPickerLocked() *picker {
+	states := make(map[string]endpointState)
+	hasEndpointConnecting := false
+	for _, epState := range b.endpointStates.All() {
+		// Copy the endpoint state to avoid races, since ring hash
+		// mutates the state, weight and hash key in place.
+		states[epState.hashKey] = *epState
+		if epState.state.ConnectivityState == connectivity.Connecting {
+			hasEndpointConnecting = true
+		}
+	}
+	return &picker{
+		ring:                         b.ring,
+		endpointStates:               states,
+		requestHashHeader:            b.config.RequestHashHeader,
+		hasEndpointInConnectingState: hasEndpointConnecting,
+		randUint64:                   rand.Uint64,
+	}
+}
+
+// aggregatedStateLocked returns the aggregated child balancers state
+// based on the following rules.
+//   - If there is at least one endpoint in READY state, report READY.
+//   - If there are 2 or more endpoints in TRANSIENT_FAILURE state, report
+//     TRANSIENT_FAILURE.
+//   - If there is at least one endpoint in CONNECTING state, report CONNECTING.
+//   - If there is one endpoint in TRANSIENT_FAILURE and there is more than one
+//     endpoint, report state CONNECTING.
+//   - If there is at least one endpoint in Idle state, report Idle.
+//   - Otherwise, report TRANSIENT_FAILURE.
+//
+// Note that if there are 1 connecting, 2 transient failure, the overall state
+// is transient failure. This is because the second transient failure is a
+// fallback of the first failing endpoint, and we want to report transient
+// failure to failover to the lower priority.
+func (b *ringhashBalancer) aggregatedStateLocked() connectivity.State {
+	var nums [5]int
+	for _, es := range b.endpointStates.All() {
+		nums[es.state.ConnectivityState]++
+	}
+
+	if nums[connectivity.Ready] > 0 {
+		return connectivity.Ready
+	}
+	if nums[connectivity.TransientFailure] > 1 {
+		return connectivity.TransientFailure
+	}
+	if nums[connectivity.Connecting] > 0 {
+		return connectivity.Connecting
+	}
+	if nums[connectivity.TransientFailure] == 1 && b.endpointStates.Len() > 1 {
+		return connectivity.Connecting
+	}
+	if nums[connectivity.Idle] > 0 {
+		return connectivity.Idle
+	}
+	return connectivity.TransientFailure
+}
+
+// getWeightAttribute is a convenience function which returns the value of the
+// weight endpoint Attribute.
+//
+// When used in the xDS context, the weight attribute is guaranteed to be
+// non-zero. But, when used in a non-xDS context, the weight attribute could be
+// unset. A Default of 1 is used in the latter case.
+func getWeightAttribute(e resolver.Endpoint) uint32 {
+	w := weight.FromEndpoint(e).Weight
+	if w == 0 {
+		return 1
+	}
+	return w
+}
+
+type endpointState struct {
+	// hashKey is the hash key of the endpoint. Per gRFC A61, each entry in the
+	// ring is an endpoint, positioned based on the hash of the endpoint's first
+	// address by default. Per gRFC A76, the hash key of an endpoint may be
+	// overridden, for example based on EDS endpoint metadata.
+	hashKey  string
+	weight   uint32
+	exitIdle func()
+
+	// state is updated by the balancer while receiving resolver updates from
+	// the channel and picker updates from its children. Access to it is guarded
+	// by ringhashBalancer.mu.
+	state balancer.State
+}

balancer/ringhash/ringhash_test.go

@@ -0,0 +1,739 @@
+/*
+ *
+ * Copyright 2021 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 ringhash
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/experimental/balancer/weight"
+	"google.golang.org/grpc/internal/grpctest"
+	iringhash "google.golang.org/grpc/internal/ringhash"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/resolver"
+)
+
+const (
+	defaultTestTimeout      = 10 * time.Second
+	defaultTestShortTimeout = 10 * time.Millisecond
+
+	testBackendAddrsCount = 12
+)
+
+var (
+	testBackendAddrStrs []string
+	testConfig          = &iringhash.LBConfig{MinRingSize: 1, MaxRingSize: 10}
+)
+
+func init() {
+	for i := 0; i < testBackendAddrsCount; i++ {
+		testBackendAddrStrs = append(testBackendAddrStrs, fmt.Sprintf("%d.%d.%d.%d:%d", i, i, i, i, i))
+	}
+}
+
+// setupTest creates the balancer, and does an initial sanity check.
+func setupTest(t *testing.T, endpoints []resolver.Endpoint) (*testutils.BalancerClientConn, balancer.Balancer, balancer.Picker) {
+	t.Helper()
+	cc := testutils.NewBalancerClientConn(t)
+	builder := balancer.Get(Name)
+	b := builder.Build(cc, balancer.BuildOptions{})
+	if b == nil {
+		t.Fatalf("builder.Build(%s) failed and returned nil", Name)
+	}
+	t.Cleanup(func() { b.Close() })
+
+	if err := b.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState:  resolver.State{Endpoints: endpoints},
+		BalancerConfig: testConfig,
+	}); err != nil {
+		t.Fatalf("UpdateClientConnState returned err: %v", err)
+	}
+
+	// The leaf pickfirst are created lazily, only when their endpoint is picked
+	// or other endpoints are in TF. No SubConns should be created immediately.
+	select {
+	case sc := <-cc.NewSubConnCh:
+		t.Errorf("unexpected SubConn creation: %v", sc)
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	// Should also have a picker, with all endpoints in Idle.
+	p1 := <-cc.NewPickerCh
+
+	ringHashPicker := p1.(*picker)
+	if got, want := len(ringHashPicker.endpointStates), len(endpoints); got != want {
+		t.Errorf("Number of child balancers = %d, want = %d", got, want)
+	}
+	for firstAddr, bs := range ringHashPicker.endpointStates {
+		if got, want := bs.state.ConnectivityState, connectivity.Idle; got != want {
+			t.Errorf("Child balancer connectivity state for address %q = %v, want = %v", firstAddr, got, want)
+		}
+	}
+	return cc, b, p1
+}
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// TestUpdateClientConnState_NewRingSize tests the scenario where the ringhash
+// LB policy receives new configuration which specifies new values for the ring
+// min and max sizes. The test verifies that a new ring is created and a new
+// picker is sent to the ClientConn.
+func (s) TestUpdateClientConnState_NewRingSize(t *testing.T) {
+	origMinRingSize, origMaxRingSize := 1, 10 // Configured from `testConfig` in `setupTest`
+	newMinRingSize, newMaxRingSize := 20, 100
+
+	endpoints := []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}}}
+	cc, b, p1 := setupTest(t, endpoints)
+	ring1 := p1.(*picker).ring
+	if ringSize := len(ring1.items); ringSize < origMinRingSize || ringSize > origMaxRingSize {
+		t.Fatalf("Ring created with size %d, want between [%d, %d]", ringSize, origMinRingSize, origMaxRingSize)
+	}
+
+	if err := b.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState: resolver.State{Endpoints: endpoints},
+		BalancerConfig: &iringhash.LBConfig{
+			MinRingSize: uint64(newMinRingSize),
+			MaxRingSize: uint64(newMaxRingSize),
+		},
+	}); err != nil {
+		t.Fatalf("UpdateClientConnState returned err: %v", err)
+	}
+
+	var ring2 *ring
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Timeout when waiting for a picker update after a configuration update")
+	case p2 := <-cc.NewPickerCh:
+		ring2 = p2.(*picker).ring
+	}
+	if ringSize := len(ring2.items); ringSize < newMinRingSize || ringSize > newMaxRingSize {
+		t.Fatalf("Ring created with size %d, want between [%d, %d]", ringSize, newMinRingSize, newMaxRingSize)
+	}
+}
+
+func (s) TestOneEndpoint(t *testing.T) {
+	wantAddr1 := resolver.Address{Addr: testBackendAddrStrs[0]}
+	cc, _, p0 := setupTest(t, []resolver.Endpoint{{Addresses: []resolver.Address{wantAddr1}}})
+	ring0 := p0.(*picker).ring
+
+	firstHash := ring0.items[0].hash
+	// firstHash-1 will pick the first (and only) SubConn from the ring.
+	testHash := firstHash - 1
+	// The first pick should be queued, and should trigger a connection to the
+	// only Endpoint which has a single address.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)}); err != balancer.ErrNoSubConnAvailable {
+		t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
+	}
+	var sc0 *testutils.TestSubConn
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case sc0 = <-cc.NewSubConnCh:
+	}
+	if got, want := sc0.Addresses[0].Addr, wantAddr1.Addr; got != want {
+		t.Fatalf("SubConn.Addresses = %v, want = %v", got, want)
+	}
+	select {
+	case <-sc0.ConnectCh:
+	case <-time.After(defaultTestTimeout):
+		t.Errorf("timeout waiting for Connect() from SubConn %v", sc0)
+	}
+
+	// Send state updates to Ready.
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	if err := cc.WaitForConnectivityState(ctx, connectivity.Ready); err != nil {
+		t.Fatal(err)
+	}
+
+	// Test pick with one backend.
+	p1 := <-cc.NewPickerCh
+	for i := 0; i < 5; i++ {
+		gotSCSt, _ := p1.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
+		if gotSCSt.SubConn != sc0 {
+			t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc0)
+		}
+	}
+}
+
+// TestThreeBackendsAffinity covers that there are 3 SubConns, RPCs with the
+// same hash always pick the same SubConn. When the one picked is down, another
+// one will be picked.
+func (s) TestThreeSubConnsAffinity(t *testing.T) {
+	endpoints := []resolver.Endpoint{
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[1]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[2]}}},
+	}
+	remainingAddrs := map[string]bool{
+		testBackendAddrStrs[0]: true,
+		testBackendAddrStrs[1]: true,
+		testBackendAddrStrs[2]: true,
+	}
+	cc, _, p0 := setupTest(t, endpoints)
+	// This test doesn't update addresses, so this ring will be used by all the
+	// pickers.
+	ring := p0.(*picker).ring
+
+	firstHash := ring.items[0].hash
+	// firstHash+1 will pick the second endpoint from the ring.
+	testHash := firstHash + 1
+	// The first pick should be queued, and should trigger Connect() on the only
+	// SubConn.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)}); err != balancer.ErrNoSubConnAvailable {
+		t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
+	}
+
+	// The picked endpoint should be the second in the ring.
+	var subConns [3]*testutils.TestSubConn
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case subConns[1] = <-cc.NewSubConnCh:
+	}
+	if got, want := subConns[1].Addresses[0].Addr, ring.items[1].hashKey; got != want {
+		t.Fatalf("SubConn.Address = %v, want = %v", got, want)
+	}
+	select {
+	case <-subConns[1].ConnectCh:
+	case <-time.After(defaultTestTimeout):
+		t.Errorf("timeout waiting for Connect() from SubConn %v", subConns[1])
+	}
+	delete(remainingAddrs, ring.items[1].hashKey)
+
+	// Turn down the subConn in use.
+	subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
+
+	// This should trigger a connection to a new endpoint.
+	<-cc.NewPickerCh
+	var sc *testutils.TestSubConn
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case sc = <-cc.NewSubConnCh:
+	}
+	scAddr := sc.Addresses[0].Addr
+	if _, ok := remainingAddrs[scAddr]; !ok {
+		t.Fatalf("New SubConn created with previously used address: %q", scAddr)
+	}
+	delete(remainingAddrs, scAddr)
+	select {
+	case <-sc.ConnectCh:
+	case <-time.After(defaultTestTimeout):
+		t.Errorf("timeout waiting for Connect() from SubConn %v", subConns[1])
+	}
+	if scAddr == ring.items[0].hashKey {
+		subConns[0] = sc
+	} else if scAddr == ring.items[2].hashKey {
+		subConns[2] = sc
+	}
+
+	// Turning down the SubConn should cause creation of a connection to the
+	// final endpoint.
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case sc = <-cc.NewSubConnCh:
+	}
+	scAddr = sc.Addresses[0].Addr
+	if _, ok := remainingAddrs[scAddr]; !ok {
+		t.Fatalf("New SubConn created with previously used address: %q", scAddr)
+	}
+	delete(remainingAddrs, scAddr)
+	select {
+	case <-sc.ConnectCh:
+	case <-time.After(defaultTestTimeout):
+		t.Errorf("timeout waiting for Connect() from SubConn %v", subConns[1])
+	}
+	if scAddr == ring.items[0].hashKey {
+		subConns[0] = sc
+	} else if scAddr == ring.items[2].hashKey {
+		subConns[2] = sc
+	}
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
+
+	// All endpoints are in TransientFailure. Make the first endpoint in the
+	// ring report Ready. All picks should go to this endpoint which is two
+	// indexes away from the endpoint with the chosen hash.
+	subConns[0].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
+	subConns[0].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	subConns[0].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	if err := cc.WaitForConnectivityState(ctx, connectivity.Ready); err != nil {
+		t.Fatalf("Context timed out while waiting for channel to report Ready.")
+	}
+	p1 := <-cc.NewPickerCh
+	for i := 0; i < 5; i++ {
+		gotSCSt, _ := p1.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
+		if gotSCSt.SubConn != subConns[0] {
+			t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, subConns[0])
+		}
+	}
+
+	// Make the last endpoint in the ring report Ready. All picks should go to
+	// this endpoint since it is one index away from the chosen hash.
+	subConns[2].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
+	subConns[2].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	subConns[2].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	p2 := <-cc.NewPickerCh
+	for i := 0; i < 5; i++ {
+		gotSCSt, _ := p2.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
+		if gotSCSt.SubConn != subConns[2] {
+			t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, subConns[2])
+		}
+	}
+
+	// Make the second endpoint in the ring report Ready. All picks should go to
+	// this endpoint as it is the one with the chosen hash.
+	subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
+	subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	p3 := <-cc.NewPickerCh
+	for i := 0; i < 5; i++ {
+		gotSCSt, _ := p3.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
+		if gotSCSt.SubConn != subConns[1] {
+			t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, subConns[1])
+		}
+	}
+}
+
+// TestThreeBackendsAffinity covers that there are 3 SubConns, RPCs with the
+// same hash always pick the same SubConn. Then try different hash to pick
+// another backend, and verify the first hash still picks the first backend.
+func (s) TestThreeBackendsAffinityMultiple(t *testing.T) {
+	wantEndpoints := []resolver.Endpoint{
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[1]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[2]}}},
+	}
+	cc, _, p0 := setupTest(t, wantEndpoints)
+	// This test doesn't update addresses, so this ring will be used by all the
+	// pickers.
+	ring0 := p0.(*picker).ring
+
+	firstHash := ring0.items[0].hash
+	// firstHash+1 will pick the second SubConn from the ring.
+	testHash := firstHash + 1
+	// The first pick should be queued, and should trigger Connect() on the only
+	// SubConn.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)}); err != balancer.ErrNoSubConnAvailable {
+		t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
+	}
+	// The picked SubConn should be the second in the ring.
+	var sc0 *testutils.TestSubConn
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case sc0 = <-cc.NewSubConnCh:
+	}
+	if got, want := sc0.Addresses[0].Addr, ring0.items[1].hashKey; got != want {
+		t.Fatalf("SubConn.Address = %v, want = %v", got, want)
+	}
+	select {
+	case <-sc0.ConnectCh:
+	case <-time.After(defaultTestTimeout):
+		t.Errorf("timeout waiting for Connect() from SubConn %v", sc0)
+	}
+
+	// Send state updates to Ready.
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	if err := cc.WaitForConnectivityState(ctx, connectivity.Ready); err != nil {
+		t.Fatal(err)
+	}
+
+	// First hash should always pick sc0.
+	p1 := <-cc.NewPickerCh
+	for i := 0; i < 5; i++ {
+		gotSCSt, _ := p1.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
+		if gotSCSt.SubConn != sc0 {
+			t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc0)
+		}
+	}
+
+	secondHash := ring0.items[1].hash
+	// secondHash+1 will pick the third SubConn from the ring.
+	testHash2 := secondHash + 1
+	if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash2)}); err != balancer.ErrNoSubConnAvailable {
+		t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
+	}
+	var sc1 *testutils.TestSubConn
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case sc1 = <-cc.NewSubConnCh:
+	}
+	if got, want := sc1.Addresses[0].Addr, ring0.items[2].hashKey; got != want {
+		t.Fatalf("SubConn.Address = %v, want = %v", got, want)
+	}
+	select {
+	case <-sc1.ConnectCh:
+	case <-time.After(defaultTestTimeout):
+		t.Errorf("timeout waiting for Connect() from SubConn %v", sc1)
+	}
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+
+	// With the new generated picker, hash2 always picks sc1.
+	p2 := <-cc.NewPickerCh
+	for i := 0; i < 5; i++ {
+		gotSCSt, _ := p2.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash2)})
+		if gotSCSt.SubConn != sc1 {
+			t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc1)
+		}
+	}
+	// But the first hash still picks sc0.
+	for i := 0; i < 5; i++ {
+		gotSCSt, _ := p2.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
+		if gotSCSt.SubConn != sc0 {
+			t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc0)
+		}
+	}
+}
+
+// TestAddrWeightChange covers the following scenarios after setting up the
+// balancer with 3 addresses [A, B, C]:
+//   - updates balancer with [A, B, C], a new Picker should not be sent.
+//   - updates balancer with [A, B] (C removed), a new Picker is sent and the
+//     ring is updated.
+//   - updates balancer with [A, B], but B has a weight of 2, a new Picker is
+//     sent.  And the new ring should contain the correct number of entries
+//     and weights.
+func (s) TestAddrWeightChange(t *testing.T) {
+	endpoints := []resolver.Endpoint{
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[1]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[2]}}},
+	}
+	cc, b, p0 := setupTest(t, endpoints)
+	ring0 := p0.(*picker).ring
+
+	// Update with the same addresses, it will result in a new picker, but with
+	// the same ring.
+	if err := b.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState:  resolver.State{Endpoints: endpoints},
+		BalancerConfig: testConfig,
+	}); err != nil {
+		t.Fatalf("UpdateClientConnState returned err: %v", err)
+	}
+	var p1 balancer.Picker
+	select {
+	case p1 = <-cc.NewPickerCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("timeout waiting for picker after UpdateClientConn with same addresses")
+	}
+	ring1 := p1.(*picker).ring
+	if ring1 != ring0 {
+		t.Fatalf("new picker with same address has a different ring than before, want same")
+	}
+
+	// Delete an address, should send a new Picker.
+	if err := b.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState:  resolver.State{Endpoints: endpoints[:2]},
+		BalancerConfig: testConfig,
+	}); err != nil {
+		t.Fatalf("UpdateClientConnState returned err: %v", err)
+	}
+	var p2 balancer.Picker
+	select {
+	case p2 = <-cc.NewPickerCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("timeout waiting for picker after UpdateClientConn with different addresses")
+	}
+	ring2 := p2.(*picker).ring
+	if ring2 == ring0 {
+		t.Fatalf("new picker after removing address has the same ring as before, want different")
+	}
+
+	// Another update with the same addresses, but different weight.
+	if err := b.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState: resolver.State{Endpoints: []resolver.Endpoint{
+			endpoints[0],
+			weight.Set(endpoints[1], weight.EndpointInfo{Weight: 2}),
+		}},
+		BalancerConfig: testConfig,
+	}); err != nil {
+		t.Fatalf("UpdateClientConnState returned err: %v", err)
+	}
+	var p3 balancer.Picker
+	select {
+	case p3 = <-cc.NewPickerCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("timeout waiting for picker after UpdateClientConn with different addresses")
+	}
+	if p3.(*picker).ring == ring2 {
+		t.Fatalf("new picker after changing address weight has the same ring as before, want different")
+	}
+	// With the new update, the ring must look like this:
+	//   [
+	//     {idx:0 endpoint: {addr: testBackendAddrStrs[0], weight: 1}},
+	//     {idx:1 endpoint: {addr: testBackendAddrStrs[1], weight: 2}},
+	//     {idx:2 endpoint: {addr: testBackendAddrStrs[2], weight: 1}},
+	//   ].
+	if len(p3.(*picker).ring.items) != 3 {
+		t.Fatalf("new picker after changing address weight has %d entries, want 3", len(p3.(*picker).ring.items))
+	}
+	for _, i := range p3.(*picker).ring.items {
+		if i.hashKey == testBackendAddrStrs[0] {
+			if i.weight != 1 {
+				t.Fatalf("new picker after changing address weight has weight %d for %v, want 1", i.weight, i.hashKey)
+			}
+		}
+		if i.hashKey == testBackendAddrStrs[1] {
+			if i.weight != 2 {
+				t.Fatalf("new picker after changing address weight has weight %d for %v, want 2", i.weight, i.hashKey)
+			}
+		}
+	}
+}
+
+// TestAutoConnectEndpointOnTransientFailure covers the situation when an
+// endpoint fails. It verifies that a new endpoint is automatically tried
+// (without a pick) when there is no endpoint already in Connecting state.
+func (s) TestAutoConnectEndpointOnTransientFailure(t *testing.T) {
+	wantEndpoints := []resolver.Endpoint{
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[1]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[2]}}},
+		{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[3]}}},
+	}
+	cc, _, p0 := setupTest(t, wantEndpoints)
+
+	// ringhash won't tell SCs to connect until there is an RPC, so simulate
+	// one now.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	ctx = iringhash.SetXDSRequestHash(ctx, 0)
+	defer cancel()
+	p0.Pick(balancer.PickInfo{Ctx: ctx})
+
+	// The picked SubConn should be the second in the ring.
+	var sc0 *testutils.TestSubConn
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case sc0 = <-cc.NewSubConnCh:
+	}
+	select {
+	case <-sc0.ConnectCh:
+	case <-time.After(defaultTestTimeout):
+		t.Errorf("timeout waiting for Connect() from SubConn %v", sc0)
+	}
+
+	// Turn the first subconn to transient failure. This should set the overall
+	// connectivity state to CONNECTING.
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
+	cc.WaitForConnectivityState(ctx, connectivity.Connecting)
+
+	// It will trigger the second subconn to connect since there is only one
+	// endpoint, which is in TF.
+	var sc1 *testutils.TestSubConn
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case sc1 = <-cc.NewSubConnCh:
+	}
+	select {
+	case <-sc1.ConnectCh:
+	case <-time.After(defaultTestShortTimeout):
+		t.Fatalf("timeout waiting for Connect() from SubConn %v", sc1)
+	}
+
+	// Turn the second subconn to TF. This will set the overall state to TF.
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
+	cc.WaitForConnectivityState(ctx, connectivity.TransientFailure)
+
+	// It will trigger the third subconn to connect.
+	var sc2 *testutils.TestSubConn
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case sc2 = <-cc.NewSubConnCh:
+	}
+	select {
+	case <-sc2.ConnectCh:
+	case <-time.After(defaultTestShortTimeout):
+		t.Fatalf("timeout waiting for Connect() from SubConn %v", sc2)
+	}
+
+	sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+
+	// Send the first SubConn into CONNECTING. To do this, first make it READY,
+	// then CONNECTING.
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	cc.WaitForConnectivityState(ctx, connectivity.Ready)
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
+	// Since one endpoint is in TF and one in CONNECTING, the aggregated state
+	// will be CONNECTING.
+	cc.WaitForConnectivityState(ctx, connectivity.Connecting)
+	p1 := <-cc.NewPickerCh
+	p1.Pick(balancer.PickInfo{Ctx: ctx})
+	select {
+	case <-sc0.ConnectCh:
+	case <-time.After(defaultTestTimeout):
+		t.Errorf("timeout waiting for Connect() from SubConn %v", sc0)
+	}
+	sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+
+	// This will not trigger any new SubCOnns to be created, because sc0 is
+	// still attempting to connect, and we only need one SubConn to connect.
+	sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
+
+	select {
+	case sc := <-cc.NewSubConnCh:
+		t.Fatalf("unexpected SubConn creation: %v", sc)
+	case <-sc0.ConnectCh:
+		t.Fatalf("unexpected Connect() from SubConn %v", sc0)
+	case <-sc1.ConnectCh:
+		t.Fatalf("unexpected Connect() from SubConn %v", sc1)
+	case <-sc2.ConnectCh:
+		t.Fatalf("unexpected Connect() from SubConn %v", sc2)
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestAggregatedConnectivityState(t *testing.T) {
+	tests := []struct {
+		name           string
+		endpointStates []connectivity.State
+		want           connectivity.State
+	}{
+		{
+			name:           "one ready",
+			endpointStates: []connectivity.State{connectivity.Ready},
+			want:           connectivity.Ready,
+		},
+		{
+			name:           "one connecting",
+			endpointStates: []connectivity.State{connectivity.Connecting},
+			want:           connectivity.Connecting,
+		},
+		{
+			name:           "one ready one transient failure",
+			endpointStates: []connectivity.State{connectivity.Ready, connectivity.TransientFailure},
+			want:           connectivity.Ready,
+		},
+		{
+			name:           "one connecting one transient failure",
+			endpointStates: []connectivity.State{connectivity.Connecting, connectivity.TransientFailure},
+			want:           connectivity.Connecting,
+		},
+		{
+			name:           "one connecting two transient failure",
+			endpointStates: []connectivity.State{connectivity.Connecting, connectivity.TransientFailure, connectivity.TransientFailure},
+			want:           connectivity.TransientFailure,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			bal := &ringhashBalancer{endpointStates: resolver.NewEndpointMap[*endpointState]()}
+			for i, cs := range tt.endpointStates {
+				es := &endpointState{
+					state: balancer.State{ConnectivityState: cs},
+				}
+				ep := resolver.Endpoint{Addresses: []resolver.Address{{Addr: fmt.Sprintf("%d.%d.%d.%d:%d", i, i, i, i, i)}}}
+				bal.endpointStates.Set(ep, es)
+			}
+			if got := bal.aggregatedStateLocked(); got != tt.want {
+				t.Errorf("recordTransition() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+type testKeyType string
+
+const testKey testKeyType = "grpc.lb.ringhash.testKey"
+
+type testAttribute struct {
+	content string
+}
+
+func setTestAttrAddr(addr resolver.Address, content string) resolver.Address {
+	addr.BalancerAttributes = addr.BalancerAttributes.WithValue(testKey, testAttribute{content})
+	return addr
+}
+
+func setTestAttrEndpoint(endpoint resolver.Endpoint, content string) resolver.Endpoint {
+	endpoint.Attributes = endpoint.Attributes.WithValue(testKey, testAttribute{content})
+	return endpoint
+}
+
+// TestAddrBalancerAttributesChange tests the case where the ringhash balancer
+// receives a ClientConnUpdate with the same config and addresses as received in
+// the previous update. Although the `BalancerAttributes` and endpoint
+// attributes contents are the same, the pointers are different. This test
+// verifies that subConns are not recreated in this scenario.
+func (s) TestAddrBalancerAttributesChange(t *testing.T) {
+	content := "test"
+	addrs1 := []resolver.Address{setTestAttrAddr(resolver.Address{Addr: testBackendAddrStrs[0]}, content)}
+	wantEndpoints1 := []resolver.Endpoint{
+		setTestAttrEndpoint(resolver.Endpoint{Addresses: addrs1}, "content"),
+	}
+	cc, b, p0 := setupTest(t, wantEndpoints1)
+	ring0 := p0.(*picker).ring
+
+	firstHash := ring0.items[0].hash
+	// The first pick should be queued, and should trigger a connection to the
+	// only Endpoint which has a single address.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, firstHash)}); err != balancer.ErrNoSubConnAvailable {
+		t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
+	}
+	select {
+	case <-ctx.Done():
+		t.Fatalf("Timed out waiting for SubConn creation.")
+	case <-cc.NewSubConnCh:
+	}
+
+	addrs2 := []resolver.Address{setTestAttrAddr(resolver.Address{Addr: testBackendAddrStrs[0]}, content)}
+	wantEndpoints2 := []resolver.Endpoint{setTestAttrEndpoint(resolver.Endpoint{Addresses: addrs2}, content)}
+	if err := b.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState:  resolver.State{Endpoints: wantEndpoints2},
+		BalancerConfig: testConfig,
+	}); err != nil {
+		t.Fatalf("UpdateClientConnState returned err: %v", err)
+	}
+	select {
+	case <-cc.NewSubConnCh:
+		t.Fatal("new subConn created for an update with the same addresses")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}

balancer/rls/balancer.go

@@ -0,0 +1,733 @@
+/*
+ *
+ * Copyright 2020 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 rls implements the RLS LB policy.
+package rls
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"time"
+	"unsafe"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	estats "google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/backoff"
+	"google.golang.org/grpc/internal/balancergroup"
+	"google.golang.org/grpc/internal/buffer"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/pretty"
+	"google.golang.org/grpc/resolver"
+)
+
+const (
+	// Name is the name of the RLS LB policy.
+	//
+	// It currently has an experimental suffix which would be removed once
+	// end-to-end testing of the policy is completed.
+	Name = internal.RLSLoadBalancingPolicyName
+	// Default frequency for data cache purging.
+	periodicCachePurgeFreq = time.Minute
+)
+
+var (
+	logger            = grpclog.Component("rls")
+	errBalancerClosed = errors.New("rls LB policy is closed")
+
+	// Below defined vars for overriding in unit tests.
+
+	// Default exponential backoff strategy for data cache entries.
+	defaultBackoffStrategy = backoff.Strategy(backoff.DefaultExponential)
+	// Ticker used for periodic data cache purging.
+	dataCachePurgeTicker = func() *time.Ticker { return time.NewTicker(periodicCachePurgeFreq) }
+	// We want every cache entry to live in the cache for at least this
+	// duration. If we encounter a cache entry whose minimum expiration time is
+	// in the future, we abort the LRU pass, which may temporarily leave the
+	// cache being too large. This is necessary to ensure that in cases where
+	// the cache is too small, when we receive an RLS Response, we keep the
+	// resulting cache entry around long enough for the pending incoming
+	// requests to be re-processed through the new Picker. If we didn't do this,
+	// then we'd risk throwing away each RLS response as we receive it, in which
+	// case we would fail to actually route any of our incoming requests.
+	minEvictDuration = 5 * time.Second
+
+	// Following functions are no-ops in actual code, but can be overridden in
+	// tests to give tests visibility into exactly when certain events happen.
+	clientConnUpdateHook = func() {}
+	dataCachePurgeHook   = func() {}
+	resetBackoffHook     = func() {}
+
+	cacheEntriesMetric = estats.RegisterInt64AsyncGauge(estats.MetricDescriptor{
+		Name:        "grpc.lb.rls.cache_entries",
+		Description: "EXPERIMENTAL. Number of entries in the RLS cache.",
+		Unit:        "{entry}",
+		Labels:      []string{"grpc.target", "grpc.lb.rls.server_target", "grpc.lb.rls.instance_uuid"},
+		Default:     false,
+	})
+	cacheSizeMetric = estats.RegisterInt64AsyncGauge(estats.MetricDescriptor{
+		Name:        "grpc.lb.rls.cache_size",
+		Description: "EXPERIMENTAL. The current size of the RLS cache.",
+		Unit:        "By",
+		Labels:      []string{"grpc.target", "grpc.lb.rls.server_target", "grpc.lb.rls.instance_uuid"},
+		Default:     false,
+	})
+	defaultTargetPicksMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:           "grpc.lb.rls.default_target_picks",
+		Description:    "EXPERIMENTAL. Number of LB picks sent to the default target.",
+		Unit:           "{pick}",
+		Labels:         []string{"grpc.target", "grpc.lb.rls.server_target", "grpc.lb.rls.data_plane_target", "grpc.lb.pick_result"},
+		OptionalLabels: []string{"grpc.client.call.custom"},
+		Default:        false,
+	})
+	targetPicksMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:           "grpc.lb.rls.target_picks",
+		Description:    "EXPERIMENTAL. Number of LB picks sent to each RLS target. Note that if the default target is also returned by the RLS server, RPCs sent to that target from the cache will be counted in this metric, not in grpc.rls.default_target_picks.",
+		Unit:           "{pick}",
+		Labels:         []string{"grpc.target", "grpc.lb.rls.server_target", "grpc.lb.rls.data_plane_target", "grpc.lb.pick_result"},
+		OptionalLabels: []string{"grpc.client.call.custom"},
+		Default:        false,
+	})
+	failedPicksMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:           "grpc.lb.rls.failed_picks",
+		Description:    "EXPERIMENTAL. Number of LB picks failed due to either a failed RLS request or the RLS channel being throttled.",
+		Unit:           "{pick}",
+		Labels:         []string{"grpc.target", "grpc.lb.rls.server_target"},
+		OptionalLabels: []string{"grpc.client.call.custom"},
+		Default:        false,
+	})
+)
+
+func init() {
+	balancer.Register(&rlsBB{})
+}
+
+type rlsBB struct{}
+
+func (rlsBB) Name() string {
+	return Name
+}
+
+func (rlsBB) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	lb := &rlsBalancer{
+		closed:             grpcsync.NewEvent(),
+		done:               grpcsync.NewEvent(),
+		cc:                 cc,
+		bopts:              opts,
+		purgeTicker:        dataCachePurgeTicker(),
+		dataCachePurgeHook: dataCachePurgeHook,
+		lbCfg:              &lbConfig{},
+		pendingMap:         make(map[cacheKey]*backoffState),
+		childPolicies:      make(map[string]*childPolicyWrapper),
+		updateCh:           buffer.NewUnbounded[any](),
+	}
+	lb.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[rls-experimental-lb %p] ", lb))
+	lb.dataCache = newDataCache(maxCacheSize, lb.logger, opts.Target.String())
+	metricsRecorder := cc.MetricsRecorder()
+	lb.unregisterMetricHandler = metricsRecorder.RegisterAsyncReporter(lb, cacheEntriesMetric, cacheSizeMetric)
+	lb.bg = balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               opts,
+		StateAggregator:         lb,
+		Logger:                  lb.logger,
+		SubBalancerCloseTimeout: time.Duration(0), // Disable caching of removed child policies
+	})
+	go lb.run()
+	return lb
+}
+
+// rlsBalancer implements the RLS LB policy.
+type rlsBalancer struct {
+	closed             *grpcsync.Event // Fires when Close() is invoked. Guarded by stateMu.
+	done               *grpcsync.Event // Fires when Close() is done.
+	cc                 balancer.ClientConn
+	bopts              balancer.BuildOptions
+	purgeTicker        *time.Ticker
+	dataCachePurgeHook func()
+	logger             *internalgrpclog.PrefixLogger
+
+	// unregisterMetricHandler is the function to deregister the async metric reporter.
+	unregisterMetricHandler func()
+
+	// If both cacheMu and stateMu need to be acquired, the former must be
+	// acquired first to prevent a deadlock. This order restriction is due to the
+	// fact that in places where we need to acquire both the locks, we always
+	// start off reading the cache.
+
+	// cacheMu guards access to the data cache and pending requests map. We
+	// cannot use an RWMutex here since even an operation like
+	// dataCache.getEntry() modifies the underlying LRU, which is implemented as
+	// a doubly linked list.
+	cacheMu    sync.Mutex
+	dataCache  *dataCache                 // Cache of RLS data.
+	pendingMap map[cacheKey]*backoffState // Map of pending RLS requests.
+
+	// stateMu guards access to all LB policy state.
+	stateMu            sync.Mutex
+	lbCfg              *lbConfig        // Most recently received service config.
+	childPolicyBuilder balancer.Builder // Cached child policy builder.
+	resolverState      resolver.State   // Cached resolver state.
+	ctrlCh             *controlChannel  // Control channel to the RLS server.
+	bg                 *balancergroup.BalancerGroup
+	childPolicies      map[string]*childPolicyWrapper
+	defaultPolicy      *childPolicyWrapper
+	// A reference to the most recent picker sent to gRPC as part of a state
+	// update is cached in this field so that we can release the reference to the
+	// default child policy wrapper when a new picker is created. See
+	// sendNewPickerLocked() for details.
+	lastPicker *rlsPicker
+	// Set during UpdateClientConnState when pushing updates to child policies.
+	// Prevents state updates from child policies causing new pickers to be sent
+	// up the channel. Cleared after all child policies have processed the
+	// updates sent to them, after which a new picker is sent up the channel.
+	inhibitPickerUpdates bool
+
+	// Channel on which all updates are pushed. Processed in run().
+	updateCh *buffer.Unbounded[any]
+}
+
+type resumePickerUpdates struct {
+	done chan struct{}
+}
+
+// childPolicyIDAndState wraps a child policy id and its state update.
+type childPolicyIDAndState struct {
+	id    string
+	state balancer.State
+}
+
+type controlChannelReady struct{}
+
+// run is a long-running goroutine which handles all the updates that the
+// balancer wishes to handle. The appropriate updateHandler will push the update
+// on to a channel that this goroutine will select on, thereby the handling of
+// the update will happen asynchronously.
+func (b *rlsBalancer) run() {
+	// We exit out of the for loop below only after `Close()` has been invoked.
+	// Firing the done event here will ensure that Close() returns only after
+	// all goroutines are done.
+	defer func() { b.done.Fire() }()
+
+	// Wait for purgeDataCache() goroutine to exit before returning from here.
+	doneCh := make(chan struct{})
+	defer func() {
+		<-doneCh
+	}()
+	go b.purgeDataCache(doneCh)
+
+	for {
+		select {
+		case u, ok := <-b.updateCh.Get():
+			if !ok {
+				return
+			}
+			b.updateCh.Load()
+			switch update := u.(type) {
+			case childPolicyIDAndState:
+				b.handleChildPolicyStateUpdate(update.id, update.state)
+			case controlChannelReady:
+				b.logger.Infof("Resetting backoff state after control channel getting back to READY")
+				b.cacheMu.Lock()
+				updatePicker := b.dataCache.resetBackoffState(&backoffState{bs: defaultBackoffStrategy})
+				b.cacheMu.Unlock()
+				if updatePicker {
+					b.sendNewPicker()
+				}
+				resetBackoffHook()
+			case resumePickerUpdates:
+				b.stateMu.Lock()
+				b.logger.Infof("Resuming picker updates after config propagation to child policies")
+				b.inhibitPickerUpdates = false
+				b.sendNewPickerLocked()
+				close(update.done)
+				b.stateMu.Unlock()
+			default:
+				b.logger.Errorf("Unsupported update type %T", update)
+			}
+		case <-b.closed.Done():
+			return
+		}
+	}
+}
+
+// purgeDataCache is a long-running goroutine which periodically deletes expired
+// entries. An expired entry is one for which both the expiryTime and
+// backoffExpiryTime are in the past.
+func (b *rlsBalancer) purgeDataCache(doneCh chan struct{}) {
+	defer close(doneCh)
+
+	for {
+		select {
+		case <-b.closed.Done():
+			return
+		case <-b.purgeTicker.C:
+			b.cacheMu.Lock()
+			updatePicker := b.dataCache.evictExpiredEntries()
+			b.cacheMu.Unlock()
+			if updatePicker {
+				b.sendNewPicker()
+			}
+			b.dataCachePurgeHook()
+		}
+	}
+}
+
+func (b *rlsBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	defer clientConnUpdateHook()
+
+	b.stateMu.Lock()
+	if b.closed.HasFired() {
+		b.stateMu.Unlock()
+		b.logger.Warningf("Received service config after balancer close: %s", pretty.ToJSON(ccs.BalancerConfig))
+		return errBalancerClosed
+	}
+
+	newCfg := ccs.BalancerConfig.(*lbConfig)
+	if b.lbCfg.Equal(newCfg) {
+		b.stateMu.Unlock()
+		b.logger.Infof("New service config matches existing config")
+		return nil
+	}
+
+	b.logger.Infof("Delaying picker updates until config is propagated to and processed by child policies")
+	b.inhibitPickerUpdates = true
+
+	// When the RLS server name changes, the old control channel needs to be
+	// swapped out for a new one. All state associated with the throttling
+	// algorithm is stored on a per-control-channel basis; when we swap out
+	// channels, we also swap out the throttling state.
+	b.handleControlChannelUpdate(newCfg)
+
+	// Any changes to child policy name or configuration needs to be handled by
+	// either creating new child policies or pushing updates to existing ones.
+	b.resolverState = ccs.ResolverState
+	b.handleChildPolicyConfigUpdate(newCfg, &ccs)
+
+	// Resize the cache if the size in the config has changed.
+	resizeCache := newCfg.cacheSizeBytes != b.lbCfg.cacheSizeBytes
+
+	// Update the copy of the config in the LB policy before releasing the lock.
+	b.lbCfg = newCfg
+	b.stateMu.Unlock()
+
+	// We cannot do cache operations above because `cacheMu` needs to be grabbed
+	// before `stateMu` if we are to hold both locks at the same time.
+	b.cacheMu.Lock()
+	b.dataCache.updateRLSServerTarget(newCfg.lookupService)
+	if resizeCache {
+		// If the new config changes reduces the size of the data cache, we
+		// might have to evict entries to get the cache size down to the newly
+		// specified size. If we do evict an entry with valid backoff timer,
+		// the new picker needs to be sent to the channel to re-process any
+		// RPCs queued as a result of this backoff timer.
+		b.dataCache.resize(newCfg.cacheSizeBytes)
+	}
+	b.cacheMu.Unlock()
+	// Enqueue an event which will notify us when the above update has been
+	// propagated to all child policies, and the child policies have all
+	// processed their updates, and we have sent a picker update.
+	done := make(chan struct{})
+	b.updateCh.Put(resumePickerUpdates{done: done})
+	<-done
+	return nil
+}
+
+// handleControlChannelUpdate handles updates to service config fields which
+// influence the control channel to the RLS server.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) handleControlChannelUpdate(newCfg *lbConfig) {
+	if newCfg.lookupService == b.lbCfg.lookupService && newCfg.lookupServiceTimeout == b.lbCfg.lookupServiceTimeout {
+		return
+	}
+
+	// Create a new control channel and close the existing one.
+	b.logger.Infof("Creating control channel to RLS server at: %v", newCfg.lookupService)
+	backToReadyFn := func() {
+		b.updateCh.Put(controlChannelReady{})
+	}
+	ctrlCh, err := newControlChannel(newCfg.lookupService, newCfg.controlChannelServiceConfig, newCfg.lookupServiceTimeout, b.bopts, backToReadyFn)
+	if err != nil {
+		// This is very uncommon and usually represents a non-transient error.
+		// There is not much we can do here other than wait for another update
+		// which might fix things.
+		b.logger.Errorf("Failed to create control channel to %q: %v", newCfg.lookupService, err)
+		return
+	}
+	if b.ctrlCh != nil {
+		b.ctrlCh.close()
+	}
+	b.ctrlCh = ctrlCh
+}
+
+// handleChildPolicyConfigUpdate handles updates to service config fields which
+// influence child policy configuration.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) handleChildPolicyConfigUpdate(newCfg *lbConfig, ccs *balancer.ClientConnState) {
+	// Update child policy builder first since other steps are dependent on this.
+	if b.childPolicyBuilder == nil || b.childPolicyBuilder.Name() != newCfg.childPolicyName {
+		b.logger.Infof("Child policy changed to %q", newCfg.childPolicyName)
+		b.childPolicyBuilder = balancer.Get(newCfg.childPolicyName)
+		for _, cpw := range b.childPolicies {
+			// If the child policy has changed, we need to remove the old policy
+			// from the BalancerGroup and add a new one. The BalancerGroup takes
+			// care of closing the old one in this case.
+			b.bg.Remove(cpw.target)
+			b.bg.Add(cpw.target, b.childPolicyBuilder)
+		}
+	}
+
+	configSentToDefault := false
+	if b.lbCfg.defaultTarget != newCfg.defaultTarget {
+		// If the default target has changed, create a new childPolicyWrapper for
+		// the new target if required. If a new wrapper is created, add it to the
+		// childPolicies map and the BalancerGroup.
+		b.logger.Infof("Default target in LB config changing from %q to %q", b.lbCfg.defaultTarget, newCfg.defaultTarget)
+		cpw := b.childPolicies[newCfg.defaultTarget]
+		if cpw == nil {
+			cpw = newChildPolicyWrapper(newCfg.defaultTarget)
+			b.childPolicies[newCfg.defaultTarget] = cpw
+			b.bg.Add(newCfg.defaultTarget, b.childPolicyBuilder)
+			b.logger.Infof("Child policy %q added to BalancerGroup", newCfg.defaultTarget)
+		}
+		if err := b.buildAndPushChildPolicyConfigs(newCfg.defaultTarget, newCfg, ccs); err != nil {
+			cpw.lamify(err)
+		}
+
+		// If an old default exists, release its reference. If this was the last
+		// reference, remove the child policy from the BalancerGroup and remove the
+		// corresponding entry the childPolicies map.
+		if b.defaultPolicy != nil {
+			if b.defaultPolicy.releaseRef() {
+				delete(b.childPolicies, b.lbCfg.defaultTarget)
+				b.bg.Remove(b.defaultPolicy.target)
+			}
+		}
+		b.defaultPolicy = cpw
+		configSentToDefault = true
+	}
+
+	// No change in configuration affecting child policies. Return early.
+	if b.lbCfg.childPolicyName == newCfg.childPolicyName && b.lbCfg.childPolicyTargetField == newCfg.childPolicyTargetField && childPolicyConfigEqual(b.lbCfg.childPolicyConfig, newCfg.childPolicyConfig) {
+		return
+	}
+
+	// If fields affecting child policy configuration have changed, the changes
+	// are pushed to the childPolicyWrapper which handles them appropriately.
+	for _, cpw := range b.childPolicies {
+		if configSentToDefault && cpw.target == newCfg.defaultTarget {
+			// Default target has already been taken care of.
+			continue
+		}
+		if err := b.buildAndPushChildPolicyConfigs(cpw.target, newCfg, ccs); err != nil {
+			cpw.lamify(err)
+		}
+	}
+}
+
+// buildAndPushChildPolicyConfigs builds the final child policy configuration by
+// adding the `targetField` to the base child policy configuration received in
+// RLS LB policy configuration. The `targetField` is set to target and
+// configuration is pushed to the child policy through the BalancerGroup.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) buildAndPushChildPolicyConfigs(target string, newCfg *lbConfig, ccs *balancer.ClientConnState) error {
+	jsonTarget, err := json.Marshal(target)
+	if err != nil {
+		return fmt.Errorf("failed to marshal child policy target %q: %v", target, err)
+	}
+
+	config := newCfg.childPolicyConfig
+	targetField := newCfg.childPolicyTargetField
+	config[targetField] = jsonTarget
+	jsonCfg, err := json.Marshal(config)
+	if err != nil {
+		return fmt.Errorf("failed to marshal child policy config %+v: %v", config, err)
+	}
+
+	parser, _ := b.childPolicyBuilder.(balancer.ConfigParser)
+	parsedCfg, err := parser.ParseConfig(jsonCfg)
+	if err != nil {
+		return fmt.Errorf("childPolicy config parsing failed: %v", err)
+	}
+
+	state := balancer.ClientConnState{ResolverState: ccs.ResolverState, BalancerConfig: parsedCfg}
+	b.logger.Infof("Pushing new state to child policy %q: %+v", target, state)
+	if err := b.bg.UpdateClientConnState(target, state); err != nil {
+		b.logger.Warningf("UpdateClientConnState(%q, %+v) failed : %v", target, ccs, err)
+	}
+	return nil
+}
+
+func (b *rlsBalancer) ResolverError(err error) {
+	b.bg.ResolverError(err)
+}
+
+func (b *rlsBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+func (b *rlsBalancer) Close() {
+	b.stateMu.Lock()
+	b.closed.Fire()
+	b.purgeTicker.Stop()
+	if b.ctrlCh != nil {
+		b.ctrlCh.close()
+	}
+	b.unregisterMetricHandler()
+	b.bg.Close()
+	b.stateMu.Unlock()
+
+	b.cacheMu.Lock()
+	b.dataCache.stop()
+	b.cacheMu.Unlock()
+
+	b.updateCh.Close()
+
+	<-b.done.Done()
+}
+
+func (b *rlsBalancer) ExitIdle() {
+	b.bg.ExitIdle()
+}
+
+// sendNewPickerLocked pushes a new picker on to the channel.
+//
+// Note that regardless of what connectivity state is reported, the policy will
+// return its own picker, and not a picker that unconditionally queues
+// (typically used for IDLE or CONNECTING) or a picker that unconditionally
+// fails (typically used for TRANSIENT_FAILURE). This is required because,
+// irrespective of the connectivity state, we need to able to perform RLS
+// lookups for incoming RPCs and affect the status of queued RPCs based on the
+// receipt of RLS responses.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) sendNewPickerLocked() {
+	aggregatedState := b.aggregatedConnectivityState()
+
+	// Acquire a separate reference for the picker. This is required to ensure
+	// that the wrapper held by the old picker is not closed when the default
+	// target changes in the config, and a new wrapper is created for the new
+	// default target. See handleChildPolicyConfigUpdate() for how config changes
+	// affecting the default target are handled.
+	if b.defaultPolicy != nil {
+		b.defaultPolicy.acquireRef()
+	}
+
+	picker := &rlsPicker{
+		kbm:             b.lbCfg.kbMap,
+		origEndpoint:    b.bopts.Target.Endpoint(),
+		lb:              b,
+		defaultPolicy:   b.defaultPolicy,
+		ctrlCh:          b.ctrlCh,
+		maxAge:          b.lbCfg.maxAge,
+		staleAge:        b.lbCfg.staleAge,
+		bg:              b.bg,
+		rlsServerTarget: b.lbCfg.lookupService,
+		grpcTarget:      b.bopts.Target.String(),
+		metricsRecorder: b.cc.MetricsRecorder(),
+	}
+	picker.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[rls-picker %p] ", picker))
+	state := balancer.State{
+		ConnectivityState: aggregatedState,
+		Picker:            picker,
+	}
+
+	if !b.inhibitPickerUpdates {
+		b.logger.Infof("New balancer.State: %+v", state)
+		b.cc.UpdateState(state)
+	} else {
+		b.logger.Infof("Delaying picker update: %+v", state)
+	}
+
+	if b.lastPicker != nil {
+		if b.defaultPolicy != nil {
+			b.defaultPolicy.releaseRef()
+		}
+	}
+	b.lastPicker = picker
+}
+
+func (b *rlsBalancer) sendNewPicker() {
+	b.stateMu.Lock()
+	defer b.stateMu.Unlock()
+	if b.closed.HasFired() {
+		return
+	}
+	b.sendNewPickerLocked()
+}
+
+// The aggregated connectivity state reported is determined as follows:
+//   - If there is at least one child policy in state READY, the connectivity
+//     state is READY.
+//   - Otherwise, if there is at least one child policy in state CONNECTING, the
+//     connectivity state is CONNECTING.
+//   - Otherwise, if there is at least one child policy in state IDLE, the
+//     connectivity state is IDLE.
+//   - Otherwise, all child policies are in TRANSIENT_FAILURE, and the
+//     connectivity state is TRANSIENT_FAILURE.
+//
+// If the RLS policy has no child policies and no configured default target,
+// then we will report connectivity state IDLE.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) aggregatedConnectivityState() connectivity.State {
+	if len(b.childPolicies) == 0 && b.lbCfg.defaultTarget == "" {
+		return connectivity.Idle
+	}
+
+	var readyN, connectingN, idleN int
+	for _, cpw := range b.childPolicies {
+		state := (*balancer.State)(atomic.LoadPointer(&cpw.state))
+		switch state.ConnectivityState {
+		case connectivity.Ready:
+			readyN++
+		case connectivity.Connecting:
+			connectingN++
+		case connectivity.Idle:
+			idleN++
+		}
+	}
+
+	switch {
+	case readyN > 0:
+		return connectivity.Ready
+	case connectingN > 0:
+		return connectivity.Connecting
+	case idleN > 0:
+		return connectivity.Idle
+	default:
+		return connectivity.TransientFailure
+	}
+}
+
+// UpdateState is a implementation of the balancergroup.BalancerStateAggregator
+// interface. The actual state aggregation functionality is handled
+// asynchronously. This method only pushes the state update on to channel read
+// and dispatched by the run() goroutine.
+func (b *rlsBalancer) UpdateState(id string, state balancer.State) {
+	b.updateCh.Put(childPolicyIDAndState{id: id, state: state})
+}
+
+// handleChildPolicyStateUpdate provides the state aggregator functionality for
+// the BalancerGroup.
+//
+// This method is invoked by the BalancerGroup whenever a child policy sends a
+// state update. We cache the child policy's connectivity state and picker for
+// two reasons:
+//   - to suppress connectivity state transitions from TRANSIENT_FAILURE to states
+//     other than READY
+//   - to delegate picks to child policies
+func (b *rlsBalancer) handleChildPolicyStateUpdate(id string, newState balancer.State) {
+	b.stateMu.Lock()
+	defer b.stateMu.Unlock()
+
+	cpw := b.childPolicies[id]
+	if cpw == nil {
+		// All child policies start with an entry in the map. If ID is not in
+		// map, it's either been removed, or never existed.
+		b.logger.Warningf("Received state update %+v for missing child policy %q", newState, id)
+		return
+	}
+
+	oldState := (*balancer.State)(atomic.LoadPointer(&cpw.state))
+	if oldState.ConnectivityState == connectivity.TransientFailure && newState.ConnectivityState == connectivity.Connecting {
+		// Ignore state transitions from TRANSIENT_FAILURE to CONNECTING, and thus
+		// fail pending RPCs instead of queuing them indefinitely when all
+		// subChannels are failing, even if the subChannels are bouncing back and
+		// forth between CONNECTING and TRANSIENT_FAILURE.
+		return
+	}
+	atomic.StorePointer(&cpw.state, unsafe.Pointer(&newState))
+	b.logger.Infof("Child policy %q has new state %+v", id, newState)
+	b.sendNewPickerLocked()
+}
+
+// acquireChildPolicyReferences attempts to acquire references to
+// childPolicyWrappers corresponding to the passed in targets. If there is no
+// childPolicyWrapper corresponding to one of the targets, a new one is created
+// and added to the BalancerGroup.
+func (b *rlsBalancer) acquireChildPolicyReferences(targets []string) []*childPolicyWrapper {
+	b.stateMu.Lock()
+	var newChildPolicies []*childPolicyWrapper
+	for _, target := range targets {
+		// If the target exists in the LB policy's childPolicies map. a new
+		// reference is taken here and added to the new list.
+		if cpw := b.childPolicies[target]; cpw != nil {
+			cpw.acquireRef()
+			newChildPolicies = append(newChildPolicies, cpw)
+			continue
+		}
+
+		// If the target does not exist in the child policy map, then a new
+		// child policy wrapper is created and added to the new list.
+		cpw := newChildPolicyWrapper(target)
+		b.childPolicies[target] = cpw
+		b.bg.Add(target, b.childPolicyBuilder)
+		b.logger.Infof("Child policy %q added to BalancerGroup", target)
+		newChildPolicies = append(newChildPolicies, cpw)
+		if err := b.buildAndPushChildPolicyConfigs(target, b.lbCfg, &balancer.ClientConnState{
+			ResolverState: b.resolverState,
+		}); err != nil {
+			cpw.lamify(err)
+		}
+	}
+	b.stateMu.Unlock()
+	return newChildPolicies
+}
+
+// releaseChildPolicyReferences releases references to childPolicyWrappers
+// corresponding to the passed in targets. If the release reference was the last
+// one, the child policy is removed from the BalancerGroup.
+func (b *rlsBalancer) releaseChildPolicyReferences(targets []string) {
+	b.stateMu.Lock()
+	for _, target := range targets {
+		if cpw := b.childPolicies[target]; cpw.releaseRef() {
+			delete(b.childPolicies, cpw.target)
+			b.bg.Remove(cpw.target)
+		}
+	}
+	b.stateMu.Unlock()
+}
+
+// Report reports the metrics data to the provided recorder.
+func (b *rlsBalancer) Report(r estats.AsyncMetricsRecorder) error {
+	b.cacheMu.Lock()
+	currentSize := b.dataCache.currentSize
+	entriesLen := int64(len(b.dataCache.entries))
+	rlsServerTarget := b.dataCache.rlsServerTarget
+	grpcTarget := b.dataCache.grpcTarget
+	uuid := b.dataCache.uuid
+	shutdown := b.dataCache.shutdown.HasFired()
+	b.cacheMu.Unlock()
+
+	if shutdown {
+		return nil
+	}
+
+	cacheSizeMetric.Record(r, currentSize, grpcTarget, rlsServerTarget, uuid)
+	cacheEntriesMetric.Record(r, entriesLen, grpcTarget, rlsServerTarget, uuid)
+	return nil
+}

balancer/rls/cache.go

@@ -0,0 +1,378 @@
+/*
+ *
+ * Copyright 2021 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 rls
+
+import (
+	"container/list"
+	"time"
+
+	"github.com/google/uuid"
+	"google.golang.org/grpc/internal/backoff"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/grpcsync"
+)
+
+// cacheKey represents the key used to uniquely identify an entry in the data
+// cache and in the pending requests map.
+type cacheKey struct {
+	// path is the full path of the incoming RPC request.
+	path string
+	// keys is a stringified version of the RLS request key map built using the
+	// RLS keyBuilder. Since maps are not a type which is comparable in Go, it
+	// cannot be part of the key for another map (entries in the data cache and
+	// pending requests map are stored in maps).
+	keys string
+}
+
+// cacheEntry wraps all the data to be stored in a data cache entry.
+type cacheEntry struct {
+	// childPolicyWrappers contains the list of child policy wrappers
+	// corresponding to the targets returned by the RLS server for this entry.
+	childPolicyWrappers []*childPolicyWrapper
+	// headerData is received in the RLS response and is to be sent in the
+	// X-Google-RLS-Data header for matching RPCs.
+	headerData string
+	// expiryTime is the absolute time at which this cache entry stops
+	// being valid. When an RLS request succeeds, this is set to the current
+	// time plus the max_age field from the LB policy config.
+	expiryTime time.Time
+	// staleTime is the absolute time after which this cache entry will be
+	// proactively refreshed if an incoming RPC matches this entry. When an RLS
+	// request succeeds, this is set to the current time plus the stale_age from
+	// the LB policy config.
+	staleTime time.Time
+	// earliestEvictTime is the absolute time before which this entry should not
+	// be evicted from the cache. When a cache entry is created, this is set to
+	// the current time plus a default value of 5 seconds. This is required to
+	// make sure that a new entry added to the cache is not evicted before the
+	// RLS response arrives (usually when the cache is too small).
+	earliestEvictTime time.Time
+
+	// status stores the RPC status of the previous RLS request for this
+	// entry. Picks for entries with a non-nil value for this field are failed
+	// with the error stored here.
+	status error
+	// backoffState contains all backoff related state. When an RLS request
+	// succeeds, backoffState is reset. This state moves between the data cache
+	// and the pending requests map.
+	backoffState *backoffState
+	// backoffTime is the absolute time at which the backoff period for this
+	// entry ends. When an RLS request fails, this is set to the current time
+	// plus the backoff value returned by the backoffState. The backoff timer is
+	// also setup with this value. No new RLS requests are sent out for this
+	// entry until the backoff period ends.
+	//
+	// Set to zero time instant upon a successful RLS response.
+	backoffTime time.Time
+	// backoffExpiryTime is the absolute time at which an entry which has gone
+	// through backoff stops being valid.  When an RLS request fails, this is
+	// set to the current time plus twice the backoff time. The cache expiry
+	// timer will only delete entries for which both expiryTime and
+	// backoffExpiryTime are in the past.
+	//
+	// Set to zero time instant upon a successful RLS response.
+	backoffExpiryTime time.Time
+
+	// size stores the size of this cache entry. Used to enforce the cache size
+	// specified in the LB policy configuration.
+	size int64
+}
+
+// backoffState wraps all backoff related state associated with a cache entry.
+type backoffState struct {
+	// retries keeps track of the number of RLS failures, to be able to
+	// determine the amount of time to backoff before the next attempt.
+	retries int
+	// bs is the exponential backoff implementation which returns the amount of
+	// time to backoff, given the number of retries.
+	bs backoff.Strategy
+	// timer fires when the backoff period ends and incoming requests after this
+	// will trigger a new RLS request.
+	timer *time.Timer
+}
+
+// lru is a cache implementation with a least recently used eviction policy.
+// Internally it uses a doubly linked list, with the least recently used element
+// at the front of the list and the most recently used element at the back of
+// the list. The value stored in this cache will be of type `cacheKey`.
+//
+// It is not safe for concurrent access.
+type lru struct {
+	ll *list.List
+
+	// A map from the value stored in the lru to its underlying list element is
+	// maintained to have a clean API. Without this, a subset of the lru's API
+	// would accept/return cacheKey while another subset would accept/return
+	// list elements.
+	m map[cacheKey]*list.Element
+}
+
+// newLRU creates a new cache with a least recently used eviction policy.
+func newLRU() *lru {
+	return &lru{
+		ll: list.New(),
+		m:  make(map[cacheKey]*list.Element),
+	}
+}
+
+func (l *lru) addEntry(key cacheKey) {
+	e := l.ll.PushBack(key)
+	l.m[key] = e
+}
+
+func (l *lru) makeRecent(key cacheKey) {
+	e := l.m[key]
+	l.ll.MoveToBack(e)
+}
+
+func (l *lru) removeEntry(key cacheKey) {
+	e := l.m[key]
+	l.ll.Remove(e)
+	delete(l.m, key)
+}
+
+func (l *lru) getLeastRecentlyUsed() cacheKey {
+	e := l.ll.Front()
+	if e == nil {
+		return cacheKey{}
+	}
+	return e.Value.(cacheKey)
+}
+
+// dataCache contains a cache of RLS data used by the LB policy to make routing
+// decisions.
+//
+// The dataCache will be keyed by the request's path and keys, represented by
+// the `cacheKey` type. It will maintain the cache keys in an `lru` and the
+// cache data, represented by the `cacheEntry` type, in a native map.
+//
+// It is not safe for concurrent access.
+type dataCache struct {
+	maxSize         int64 // Maximum allowed size.
+	currentSize     int64 // Current size.
+	keys            *lru  // Cache keys maintained in lru order.
+	entries         map[cacheKey]*cacheEntry
+	logger          *internalgrpclog.PrefixLogger
+	shutdown        *grpcsync.Event
+	rlsServerTarget string
+
+	// Read only after initialization.
+	grpcTarget string
+	uuid       string
+}
+
+func newDataCache(size int64, logger *internalgrpclog.PrefixLogger, grpcTarget string) *dataCache {
+	return &dataCache{
+		maxSize:    size,
+		keys:       newLRU(),
+		entries:    make(map[cacheKey]*cacheEntry),
+		logger:     logger,
+		shutdown:   grpcsync.NewEvent(),
+		grpcTarget: grpcTarget,
+		uuid:       uuid.New().String(),
+	}
+}
+
+// updateRLSServerTarget updates the RLS Server Target the RLS Balancer is
+// configured with.
+func (dc *dataCache) updateRLSServerTarget(rlsServerTarget string) {
+	dc.rlsServerTarget = rlsServerTarget
+}
+
+// resize changes the maximum allowed size of the data cache.
+//
+// The return value indicates if an entry with a valid backoff timer was
+// evicted. This is important to the RLS LB policy which would send a new picker
+// on the channel to re-process any RPCs queued as a result of this backoff
+// timer.
+func (dc *dataCache) resize(size int64) (backoffCancelled bool) {
+	if dc.shutdown.HasFired() {
+		return false
+	}
+
+	backoffCancelled = false
+	for dc.currentSize > size {
+		key := dc.keys.getLeastRecentlyUsed()
+		entry, ok := dc.entries[key]
+		if !ok {
+			// This should never happen.
+			dc.logger.Errorf("cacheKey %+v not found in the cache while attempting to resize it", key)
+			break
+		}
+
+		// When we encounter a cache entry whose minimum expiration time is in
+		// the future, we abort the LRU pass, which may temporarily leave the
+		// cache being too large. This is necessary to ensure that in cases
+		// where the cache is too small, when we receive an RLS Response, we
+		// keep the resulting cache entry around long enough for the pending
+		// incoming requests to be re-processed through the new Picker. If we
+		// didn't do this, then we'd risk throwing away each RLS response as we
+		// receive it, in which case we would fail to actually route any of our
+		// incoming requests.
+		if entry.earliestEvictTime.After(time.Now()) {
+			dc.logger.Warningf("cachekey %+v is too recent to be evicted. Stopping cache resizing for now", key)
+			break
+		}
+
+		// Stop the backoff timer before evicting the entry.
+		if entry.backoffState != nil && entry.backoffState.timer != nil {
+			if entry.backoffState.timer.Stop() {
+				entry.backoffState.timer = nil
+				backoffCancelled = true
+			}
+		}
+		dc.deleteAndCleanup(key, entry)
+	}
+	dc.maxSize = size
+	return backoffCancelled
+}
+
+// evictExpiredEntries sweeps through the cache and deletes expired entries. An
+// expired entry is one for which both the `expiryTime` and `backoffExpiryTime`
+// fields are in the past.
+//
+// The return value indicates if any expired entries were evicted.
+//
+// The LB policy invokes this method periodically to purge expired entries.
+func (dc *dataCache) evictExpiredEntries() bool {
+	if dc.shutdown.HasFired() {
+		return false
+	}
+
+	evicted := false
+	for key, entry := range dc.entries {
+		// Only evict entries for which both the data expiration time and
+		// backoff expiration time fields are in the past.
+		now := time.Now()
+		if entry.expiryTime.After(now) || entry.backoffExpiryTime.After(now) {
+			continue
+		}
+		dc.deleteAndCleanup(key, entry)
+		evicted = true
+	}
+	return evicted
+}
+
+// resetBackoffState sweeps through the cache and for entries with a backoff
+// state, the backoff timer is cancelled and the backoff state is reset. The
+// return value indicates if any entries were mutated in this fashion.
+//
+// The LB policy invokes this method when the control channel moves from READY
+// to TRANSIENT_FAILURE back to READY. See `monitorConnectivityState` method on
+// the `controlChannel` type for more details.
+func (dc *dataCache) resetBackoffState(newBackoffState *backoffState) bool {
+	if dc.shutdown.HasFired() {
+		return false
+	}
+
+	backoffReset := false
+	for _, entry := range dc.entries {
+		if entry.backoffState == nil {
+			continue
+		}
+		if entry.backoffState.timer != nil {
+			entry.backoffState.timer.Stop()
+			entry.backoffState.timer = nil
+		}
+		entry.backoffState = &backoffState{bs: newBackoffState.bs}
+		entry.backoffTime = time.Time{}
+		entry.backoffExpiryTime = time.Time{}
+		backoffReset = true
+	}
+	return backoffReset
+}
+
+// addEntry adds a cache entry for the given key.
+//
+// Return value backoffCancelled indicates if a cache entry with a valid backoff
+// timer was evicted to make space for the current entry. This is important to
+// the RLS LB policy which would send a new picker on the channel to re-process
+// any RPCs queued as a result of this backoff timer.
+//
+// Return value ok indicates if entry was successfully added to the cache.
+func (dc *dataCache) addEntry(key cacheKey, entry *cacheEntry) (backoffCancelled bool, ok bool) {
+	if dc.shutdown.HasFired() {
+		return false, false
+	}
+
+	// Handle the extremely unlikely case that a single entry is bigger than the
+	// size of the cache.
+	if entry.size > dc.maxSize {
+		return false, false
+	}
+	dc.entries[key] = entry
+	dc.currentSize += entry.size
+	dc.keys.addEntry(key)
+	// If the new entry makes the cache go over its configured size, remove some
+	// old entries.
+	if dc.currentSize > dc.maxSize {
+		backoffCancelled = dc.resize(dc.maxSize)
+	}
+
+	return backoffCancelled, true
+}
+
+// updateEntrySize updates the size of a cache entry and the current size of the
+// data cache. An entry's size can change upon receipt of an RLS response.
+func (dc *dataCache) updateEntrySize(entry *cacheEntry, newSize int64) {
+	dc.currentSize -= entry.size
+	entry.size = newSize
+	dc.currentSize += entry.size
+
+}
+
+func (dc *dataCache) getEntry(key cacheKey) *cacheEntry {
+	if dc.shutdown.HasFired() {
+		return nil
+	}
+
+	entry, ok := dc.entries[key]
+	if !ok {
+		return nil
+	}
+	dc.keys.makeRecent(key)
+	return entry
+}
+
+func (dc *dataCache) removeEntryForTesting(key cacheKey) {
+	entry, ok := dc.entries[key]
+	if !ok {
+		return
+	}
+	dc.deleteAndCleanup(key, entry)
+}
+
+// deleteAndCleanup performs actions required at the time of deleting an entry
+// from the data cache.
+// - the entry is removed from the map of entries
+// - current size of the data cache is update
+// - the key is removed from the LRU
+func (dc *dataCache) deleteAndCleanup(key cacheKey, entry *cacheEntry) {
+	delete(dc.entries, key)
+	dc.currentSize -= entry.size
+	dc.keys.removeEntry(key)
+
+}
+
+func (dc *dataCache) stop() {
+	for key, entry := range dc.entries {
+		dc.deleteAndCleanup(key, entry)
+	}
+	dc.shutdown.Fire()
+}

balancer/rls/cache_test.go

@@ -0,0 +1,243 @@
+/*
+ *
+ * Copyright 2021 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 rls
+
+import (
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"google.golang.org/grpc/internal/backoff"
+)
+
+var (
+	cacheKeys = []cacheKey{
+		{path: "0", keys: "a"},
+		{path: "1", keys: "b"},
+		{path: "2", keys: "c"},
+		{path: "3", keys: "d"},
+		{path: "4", keys: "e"},
+	}
+
+	longDuration  = 10 * time.Minute
+	shortDuration = 1 * time.Millisecond
+	cacheEntries  []*cacheEntry
+)
+
+func initCacheEntries() {
+	// All entries have a dummy size of 1 to simplify resize operations.
+	cacheEntries = []*cacheEntry{
+		{
+			// Entry is valid and minimum expiry time has not expired.
+			expiryTime:        time.Now().Add(longDuration),
+			earliestEvictTime: time.Now().Add(longDuration),
+			size:              1,
+		},
+		{
+			// Entry is valid and is in backoff.
+			expiryTime:   time.Now().Add(longDuration),
+			backoffTime:  time.Now().Add(longDuration),
+			backoffState: &backoffState{timer: time.NewTimer(longDuration)},
+			size:         1,
+		},
+		{
+			// Entry is valid, and not in backoff.
+			expiryTime: time.Now().Add(longDuration),
+			size:       1,
+		},
+		{
+			// Entry is invalid.
+			expiryTime: time.Time{}.Add(shortDuration),
+			size:       1,
+		},
+		{
+			// Entry is invalid valid and backoff has expired.
+			expiryTime:        time.Time{}.Add(shortDuration),
+			backoffExpiryTime: time.Time{}.Add(shortDuration),
+			size:              1,
+		},
+	}
+}
+
+func (s) TestLRU_BasicOperations(t *testing.T) {
+	initCacheEntries()
+	// Create an LRU and add some entries to it.
+	lru := newLRU()
+	for _, k := range cacheKeys {
+		lru.addEntry(k)
+	}
+
+	// Get the least recent entry. This should be the first entry we added.
+	if got, want := lru.getLeastRecentlyUsed(), cacheKeys[0]; got != want {
+		t.Fatalf("lru.getLeastRecentlyUsed() = %v, want %v", got, want)
+	}
+
+	// Iterate through the slice of keys we added earlier, making them the most
+	// recent entry, one at a time. The least recent entry at that point should
+	// be the next entry from our slice of keys.
+	for i, k := range cacheKeys {
+		lru.makeRecent(k)
+
+		lruIndex := (i + 1) % len(cacheKeys)
+		if got, want := lru.getLeastRecentlyUsed(), cacheKeys[lruIndex]; got != want {
+			t.Fatalf("lru.getLeastRecentlyUsed() = %v, want %v", got, want)
+		}
+	}
+
+	// Iterate through the slice of keys we added earlier, removing them one at
+	// a time The least recent entry at that point should be the next entry from
+	// our slice of keys, except for the last one because the lru will be empty.
+	for i, k := range cacheKeys {
+		lru.removeEntry(k)
+
+		var want cacheKey
+		if i < len(cacheKeys)-1 {
+			want = cacheKeys[i+1]
+		}
+		if got := lru.getLeastRecentlyUsed(); got != want {
+			t.Fatalf("lru.getLeastRecentlyUsed() = %v, want %v", got, want)
+		}
+	}
+}
+
+func (s) TestDataCache_BasicOperations(t *testing.T) {
+	initCacheEntries()
+	dc := newDataCache(5, nil, "")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntries[i])
+	}
+	for i, k := range cacheKeys {
+		entry := dc.getEntry(k)
+		if !cmp.Equal(entry, cacheEntries[i], cmp.AllowUnexported(cacheEntry{}, backoffState{}), cmpopts.IgnoreUnexported(time.Timer{})) {
+			t.Fatalf("Data cache lookup for key %v returned entry %v, want %v", k, entry, cacheEntries[i])
+		}
+	}
+}
+
+func (s) TestDataCache_AddForcesResize(t *testing.T) {
+	initCacheEntries()
+	dc := newDataCache(1, nil, "")
+
+	// The first entry in cacheEntries has a minimum expiry time in the future.
+	// This entry would stop the resize operation since we do not evict entries
+	// whose minimum expiration time is in the future. So, we do not use that
+	// entry in this test. The entry being added has a running backoff timer.
+	evicted, ok := dc.addEntry(cacheKeys[1], cacheEntries[1])
+	if evicted || !ok {
+		t.Fatalf("dataCache.addEntry() returned (%v, %v) want (false, true)", evicted, ok)
+	}
+
+	// Add another entry leading to the eviction of the above entry which has a
+	// running backoff timer. The first return value is expected to be true.
+	backoffCancelled, ok := dc.addEntry(cacheKeys[2], cacheEntries[2])
+	if !backoffCancelled || !ok {
+		t.Fatalf("dataCache.addEntry() returned (%v, %v) want (true, true)", backoffCancelled, ok)
+	}
+
+	// Add another entry leading to the eviction of the above entry which does not
+	// have a running backoff timer. This should evict the above entry, but the
+	// first return value is expected to be false.
+	backoffCancelled, ok = dc.addEntry(cacheKeys[3], cacheEntries[3])
+	if backoffCancelled || !ok {
+		t.Fatalf("dataCache.addEntry() returned (%v, %v) want (false, true)", backoffCancelled, ok)
+	}
+}
+
+func (s) TestDataCache_Resize(t *testing.T) {
+	initCacheEntries()
+	dc := newDataCache(5, nil, "")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntries[i])
+	}
+
+	// The first cache entry (with a key of cacheKeys[0]) that we added has an
+	// earliestEvictTime in the future. As part of the resize operation, we
+	// traverse the cache in least recently used order, and this will be first
+	// entry that we will encounter. And since the earliestEvictTime is in the
+	// future, the resize operation will stop, leaving the cache bigger than
+	// what was asked for.
+	if dc.resize(1) {
+		t.Fatalf("dataCache.resize() returned true, want false")
+	}
+	if dc.currentSize != 5 {
+		t.Fatalf("dataCache.size is %d, want 5", dc.currentSize)
+	}
+
+	// Remove the entry with earliestEvictTime in the future and retry the
+	// resize operation.
+	dc.removeEntryForTesting(cacheKeys[0])
+	if !dc.resize(1) {
+		t.Fatalf("dataCache.resize() returned false, want true")
+	}
+	if dc.currentSize != 1 {
+		t.Fatalf("dataCache.size is %d, want 1", dc.currentSize)
+	}
+}
+
+func (s) TestDataCache_EvictExpiredEntries(t *testing.T) {
+	initCacheEntries()
+	dc := newDataCache(5, nil, "")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntries[i])
+	}
+
+	// The last two entries in the cacheEntries list have expired, and will be
+	// evicted. The first three should still remain in the cache.
+	if !dc.evictExpiredEntries() {
+		t.Fatal("dataCache.evictExpiredEntries() returned false, want true")
+	}
+	if dc.currentSize != 3 {
+		t.Fatalf("dataCache.size is %d, want 3", dc.currentSize)
+	}
+	for i := 0; i < 3; i++ {
+		entry := dc.getEntry(cacheKeys[i])
+		if !cmp.Equal(entry, cacheEntries[i], cmp.AllowUnexported(cacheEntry{}, backoffState{}), cmpopts.IgnoreUnexported(time.Timer{})) {
+			t.Fatalf("Data cache lookup for key %v returned entry %v, want %v", cacheKeys[i], entry, cacheEntries[i])
+		}
+	}
+}
+
+func (s) TestDataCache_ResetBackoffState(t *testing.T) {
+	type fakeBackoff struct {
+		backoff.Strategy
+	}
+
+	initCacheEntries()
+	dc := newDataCache(5, nil, "")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntries[i])
+	}
+
+	newBackoffState := &backoffState{bs: &fakeBackoff{}}
+	if updatePicker := dc.resetBackoffState(newBackoffState); !updatePicker {
+		t.Fatal("dataCache.resetBackoffState() returned updatePicker is false, want true")
+	}
+
+	// Make sure that the entry with no backoff state was not touched.
+	if entry := dc.getEntry(cacheKeys[0]); cmp.Equal(entry.backoffState, newBackoffState, cmp.AllowUnexported(backoffState{})) {
+		t.Fatal("dataCache.resetBackoffState() touched entries without a valid backoffState")
+	}
+
+	// Make sure that the entry with a valid backoff state was reset.
+	entry := dc.getEntry(cacheKeys[1])
+	if diff := cmp.Diff(entry.backoffState, newBackoffState, cmp.AllowUnexported(backoffState{})); diff != "" {
+		t.Fatalf("unexpected diff in backoffState for cache entry after dataCache.resetBackoffState(): %s", diff)
+	}
+}

balancer/rls/child_policy.go

@@ -0,0 +1,109 @@
+/*
+ *
+ * Copyright 2021 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 rls
+
+import (
+	"fmt"
+	"sync/atomic"
+	"unsafe"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/connectivity"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+)
+
+// childPolicyWrapper is a reference counted wrapper around a child policy.
+//
+// The LB policy maintains a map of these wrappers keyed by the target returned
+// by RLS. When a target is seen for the first time, a child policy wrapper is
+// created for it and the wrapper is added to the child policy map. Each entry
+// in the data cache holds references to the corresponding child policy
+// wrappers. The LB policy also holds a reference to the child policy wrapper
+// for the default target specified in the LB Policy Configuration
+//
+// When a cache entry is evicted, it releases references to the child policy
+// wrappers that it contains. When all references have been released, the
+// wrapper is removed from the child policy map and is destroyed.
+//
+// The child policy wrapper also caches the connectivity state and most recent
+// picker from the child policy. Once the child policy wrapper reports
+// TRANSIENT_FAILURE, it will continue reporting that state until it goes READY;
+// transitions from TRANSIENT_FAILURE to CONNECTING are ignored.
+//
+// Whenever a child policy wrapper changes its connectivity state, the LB policy
+// returns a new picker to the channel, since the channel may need to re-process
+// the picks for queued RPCs.
+//
+// It is not safe for concurrent access.
+type childPolicyWrapper struct {
+	logger *internalgrpclog.PrefixLogger
+	target string // RLS target corresponding to this child policy.
+	refCnt int    // Reference count.
+
+	// Balancer state reported by the child policy. The RLS LB policy maintains
+	// these child policies in a BalancerGroup. The state reported by the child
+	// policy is pushed to the state aggregator (which is also implemented by the
+	// RLS LB policy) and cached here. See handleChildPolicyStateUpdate() for
+	// details on how the state aggregation is performed.
+	//
+	// While this field is written to by the LB policy, it is read by the picker
+	// at Pick time. Making this an atomic to enable the picker to read this value
+	// without a mutex.
+	state unsafe.Pointer // *balancer.State
+}
+
+// newChildPolicyWrapper creates a child policy wrapper for the given target,
+// and is initialized with one reference and starts off in CONNECTING state.
+func newChildPolicyWrapper(target string) *childPolicyWrapper {
+	c := &childPolicyWrapper{
+		target: target,
+		refCnt: 1,
+		state: unsafe.Pointer(&balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            base.NewErrPicker(balancer.ErrNoSubConnAvailable),
+		}),
+	}
+	c.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[rls-child-policy-wrapper %s %p] ", c.target, c))
+	c.logger.Infof("Created")
+	return c
+}
+
+// acquireRef increments the reference count on the child policy wrapper.
+func (c *childPolicyWrapper) acquireRef() {
+	c.refCnt++
+}
+
+// releaseRef decrements the reference count on the child policy wrapper. The
+// return value indicates whether the released reference was the last one.
+func (c *childPolicyWrapper) releaseRef() bool {
+	c.refCnt--
+	return c.refCnt == 0
+}
+
+// lamify causes the child policy wrapper to return a picker which will always
+// fail requests. This is used when the wrapper runs into errors when trying to
+// build and parse the child policy configuration.
+func (c *childPolicyWrapper) lamify(err error) {
+	c.logger.Warningf("Entering lame mode: %v", err)
+	atomic.StorePointer(&c.state, unsafe.Pointer(&balancer.State{
+		ConnectivityState: connectivity.TransientFailure,
+		Picker:            base.NewErrPicker(err),
+	}))
+}

balancer/rls/config.go

@@ -0,0 +1,330 @@
+/*
+ *
+ * Copyright 2020 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 rls
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/url"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/rls/internal/keys"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/pretty"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/types/known/durationpb"
+)
+
+const (
+	// Default max_age if not specified (or greater than this value) in the
+	// service config.
+	maxMaxAge = 5 * time.Minute
+	// Upper limit for cache_size since we don't fully trust the service config.
+	maxCacheSize = 5 * 1024 * 1024 * 8 // 5MB in bytes
+	// Default lookup_service_timeout if not specified in the service config.
+	defaultLookupServiceTimeout = 10 * time.Second
+	// Default value for targetNameField in the child policy config during
+	// service config validation.
+	dummyChildPolicyTarget = "target_name_to_be_filled_in_later"
+)
+
+// lbConfig is the internal representation of the RLS LB policy's config.
+type lbConfig struct {
+	serviceconfig.LoadBalancingConfig
+
+	cacheSizeBytes       int64 // Keep this field 64-bit aligned.
+	kbMap                keys.BuilderMap
+	lookupService        string
+	lookupServiceTimeout time.Duration
+	maxAge               time.Duration
+	staleAge             time.Duration
+	defaultTarget        string
+
+	childPolicyName             string
+	childPolicyConfig           map[string]json.RawMessage
+	childPolicyTargetField      string
+	controlChannelServiceConfig string
+}
+
+func (lbCfg *lbConfig) Equal(other *lbConfig) bool {
+	return lbCfg.kbMap.Equal(other.kbMap) &&
+		lbCfg.lookupService == other.lookupService &&
+		lbCfg.lookupServiceTimeout == other.lookupServiceTimeout &&
+		lbCfg.maxAge == other.maxAge &&
+		lbCfg.staleAge == other.staleAge &&
+		lbCfg.cacheSizeBytes == other.cacheSizeBytes &&
+		lbCfg.defaultTarget == other.defaultTarget &&
+		lbCfg.childPolicyName == other.childPolicyName &&
+		lbCfg.childPolicyTargetField == other.childPolicyTargetField &&
+		lbCfg.controlChannelServiceConfig == other.controlChannelServiceConfig &&
+		childPolicyConfigEqual(lbCfg.childPolicyConfig, other.childPolicyConfig)
+}
+
+func childPolicyConfigEqual(a, b map[string]json.RawMessage) bool {
+	if (b == nil) != (a == nil) {
+		return false
+	}
+	if len(b) != len(a) {
+		return false
+	}
+	for k, jsonA := range a {
+		jsonB, ok := b[k]
+		if !ok {
+			return false
+		}
+		if !bytes.Equal(jsonA, jsonB) {
+			return false
+		}
+	}
+	return true
+}
+
+// This struct resembles the JSON representation of the loadBalancing config
+// and makes it easier to unmarshal.
+type lbConfigJSON struct {
+	RouteLookupConfig                json.RawMessage
+	RouteLookupChannelServiceConfig  json.RawMessage
+	ChildPolicy                      []map[string]json.RawMessage
+	ChildPolicyConfigTargetFieldName string
+}
+
+// ParseConfig parses the JSON load balancer config provided into an
+// internal form or returns an error if the config is invalid.
+//
+//	 When parsing a config update, the following validations are performed:
+//	 - routeLookupConfig:
+//	   - grpc_keybuilders field:
+//	     - must have at least one entry
+//	     - must not have two entries with the same `Name`
+//	     - within each entry:
+//	       - must have at least one `Name`
+//	       - must not have a `Name` with the `service` field unset or empty
+//	       - within each `headers` entry:
+//	         - must not have `required_match` set
+//	         - must not have `key` unset or empty
+//	       - across all `headers`, `constant_keys` and `extra_keys` fields:
+//	         - must not have the same `key` specified twice
+//	         - no `key` must be the empty string
+//	   - `lookup_service` field must be set and must parse as a target URI
+//	   - if `max_age` > 5m, it should be set to 5 minutes
+//	   - if `stale_age` > `max_age`, ignore it
+//	   - if `stale_age` is set, then `max_age` must also be set
+//	   - ignore `valid_targets` field
+//	   - `cache_size_bytes` field must have a value greater than 0, and if its
+//	     value is greater than 5M, we cap it at 5M
+//
+//	- routeLookupChannelServiceConfig:
+//	  - if specified, must parse as valid service config
+//
+//	- childPolicy:
+//	  - must find a valid child policy with a valid config
+//
+//	- childPolicyConfigTargetFieldName:
+//	  - must be set and non-empty
+func (rlsBB) ParseConfig(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	if logger.V(2) {
+		logger.Infof("Received JSON service config: %v", pretty.ToJSON(c))
+	}
+
+	cfgJSON := &lbConfigJSON{}
+	if err := json.Unmarshal(c, cfgJSON); err != nil {
+		return nil, fmt.Errorf("rls: json unmarshal failed for service config %+v: %v", string(c), err)
+	}
+
+	m := protojson.UnmarshalOptions{DiscardUnknown: true}
+	rlsProto := &rlspb.RouteLookupConfig{}
+	if err := m.Unmarshal(cfgJSON.RouteLookupConfig, rlsProto); err != nil {
+		return nil, fmt.Errorf("rls: bad RouteLookupConfig proto %+v: %v", string(cfgJSON.RouteLookupConfig), err)
+	}
+	lbCfg, err := parseRLSProto(rlsProto)
+	if err != nil {
+		return nil, err
+	}
+
+	if sc := string(cfgJSON.RouteLookupChannelServiceConfig); sc != "" {
+		parsed := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(sc)
+		if parsed.Err != nil {
+			return nil, fmt.Errorf("rls: bad control channel service config %q: %v", sc, parsed.Err)
+		}
+		lbCfg.controlChannelServiceConfig = sc
+	}
+
+	if cfgJSON.ChildPolicyConfigTargetFieldName == "" {
+		return nil, fmt.Errorf("rls: childPolicyConfigTargetFieldName field is not set in service config %+v", string(c))
+	}
+	name, config, err := parseChildPolicyConfigs(cfgJSON.ChildPolicy, cfgJSON.ChildPolicyConfigTargetFieldName)
+	if err != nil {
+		return nil, err
+	}
+	lbCfg.childPolicyName = name
+	lbCfg.childPolicyConfig = config
+	lbCfg.childPolicyTargetField = cfgJSON.ChildPolicyConfigTargetFieldName
+	return lbCfg, nil
+}
+
+func parseRLSProto(rlsProto *rlspb.RouteLookupConfig) (*lbConfig, error) {
+	// Validations specified on the `grpc_keybuilders` field are performed here.
+	kbMap, err := keys.MakeBuilderMap(rlsProto)
+	if err != nil {
+		return nil, err
+	}
+
+	// `lookup_service` field must be set and must parse as a target URI.
+	lookupService := rlsProto.GetLookupService()
+	if lookupService == "" {
+		return nil, fmt.Errorf("rls: empty lookup_service in route lookup config %+v", rlsProto)
+	}
+	parsedTarget, err := url.Parse(lookupService)
+	if err != nil {
+		// url.Parse() fails if scheme is missing. Retry with default scheme.
+		parsedTarget, err = url.Parse(resolver.GetDefaultScheme() + ":///" + lookupService)
+		if err != nil {
+			return nil, fmt.Errorf("rls: invalid target URI in lookup_service %s", lookupService)
+		}
+	}
+	if parsedTarget.Scheme == "" {
+		parsedTarget.Scheme = resolver.GetDefaultScheme()
+	}
+	if resolver.Get(parsedTarget.Scheme) == nil {
+		return nil, fmt.Errorf("rls: unregistered scheme in lookup_service %s", lookupService)
+	}
+
+	lookupServiceTimeout, err := convertDuration(rlsProto.GetLookupServiceTimeout())
+	if err != nil {
+		return nil, fmt.Errorf("rls: failed to parse lookup_service_timeout in route lookup config %+v: %v", rlsProto, err)
+	}
+	if lookupServiceTimeout == 0 {
+		lookupServiceTimeout = defaultLookupServiceTimeout
+	}
+
+	// Validations performed here:
+	// - if `max_age` > 5m, it should be set to 5 minutes
+	//   only if stale age is not set
+	// - if `stale_age` > `max_age`, ignore it
+	// - if `stale_age` is set, then `max_age` must also be set
+	maxAgeSet := false
+	maxAge, err := convertDuration(rlsProto.GetMaxAge())
+	if err != nil {
+		return nil, fmt.Errorf("rls: failed to parse max_age in route lookup config %+v: %v", rlsProto, err)
+	}
+	if maxAge == 0 {
+		maxAge = maxMaxAge
+	} else {
+		maxAgeSet = true
+	}
+
+	staleAgeSet := false
+	staleAge, err := convertDuration(rlsProto.GetStaleAge())
+	if err != nil {
+		return nil, fmt.Errorf("rls: failed to parse staleAge in route lookup config %+v: %v", rlsProto, err)
+	}
+	if staleAge == 0 {
+		staleAge = maxMaxAge
+	} else {
+		staleAgeSet = true
+	}
+
+	if staleAgeSet && !maxAgeSet {
+		return nil, fmt.Errorf("rls: stale_age is set, but max_age is not in route lookup config %+v", rlsProto)
+	}
+	if staleAge > maxMaxAge {
+		staleAge = maxMaxAge
+	}
+	if !staleAgeSet && maxAge > maxMaxAge {
+		maxAge = maxMaxAge
+	}
+	if staleAge > maxAge {
+		staleAge = maxAge
+	}
+
+	// `cache_size_bytes` field must have a value greater than 0, and if its
+	// value is greater than 5M, we cap it at 5M
+	cacheSizeBytes := rlsProto.GetCacheSizeBytes()
+	if cacheSizeBytes <= 0 {
+		return nil, fmt.Errorf("rls: cache_size_bytes must be set to a non-zero value: %+v", rlsProto)
+	}
+	if cacheSizeBytes > maxCacheSize {
+		logger.Infof("rls: cache_size_bytes %v is too large, setting it to: %v", cacheSizeBytes, maxCacheSize)
+		cacheSizeBytes = maxCacheSize
+	}
+	return &lbConfig{
+		kbMap:                kbMap,
+		lookupService:        lookupService,
+		lookupServiceTimeout: lookupServiceTimeout,
+		maxAge:               maxAge,
+		staleAge:             staleAge,
+		cacheSizeBytes:       cacheSizeBytes,
+		defaultTarget:        rlsProto.GetDefaultTarget(),
+	}, nil
+}
+
+// parseChildPolicyConfigs iterates through the list of child policies and picks
+// the first registered policy and validates its config.
+func parseChildPolicyConfigs(childPolicies []map[string]json.RawMessage, targetFieldName string) (string, map[string]json.RawMessage, error) {
+	for i, config := range childPolicies {
+		if len(config) != 1 {
+			return "", nil, fmt.Errorf("rls: invalid childPolicy: entry %v does not contain exactly 1 policy/config pair: %q", i, config)
+		}
+
+		var name string
+		var rawCfg json.RawMessage
+		for name, rawCfg = range config {
+		}
+		builder := balancer.Get(name)
+		if builder == nil {
+			continue
+		}
+		parser, ok := builder.(balancer.ConfigParser)
+		if !ok {
+			return "", nil, fmt.Errorf("rls: childPolicy %q with config %q does not support config parsing", name, string(rawCfg))
+		}
+
+		// To validate child policy configs we do the following:
+		// - unmarshal the raw JSON bytes of the child policy config into a map
+		// - add an entry with key set to `target_field_name` and a dummy value
+		// - marshal the map back to JSON and parse the config using the parser
+		// retrieved previously
+		var childConfig map[string]json.RawMessage
+		if err := json.Unmarshal(rawCfg, &childConfig); err != nil {
+			return "", nil, fmt.Errorf("rls: json unmarshal failed for child policy config %q: %v", string(rawCfg), err)
+		}
+		childConfig[targetFieldName], _ = json.Marshal(dummyChildPolicyTarget)
+		jsonCfg, err := json.Marshal(childConfig)
+		if err != nil {
+			return "", nil, fmt.Errorf("rls: json marshal failed for child policy config {%+v}: %v", childConfig, err)
+		}
+		if _, err := parser.ParseConfig(jsonCfg); err != nil {
+			return "", nil, fmt.Errorf("rls: childPolicy config validation failed: %v", err)
+		}
+		return name, childConfig, nil
+	}
+	return "", nil, fmt.Errorf("rls: invalid childPolicy config: no supported policies found in %+v", childPolicies)
+}
+
+func convertDuration(d *durationpb.Duration) (time.Duration, error) {
+	if d == nil {
+		return 0, nil
+	}
+	return d.AsDuration(), d.CheckValid()
+}

balancer/rls/config_test.go

@@ -0,0 +1,515 @@
+/*
+ *
+ * Copyright 2020 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 rls
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	_ "google.golang.org/grpc/balancer/grpclb"               // grpclb for config parsing.
+	_ "google.golang.org/grpc/internal/resolver/passthrough" // passthrough resolver.
+)
+
+// testEqual reports whether the lbCfgs a and b are equal. This is to be used
+// only from tests. This ignores the keyBuilderMap field because its internals
+// are not exported, and hence not possible to specify in the want section of
+// the test. This is fine because we already have tests to make sure that the
+// keyBuilder is parsed properly from the service config.
+func testEqual(a, b *lbConfig) bool {
+	return a.lookupService == b.lookupService &&
+		a.lookupServiceTimeout == b.lookupServiceTimeout &&
+		a.maxAge == b.maxAge &&
+		a.staleAge == b.staleAge &&
+		a.cacheSizeBytes == b.cacheSizeBytes &&
+		a.defaultTarget == b.defaultTarget &&
+		a.controlChannelServiceConfig == b.controlChannelServiceConfig &&
+		a.childPolicyName == b.childPolicyName &&
+		a.childPolicyTargetField == b.childPolicyTargetField &&
+		childPolicyConfigEqual(a.childPolicyConfig, b.childPolicyConfig)
+}
+
+// TestParseConfig verifies successful config parsing scenarios.
+func (s) TestParseConfig(t *testing.T) {
+	childPolicyTargetFieldVal, _ := json.Marshal(dummyChildPolicyTarget)
+	tests := []struct {
+		desc    string
+		input   []byte
+		wantCfg *lbConfig
+	}{
+		{
+			// This input validates a few cases:
+			// - A top-level unknown field should not fail.
+			// - An unknown field in routeLookupConfig proto should not fail.
+			// - lookupServiceTimeout is set to its default value, since it is not specified in the input.
+			// - maxAge is clamped to maxMaxAge if staleAge is not set.
+			// - staleAge is ignored because it is higher than maxAge in the input.
+			// - cacheSizeBytes is greater than the hard upper limit of 5MB
+			desc: "with transformations 1",
+			input: []byte(`{
+				"top-level-unknown-field": "unknown-value",
+				"routeLookupConfig": {
+					"unknown-field": "unknown-value",
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": ":///target",
+					"maxAge" : "500s",
+					"staleAge": "600s",
+					"cacheSizeBytes": 100000000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicy": [
+					{"cds_experimental": {"Cluster": "my-fav-cluster"}},
+					{"unknown-policy": {"unknown-field": "unknown-value"}},
+					{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantCfg: &lbConfig{
+				lookupService:          ":///target",
+				lookupServiceTimeout:   10 * time.Second,  // This is the default value.
+				maxAge:                 500 * time.Second, // Max age is not clamped when stale age is set.
+				staleAge:               300 * time.Second, // StaleAge is clamped because it was higher than maxMaxAge.
+				cacheSizeBytes:         maxCacheSize,
+				defaultTarget:          "passthrough:///default",
+				childPolicyName:        "grpclb",
+				childPolicyTargetField: "serviceName",
+				childPolicyConfig: map[string]json.RawMessage{
+					"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`),
+					"serviceName": json.RawMessage(childPolicyTargetFieldVal),
+				},
+			},
+		},
+		{
+			desc: "maxAge not clamped when staleAge is set",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": ":///target",
+					"maxAge" : "500s",
+					"staleAge": "200s",
+					"cacheSizeBytes": 100000000
+				},
+				"childPolicy": [
+					{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantCfg: &lbConfig{
+				lookupService:          ":///target",
+				lookupServiceTimeout:   10 * time.Second,  // This is the default value.
+				maxAge:                 500 * time.Second, // Max age is not clamped when stale age is set.
+				staleAge:               200 * time.Second, // This is stale age within maxMaxAge.
+				cacheSizeBytes:         maxCacheSize,
+				childPolicyName:        "grpclb",
+				childPolicyTargetField: "serviceName",
+				childPolicyConfig: map[string]json.RawMessage{
+					"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`),
+					"serviceName": json.RawMessage(childPolicyTargetFieldVal),
+				},
+			},
+		},
+		{
+			desc: "maxAge clamped when staleAge is not set",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": ":///target",
+					"maxAge" : "500s",
+					"cacheSizeBytes": 100000000
+				},
+				"childPolicy": [
+					{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantCfg: &lbConfig{
+				lookupService:          ":///target",
+				lookupServiceTimeout:   10 * time.Second,  // This is the default value.
+				maxAge:                 300 * time.Second, // Max age is clamped when stale age is not set.
+				staleAge:               300 * time.Second,
+				cacheSizeBytes:         maxCacheSize,
+				childPolicyName:        "grpclb",
+				childPolicyTargetField: "serviceName",
+				childPolicyConfig: map[string]json.RawMessage{
+					"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`),
+					"serviceName": json.RawMessage(childPolicyTargetFieldVal),
+				},
+			},
+		},
+		{
+			desc: "without transformations",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "target",
+					"lookupServiceTimeout" : "100s",
+					"maxAge": "60s",
+					"staleAge" : "50s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"routeLookupChannelServiceConfig": {"loadBalancingConfig": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}]},
+				"childPolicy": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantCfg: &lbConfig{
+				lookupService:               "target",
+				lookupServiceTimeout:        100 * time.Second,
+				maxAge:                      60 * time.Second,
+				staleAge:                    50 * time.Second,
+				cacheSizeBytes:              1000,
+				defaultTarget:               "passthrough:///default",
+				controlChannelServiceConfig: `{"loadBalancingConfig": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}]}`,
+				childPolicyName:             "grpclb",
+				childPolicyTargetField:      "serviceName",
+				childPolicyConfig: map[string]json.RawMessage{
+					"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`),
+					"serviceName": json.RawMessage(childPolicyTargetFieldVal),
+				},
+			},
+		},
+	}
+
+	builder := rlsBB{}
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			lbCfg, err := builder.ParseConfig(test.input)
+			if err != nil || !testEqual(lbCfg.(*lbConfig), test.wantCfg) {
+				t.Errorf("ParseConfig(%s) = {%+v, %v}, want {%+v, nil}", string(test.input), lbCfg, err, test.wantCfg)
+			}
+		})
+	}
+}
+
+// TestParseConfigErrors verifies config parsing failure scenarios.
+func (s) TestParseConfigErrors(t *testing.T) {
+	tests := []struct {
+		desc    string
+		input   []byte
+		wantErr string
+	}{
+		{
+			desc:    "empty input",
+			input:   nil,
+			wantErr: "rls: json unmarshal failed for service config",
+		},
+		{
+			desc:    "bad json",
+			input:   []byte(`bad bad json`),
+			wantErr: "rls: json unmarshal failed for service config",
+		},
+		{
+			desc: "bad grpcKeyBuilder",
+			input: []byte(`{
+					"routeLookupConfig": {
+						"grpcKeybuilders": [{
+							"names": [{"service": "service", "method": "method"}],
+							"headers": [{"key": "k1", "requiredMatch": true, "names": ["v1"]}]
+						}]
+					}
+				}`),
+			wantErr: "rls: GrpcKeyBuilder in RouteLookupConfig has required_match field set",
+		},
+		{
+			desc: "empty lookup service",
+			input: []byte(`{
+					"routeLookupConfig": {
+						"grpcKeybuilders": [{
+							"names": [{"service": "service", "method": "method"}],
+							"headers": [{"key": "k1", "names": ["v1"]}]
+						}]
+					}
+				}`),
+			wantErr: "rls: empty lookup_service in route lookup config",
+		},
+		{
+			desc: "unregistered scheme in lookup service URI",
+			input: []byte(`{
+					"routeLookupConfig": {
+						"grpcKeybuilders": [{
+							"names": [{"service": "service", "method": "method"}],
+							"headers": [{"key": "k1", "names": ["v1"]}]
+						}],
+						"lookupService": "badScheme:///target"
+					}
+				}`),
+			wantErr: "rls: unregistered scheme in lookup_service",
+		},
+		{
+			desc: "invalid lookup service timeout",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "315576000001s"
+				}
+			}`),
+			wantErr: "google.protobuf.Duration value out of range",
+		},
+		{
+			desc: "invalid max age",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge" : "315576000001s"
+				}
+			}`),
+			wantErr: "google.protobuf.Duration value out of range",
+		},
+		{
+			desc: "invalid stale age",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge" : "10s",
+					"staleAge" : "315576000001s"
+				}
+			}`),
+			wantErr: "google.protobuf.Duration value out of range",
+		},
+		{
+			desc: "invalid max age stale age combo",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"staleAge" : "10s"
+				}
+			}`),
+			wantErr: "rls: stale_age is set, but max_age is not in route lookup config",
+		},
+		{
+			desc: "cache_size_bytes field is not set",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge": "30s",
+					"staleAge" : "25s",
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "rls: cache_size_bytes must be set to a non-zero value",
+		},
+		{
+			desc: "routeLookupChannelServiceConfig is not in service config format",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "target",
+					"lookupServiceTimeout" : "100s",
+					"maxAge": "60s",
+					"staleAge" : "50s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"routeLookupChannelServiceConfig": "unknown",
+				"childPolicy": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "cannot unmarshal string into Go value of type grpc.jsonSC",
+		},
+		{
+			desc: "routeLookupChannelServiceConfig contains unknown LB policy",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "target",
+					"lookupServiceTimeout" : "100s",
+					"maxAge": "60s",
+					"staleAge" : "50s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"routeLookupChannelServiceConfig": {
+					"loadBalancingConfig": [{"not_a_balancer1": {} }, {"not_a_balancer2": {}}]
+				},
+				"childPolicy": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "no supported policies found in config",
+		},
+		{
+			desc: "no child policy",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge": "30s",
+					"staleAge" : "25s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "rls: invalid childPolicy config: no supported policies found",
+		},
+		{
+			desc: "no known child policy",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge": "30s",
+					"staleAge" : "25s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicy": [
+					{"cds_experimental": {"Cluster": "my-fav-cluster"}},
+					{"unknown-policy": {"unknown-field": "unknown-value"}}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "rls: invalid childPolicy config: no supported policies found",
+		},
+		{
+			desc: "invalid child policy config - more than one entry in map",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge": "30s",
+					"staleAge" : "25s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicy": [
+					{
+						"cds_experimental": {"Cluster": "my-fav-cluster"},
+						"unknown-policy": {"unknown-field": "unknown-value"}
+					}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "does not contain exactly 1 policy/config pair",
+		},
+		{
+			desc: "no childPolicyConfigTargetFieldName",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge": "30s",
+					"staleAge" : "25s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicy": [
+					{"cds_experimental": {"Cluster": "my-fav-cluster"}},
+					{"unknown-policy": {"unknown-field": "unknown-value"}},
+					{"grpclb": {}}
+				]
+			}`),
+			wantErr: "rls: childPolicyConfigTargetFieldName field is not set in service config",
+		},
+		{
+			desc: "child policy config validation failure",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge": "30s",
+					"staleAge" : "25s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicy": [
+					{"cds_experimental": {"Cluster": "my-fav-cluster"}},
+					{"unknown-policy": {"unknown-field": "unknown-value"}},
+					{"grpclb": {"childPolicy": "not-an-array"}}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "rls: childPolicy config validation failed",
+		},
+	}
+
+	builder := rlsBB{}
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			lbCfg, err := builder.ParseConfig(test.input)
+			if lbCfg != nil || !strings.Contains(fmt.Sprint(err), test.wantErr) {
+				t.Errorf("ParseConfig(%s) = {%+v, %v}, want {nil, %s}", string(test.input), lbCfg, err, test.wantErr)
+			}
+		})
+	}
+}

balancer/rls/control_channel.go

@@ -0,0 +1,219 @@
+/*
+ *
+ * Copyright 2021 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 rls
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/rls/internal/adaptive"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/pretty"
+	rlsgrpc "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+)
+
+var newAdaptiveThrottler = func() adaptiveThrottler { return adaptive.New() }
+
+// newConnectivityStateSubscriber is a variable that can be overridden in tests
+// to wrap the connectivity state subscriber for testing purposes.
+var newConnectivityStateSubscriber = func(sub grpcsync.Subscriber) grpcsync.Subscriber {
+	return sub
+}
+
+type adaptiveThrottler interface {
+	ShouldThrottle() bool
+	RegisterBackendResponse(throttled bool)
+}
+
+// controlChannel is a wrapper around the gRPC channel to the RLS server
+// specified in the service config.
+type controlChannel struct {
+	// rpcTimeout specifies the timeout for the RouteLookup RPC call. The LB
+	// policy receives this value in its service config.
+	rpcTimeout time.Duration
+	// backToReadyFunc is a callback to be invoked when the connectivity state
+	// changes from READY --> TRANSIENT_FAILURE --> READY.
+	backToReadyFunc func()
+	// throttler in an adaptive throttling implementation used to avoid
+	// hammering the RLS service while it is overloaded or down.
+	throttler adaptiveThrottler
+
+	cc                      *grpc.ClientConn
+	client                  rlsgrpc.RouteLookupServiceClient
+	logger                  *internalgrpclog.PrefixLogger
+	dropConnStateSubscriber func()
+	seenTransientFailure    bool
+}
+
+// newControlChannel creates a controlChannel to rlsServerName and uses
+// serviceConfig, if non-empty, as the default service config for the underlying
+// gRPC channel.
+func newControlChannel(rlsServerName, serviceConfig string, rpcTimeout time.Duration, bOpts balancer.BuildOptions, backToReadyFunc func()) (*controlChannel, error) {
+	ctrlCh := &controlChannel{
+		rpcTimeout:      rpcTimeout,
+		backToReadyFunc: backToReadyFunc,
+		throttler:       newAdaptiveThrottler(),
+	}
+	ctrlCh.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[rls-control-channel %p] ", ctrlCh))
+
+	dopts, err := ctrlCh.dialOpts(bOpts, serviceConfig)
+	if err != nil {
+		return nil, err
+	}
+	ctrlCh.cc, err = grpc.NewClient(rlsServerName, dopts...)
+	if err != nil {
+		return nil, err
+	}
+	// Subscribe to connectivity state before connecting to avoid missing initial
+	// updates, which are only delivered to active subscribers.
+	subscribe := internal.SubscribeToConnectivityStateChanges.(func(cc *grpc.ClientConn, s grpcsync.Subscriber) func())
+	ctrlCh.dropConnStateSubscriber = subscribe(ctrlCh.cc, newConnectivityStateSubscriber(ctrlCh))
+	ctrlCh.cc.Connect()
+	ctrlCh.client = rlsgrpc.NewRouteLookupServiceClient(ctrlCh.cc)
+	ctrlCh.logger.Infof("Control channel created to RLS server at: %v", rlsServerName)
+	return ctrlCh, nil
+}
+
+func (cc *controlChannel) OnMessage(msg any) {
+	st, ok := msg.(connectivity.State)
+	if !ok {
+		panic(fmt.Sprintf("Unexpected message type %T , wanted connectectivity.State type", msg))
+	}
+
+	switch st {
+	case connectivity.Ready:
+		// Only reset backoff when transitioning from TRANSIENT_FAILURE to READY.
+		// This indicates the RLS server has recovered from being unreachable, so
+		// we reset backoff state in all cache entries to allow pending RPCs to
+		// proceed immediately. We skip benign transitions like READY → IDLE → READY
+		// since those don't represent actual failures.
+		if cc.seenTransientFailure {
+			if cc.logger.V(2) {
+				cc.logger.Infof("Control channel back to READY after TRANSIENT_FAILURE")
+			}
+			cc.seenTransientFailure = false
+			if cc.backToReadyFunc != nil {
+				cc.backToReadyFunc()
+			}
+		} else {
+			if cc.logger.V(2) {
+				cc.logger.Infof("Control channel is READY")
+			}
+		}
+	case connectivity.TransientFailure:
+		// Track that we've entered TRANSIENT_FAILURE state so we know to reset
+		// backoffs when we recover to READY.
+		cc.logger.Warningf("Control channel is TRANSIENT_FAILURE")
+		cc.seenTransientFailure = true
+	default:
+		if cc.logger.V(2) {
+			cc.logger.Infof("Control channel connectivity state is %s", st)
+		}
+	}
+}
+
+// dialOpts constructs the dial options for the control plane channel.
+func (cc *controlChannel) dialOpts(bOpts balancer.BuildOptions, serviceConfig string) ([]grpc.DialOption, error) {
+	// The control plane channel will use the same authority as the parent
+	// channel for server authorization. This ensures that the identity of the
+	// RLS server and the identity of the backends is the same, so if the RLS
+	// config is injected by an attacker, it cannot cause leakage of private
+	// information contained in headers set by the application.
+	dopts := []grpc.DialOption{grpc.WithAuthority(bOpts.Authority)}
+	if bOpts.Dialer != nil {
+		dopts = append(dopts, grpc.WithContextDialer(bOpts.Dialer))
+	}
+	// The control channel will use the channel credentials from the parent
+	// channel, including any call creds associated with the channel creds.
+	var credsOpt grpc.DialOption
+	switch {
+	case bOpts.DialCreds != nil:
+		credsOpt = grpc.WithTransportCredentials(bOpts.DialCreds.Clone())
+	case bOpts.CredsBundle != nil:
+		// The "fallback" mode in google default credentials (which is the only
+		// type of credentials we expect to be used with RLS) uses TLS/ALTS
+		// creds for transport and uses the same call creds as that on the
+		// parent bundle.
+		bundle, err := bOpts.CredsBundle.NewWithMode(internal.CredsBundleModeFallback)
+		if err != nil {
+			return nil, err
+		}
+		credsOpt = grpc.WithCredentialsBundle(bundle)
+	default:
+		cc.logger.Warningf("no credentials available, using Insecure")
+		credsOpt = grpc.WithTransportCredentials(insecure.NewCredentials())
+	}
+	dopts = append(dopts, credsOpt)
+
+	// If the RLS LB policy's configuration specified a service config for the
+	// control channel, use that and disable service config fetching via the name
+	// resolver for the control channel.
+	if serviceConfig != "" {
+		cc.logger.Infof("Disabling service config from the name resolver and instead using: %s", serviceConfig)
+		dopts = append(dopts, grpc.WithDisableServiceConfig(), grpc.WithDefaultServiceConfig(serviceConfig))
+	}
+
+	return dopts, nil
+}
+
+func (cc *controlChannel) close() {
+	cc.dropConnStateSubscriber()
+	cc.cc.Close()
+	cc.logger.Infof("Shutdown")
+}
+
+type lookupCallback func(targets []string, headerData string, err error)
+
+// lookup starts a RouteLookup RPC in a separate goroutine and returns the
+// results (and error, if any) in the provided callback.
+//
+// The returned boolean indicates whether the request was throttled by the
+// client-side adaptive throttling algorithm in which case the provided callback
+// will not be invoked.
+func (cc *controlChannel) lookup(reqKeys map[string]string, reason rlspb.RouteLookupRequest_Reason, staleHeaders string, cb lookupCallback) (throttled bool) {
+	if cc.throttler.ShouldThrottle() {
+		cc.logger.Infof("RLS request throttled by client-side adaptive throttling")
+		return true
+	}
+	go func() {
+		req := &rlspb.RouteLookupRequest{
+			TargetType:      "grpc",
+			KeyMap:          reqKeys,
+			Reason:          reason,
+			StaleHeaderData: staleHeaders,
+		}
+		if cc.logger.V(2) {
+			cc.logger.Infof("Sending RLS request %+v", pretty.ToJSON(req))
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), cc.rpcTimeout)
+		defer cancel()
+		resp, err := cc.client.RouteLookup(ctx, req)
+		cb(resp.GetTargets(), resp.GetHeaderData(), err)
+	}()
+	return false
+}

balancer/rls/control_channel_test.go

@@ -0,0 +1,465 @@
+/*
+ *
+ * Copyright 2021 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 rls
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"fmt"
+	"os"
+	"regexp"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/internal"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	rlstest "google.golang.org/grpc/internal/testutils/rls"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/testdata"
+	"google.golang.org/protobuf/proto"
+)
+
+// TestControlChannelThrottled tests the case where the adaptive throttler
+// indicates that the control channel needs to be throttled.
+func (s) TestControlChannelThrottled(t *testing.T) {
+	// Start an RLS server and set the throttler to always throttle requests.
+	rlsServer, rlsReqCh := rlstest.SetupFakeRLSServer(t, nil)
+	overrideAdaptiveThrottler(t, alwaysThrottlingThrottler())
+
+	// Create a control channel to the fake RLS server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, balancer.BuildOptions{}, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect the attempt to be throttled.
+	ctrlCh.lookup(nil, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, nil)
+
+	select {
+	case <-rlsReqCh:
+		t.Fatal("RouteLookup RPC invoked when control channel is throttled")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+// TestLookupFailure tests the case where the RLS server responds with an error.
+func (s) TestLookupFailure(t *testing.T) {
+	// Start an RLS server and set the throttler to never throttle requests.
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil)
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Setup the RLS server to respond with errors.
+	rlsServer.SetResponseCallback(func(context.Context, *rlspb.RouteLookupRequest) *rlstest.RouteLookupResponse {
+		return &rlstest.RouteLookupResponse{Err: errors.New("rls failure")}
+	})
+
+	// Create a control channel to the fake RLS server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, balancer.BuildOptions{}, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect the callback to be invoked with an error.
+	errCh := make(chan error, 1)
+	ctrlCh.lookup(nil, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, func(_ []string, _ string, err error) {
+		if err == nil {
+			errCh <- errors.New("rlsClient.lookup() succeeded, should have failed")
+			return
+		}
+		errCh <- nil
+	})
+
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("timeout when waiting for lookup callback to be invoked")
+	case err := <-errCh:
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// TestLookupDeadlineExceeded tests the case where the RLS server does not
+// respond within the configured rpc timeout.
+func (s) TestLookupDeadlineExceeded(t *testing.T) {
+	// A unary interceptor which returns a status error with DeadlineExceeded.
+	interceptor := func(context.Context, any, *grpc.UnaryServerInfo, grpc.UnaryHandler) (resp any, err error) {
+		return nil, status.Error(codes.DeadlineExceeded, "deadline exceeded")
+	}
+
+	// Start an RLS server and set the throttler to never throttle.
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil, grpc.UnaryInterceptor(interceptor))
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Create a control channel with a small deadline.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestShortTimeout, balancer.BuildOptions{}, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect the callback to be invoked with an error.
+	errCh := make(chan error, 1)
+	ctrlCh.lookup(nil, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, func(_ []string, _ string, err error) {
+		if st, ok := status.FromError(err); !ok || st.Code() != codes.DeadlineExceeded {
+			errCh <- fmt.Errorf("rlsClient.lookup() returned error: %v, want %v", err, codes.DeadlineExceeded)
+			return
+		}
+		errCh <- nil
+	})
+
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("timeout when waiting for lookup callback to be invoked")
+	case err := <-errCh:
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// testCredsBundle wraps a test call creds and real transport creds.
+type testCredsBundle struct {
+	transportCreds credentials.TransportCredentials
+	callCreds      credentials.PerRPCCredentials
+}
+
+func (f *testCredsBundle) TransportCredentials() credentials.TransportCredentials {
+	return f.transportCreds
+}
+
+func (f *testCredsBundle) PerRPCCredentials() credentials.PerRPCCredentials {
+	return f.callCreds
+}
+
+func (f *testCredsBundle) NewWithMode(mode string) (credentials.Bundle, error) {
+	if mode != internal.CredsBundleModeFallback {
+		return nil, fmt.Errorf("unsupported mode: %v", mode)
+	}
+	return &testCredsBundle{
+		transportCreds: f.transportCreds,
+		callCreds:      f.callCreds,
+	}, nil
+}
+
+var (
+	// Call creds sent by the testPerRPCCredentials on the client, and verified
+	// by an interceptor on the server.
+	perRPCCredsData = map[string]string{
+		"test-key":     "test-value",
+		"test-key-bin": string([]byte{1, 2, 3}),
+	}
+)
+
+type testPerRPCCredentials struct {
+	callCreds map[string]string
+}
+
+func (f *testPerRPCCredentials) GetRequestMetadata(context.Context, ...string) (map[string]string, error) {
+	return f.callCreds, nil
+}
+
+func (f *testPerRPCCredentials) RequireTransportSecurity() bool {
+	return true
+}
+
+// Unary server interceptor which validates if the RPC contains call credentials
+// which match `perRPCCredsData
+func callCredsValidatingServerInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return nil, status.Error(codes.PermissionDenied, "didn't find metadata in context")
+	}
+	for k, want := range perRPCCredsData {
+		got, ok := md[k]
+		if !ok {
+			return ctx, status.Errorf(codes.PermissionDenied, "didn't find call creds key %v in context", k)
+		}
+		if got[0] != want {
+			return ctx, status.Errorf(codes.PermissionDenied, "for key %v, got value %v, want %v", k, got, want)
+		}
+	}
+	return handler(ctx, req)
+}
+
+// makeTLSCreds is a test helper which creates a TLS based transport credentials
+// from files specified in the arguments.
+func makeTLSCreds(t *testing.T, certPath, keyPath, rootsPath string) credentials.TransportCredentials {
+	cert, err := tls.LoadX509KeyPair(testdata.Path(certPath), testdata.Path(keyPath))
+	if err != nil {
+		t.Fatalf("tls.LoadX509KeyPair(%q, %q) failed: %v", certPath, keyPath, err)
+	}
+	b, err := os.ReadFile(testdata.Path(rootsPath))
+	if err != nil {
+		t.Fatalf("os.ReadFile(%q) failed: %v", rootsPath, err)
+	}
+	roots := x509.NewCertPool()
+	if !roots.AppendCertsFromPEM(b) {
+		t.Fatal("failed to append certificates")
+	}
+	return credentials.NewTLS(&tls.Config{
+		Certificates: []tls.Certificate{cert},
+		RootCAs:      roots,
+	})
+}
+
+const (
+	wantHeaderData  = "headerData"
+	staleHeaderData = "staleHeaderData"
+)
+
+var (
+	keyMap = map[string]string{
+		"k1": "v1",
+		"k2": "v2",
+	}
+	wantTargets   = []string{"us_east_1.firestore.googleapis.com"}
+	lookupRequest = &rlspb.RouteLookupRequest{
+		TargetType:      "grpc",
+		KeyMap:          keyMap,
+		Reason:          rlspb.RouteLookupRequest_REASON_MISS,
+		StaleHeaderData: staleHeaderData,
+	}
+	lookupResponse = &rlstest.RouteLookupResponse{
+		Resp: &rlspb.RouteLookupResponse{
+			Targets:    wantTargets,
+			HeaderData: wantHeaderData,
+		},
+	}
+)
+
+func testControlChannelCredsSuccess(t *testing.T, sopts []grpc.ServerOption, bopts balancer.BuildOptions) {
+	// Start an RLS server and set the throttler to never throttle requests.
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil, sopts...)
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Setup the RLS server to respond with a valid response.
+	rlsServer.SetResponseCallback(func(context.Context, *rlspb.RouteLookupRequest) *rlstest.RouteLookupResponse {
+		return lookupResponse
+	})
+
+	// Verify that the request received by the RLS matches the expected one.
+	rlsServer.SetRequestCallback(func(got *rlspb.RouteLookupRequest) {
+		if diff := cmp.Diff(lookupRequest, got, cmp.Comparer(proto.Equal)); diff != "" {
+			t.Errorf("RouteLookupRequest diff (-want, +got):\n%s", diff)
+		}
+	})
+
+	// Create a control channel to the fake server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, bopts, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect a successful callback invocation.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	errCh := make(chan error, 1)
+	ctrlCh.lookup(keyMap, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, func(targets []string, headerData string, err error) {
+		if err != nil {
+			errCh <- fmt.Errorf("rlsClient.lookup() failed with err: %v", err)
+			return
+		}
+		if !cmp.Equal(targets, wantTargets) || headerData != wantHeaderData {
+			errCh <- fmt.Errorf("rlsClient.lookup() = (%v, %s), want (%v, %s)", targets, headerData, wantTargets, wantHeaderData)
+			return
+		}
+		errCh <- nil
+	})
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout when waiting for lookup callback to be invoked")
+	case err := <-errCh:
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// TestControlChannelCredsSuccess tests creation of the control channel with
+// different credentials, which are expected to succeed.
+func (s) TestControlChannelCredsSuccess(t *testing.T) {
+	serverCreds := makeTLSCreds(t, "x509/server1_cert.pem", "x509/server1_key.pem", "x509/client_ca_cert.pem")
+	clientCreds := makeTLSCreds(t, "x509/client1_cert.pem", "x509/client1_key.pem", "x509/server_ca_cert.pem")
+
+	tests := []struct {
+		name  string
+		sopts []grpc.ServerOption
+		bopts balancer.BuildOptions
+	}{
+		{
+			name:  "insecure",
+			sopts: nil,
+			bopts: balancer.BuildOptions{},
+		},
+		{
+			name:  "transport creds only",
+			sopts: []grpc.ServerOption{grpc.Creds(serverCreds)},
+			bopts: balancer.BuildOptions{
+				DialCreds: clientCreds,
+				Authority: "x.test.example.com",
+			},
+		},
+		{
+			name: "creds bundle",
+			sopts: []grpc.ServerOption{
+				grpc.Creds(serverCreds),
+				grpc.UnaryInterceptor(callCredsValidatingServerInterceptor),
+			},
+			bopts: balancer.BuildOptions{
+				CredsBundle: &testCredsBundle{
+					transportCreds: clientCreds,
+					callCreds:      &testPerRPCCredentials{callCreds: perRPCCredsData},
+				},
+				Authority: "x.test.example.com",
+			},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			testControlChannelCredsSuccess(t, test.sopts, test.bopts)
+		})
+	}
+}
+
+func testControlChannelCredsFailure(t *testing.T, sopts []grpc.ServerOption, bopts balancer.BuildOptions, wantCode codes.Code, wantErrRegex *regexp.Regexp) {
+	// StartFakeRouteLookupServer a fake server.
+	//
+	// Start an RLS server and set the throttler to never throttle requests. The
+	// creds failures happen before the RPC handler on the server is invoked.
+	// So, there is need to setup the request and responses on the fake server.
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil, sopts...)
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Create the control channel to the fake server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, bopts, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect the callback to be invoked with an error.
+	errCh := make(chan error, 1)
+	ctrlCh.lookup(nil, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, func(_ []string, _ string, err error) {
+		if st, ok := status.FromError(err); !ok || st.Code() != wantCode || !wantErrRegex.MatchString(st.String()) {
+			errCh <- fmt.Errorf("rlsClient.lookup() returned error: %v, wantCode: %v, wantErr: %s", err, wantCode, wantErrRegex.String())
+			return
+		}
+		errCh <- nil
+	})
+
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("timeout when waiting for lookup callback to be invoked")
+	case err := <-errCh:
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// TestControlChannelCredsFailure tests creation of the control channel with
+// different credentials, which are expected to fail.
+func (s) TestControlChannelCredsFailure(t *testing.T) {
+	serverCreds := makeTLSCreds(t, "x509/server1_cert.pem", "x509/server1_key.pem", "x509/client_ca_cert.pem")
+	clientCreds := makeTLSCreds(t, "x509/client1_cert.pem", "x509/client1_key.pem", "x509/server_ca_cert.pem")
+
+	tests := []struct {
+		name         string
+		sopts        []grpc.ServerOption
+		bopts        balancer.BuildOptions
+		wantCode     codes.Code
+		wantErrRegex *regexp.Regexp
+	}{
+		{
+			name:  "transport creds authority mismatch",
+			sopts: []grpc.ServerOption{grpc.Creds(serverCreds)},
+			bopts: balancer.BuildOptions{
+				DialCreds: clientCreds,
+				Authority: "authority-mismatch",
+			},
+			wantCode:     codes.Unavailable,
+			wantErrRegex: regexp.MustCompile(`transport: authentication handshake failed: .* \*\.test\.example\.com.*authority-mismatch`),
+		},
+		{
+			name:  "transport creds handshake failure",
+			sopts: nil, // server expects insecure connection
+			bopts: balancer.BuildOptions{
+				DialCreds: clientCreds,
+				Authority: "x.test.example.com",
+			},
+			wantCode:     codes.Unavailable,
+			wantErrRegex: regexp.MustCompile("transport: authentication handshake failed: .*"),
+		},
+		{
+			name: "call creds mismatch",
+			sopts: []grpc.ServerOption{
+				grpc.Creds(serverCreds),
+				grpc.UnaryInterceptor(callCredsValidatingServerInterceptor), // server expects call creds
+			},
+			bopts: balancer.BuildOptions{
+				CredsBundle: &testCredsBundle{
+					transportCreds: clientCreds,
+					callCreds:      &testPerRPCCredentials{}, // sends no call creds
+				},
+				Authority: "x.test.example.com",
+			},
+			wantCode:     codes.PermissionDenied,
+			wantErrRegex: regexp.MustCompile("didn't find call creds"),
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			testControlChannelCredsFailure(t, test.sopts, test.bopts, test.wantCode, test.wantErrRegex)
+		})
+	}
+}
+
+type unsupportedCredsBundle struct {
+	credentials.Bundle
+}
+
+func (*unsupportedCredsBundle) NewWithMode(mode string) (credentials.Bundle, error) {
+	return nil, fmt.Errorf("unsupported mode: %v", mode)
+}
+
+// TestNewControlChannelUnsupportedCredsBundle tests the case where the control
+// channel is configured with a bundle which does not support the mode we use.
+func (s) TestNewControlChannelUnsupportedCredsBundle(t *testing.T) {
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil)
+
+	// Create the control channel to the fake server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, balancer.BuildOptions{CredsBundle: &unsupportedCredsBundle{}}, nil)
+	if err == nil {
+		ctrlCh.close()
+		t.Fatal("newControlChannel succeeded when expected to fail")
+	}
+}

balancer/rls/helpers_test.go

@@ -0,0 +1,299 @@
+/*
+ *
+ * Copyright 2021 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 rls
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer/rls/internal/test/e2e"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/grpctest"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/internal/stubserver"
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/grpc/serviceconfig"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/durationpb"
+)
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// fakeBackoffStrategy is a fake implementation of the backoff.Strategy
+// interface, for tests to inject the backoff duration.
+type fakeBackoffStrategy struct {
+	backoff time.Duration
+}
+
+func (f *fakeBackoffStrategy) Backoff(int) time.Duration {
+	return f.backoff
+}
+
+// fakeThrottler is a fake implementation of the adaptiveThrottler interface.
+type fakeThrottler struct {
+	throttleFunc func() bool   // Fake throttler implementation.
+	throttleCh   chan struct{} // Invocation of ShouldThrottle signals here.
+}
+
+func (f *fakeThrottler) ShouldThrottle() bool {
+	select {
+	case <-f.throttleCh:
+	default:
+	}
+	f.throttleCh <- struct{}{}
+
+	return f.throttleFunc()
+}
+
+func (f *fakeThrottler) RegisterBackendResponse(bool) {}
+
+// alwaysThrottlingThrottler returns a fake throttler which always throttles.
+func alwaysThrottlingThrottler() *fakeThrottler {
+	return &fakeThrottler{
+		throttleFunc: func() bool { return true },
+		throttleCh:   make(chan struct{}, 1),
+	}
+}
+
+// neverThrottlingThrottler returns a fake throttler which never throttles.
+func neverThrottlingThrottler() *fakeThrottler {
+	return &fakeThrottler{
+		throttleFunc: func() bool { return false },
+		throttleCh:   make(chan struct{}, 1),
+	}
+}
+
+// oneTimeAllowingThrottler returns a fake throttler which does not throttle
+// requests until the client RPC succeeds, but throttles everything that comes
+// after. This is useful for tests which need to set up a valid cache entry
+// before testing other cases.
+func oneTimeAllowingThrottler(firstRPCDone *grpcsync.Event) *fakeThrottler {
+	return &fakeThrottler{
+		throttleFunc: firstRPCDone.HasFired,
+		throttleCh:   make(chan struct{}, 1),
+	}
+}
+
+func overrideAdaptiveThrottler(t *testing.T, f *fakeThrottler) {
+	origAdaptiveThrottler := newAdaptiveThrottler
+	newAdaptiveThrottler = func() adaptiveThrottler { return f }
+	t.Cleanup(func() { newAdaptiveThrottler = origAdaptiveThrottler })
+}
+
+// buildBasicRLSConfig constructs a basic service config for the RLS LB policy
+// with header matching rules. This expects the passed child policy name to
+// have been registered by the caller.
+func buildBasicRLSConfig(childPolicyName, rlsServerAddress string) *e2e.RLSConfig {
+	return &e2e.RLSConfig{
+		RouteLookupConfig: &rlspb.RouteLookupConfig{
+			GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+				{
+					Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "grpc.testing.TestService"}},
+					Headers: []*rlspb.NameMatcher{
+						{Key: "k1", Names: []string{"n1"}},
+						{Key: "k2", Names: []string{"n2"}},
+					},
+				},
+			},
+			LookupService:        rlsServerAddress,
+			LookupServiceTimeout: durationpb.New(defaultTestTimeout),
+			CacheSizeBytes:       1024,
+		},
+		RouteLookupChannelServiceConfig:  `{"loadBalancingConfig": [{"pick_first": {}}]}`,
+		ChildPolicy:                      &internalserviceconfig.BalancerConfig{Name: childPolicyName},
+		ChildPolicyConfigTargetFieldName: e2e.RLSChildPolicyTargetNameField,
+	}
+}
+
+// buildBasicRLSConfigWithChildPolicy constructs a very basic service config for
+// the RLS LB policy. It also registers a test LB policy which is capable of
+// being a child of the RLS LB policy.
+func buildBasicRLSConfigWithChildPolicy(t *testing.T, childPolicyName, rlsServerAddress string) *e2e.RLSConfig {
+	childPolicyName = "test-child-policy" + childPolicyName
+	e2e.RegisterRLSChildPolicy(childPolicyName, nil)
+	t.Logf("Registered child policy with name %q", childPolicyName)
+
+	return &e2e.RLSConfig{
+		RouteLookupConfig: &rlspb.RouteLookupConfig{
+			GrpcKeybuilders:      []*rlspb.GrpcKeyBuilder{{Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "grpc.testing.TestService"}}}},
+			LookupService:        rlsServerAddress,
+			LookupServiceTimeout: durationpb.New(defaultTestTimeout),
+			CacheSizeBytes:       1024,
+		},
+		RouteLookupChannelServiceConfig:  `{"loadBalancingConfig": [{"pick_first": {}}]}`,
+		ChildPolicy:                      &internalserviceconfig.BalancerConfig{Name: childPolicyName},
+		ChildPolicyConfigTargetFieldName: e2e.RLSChildPolicyTargetNameField,
+	}
+}
+
+// startBackend starts a backend implementing the TestService on a local port.
+// It returns a channel for tests to get notified whenever an RPC is invoked on
+// the backend. This allows tests to ensure that RPCs reach expected backends.
+// Also returns the address of the backend.
+func startBackend(t *testing.T, sopts ...grpc.ServerOption) (rpcCh chan struct{}, address string) {
+	t.Helper()
+
+	rpcCh = make(chan struct{}, 1)
+	backend := &stubserver.StubServer{
+		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+			select {
+			case rpcCh <- struct{}{}:
+			default:
+			}
+			return &testpb.Empty{}, nil
+		},
+	}
+	if err := backend.StartServer(sopts...); err != nil {
+		t.Fatalf("Failed to start backend: %v", err)
+	}
+	t.Logf("Started TestService backend at: %q", backend.Address)
+	t.Cleanup(func() { backend.Stop() })
+	return rpcCh, backend.Address
+}
+
+// startManualResolverWithConfig registers and returns a manual resolver which
+// pushes the RLS LB policy's service config on the channel.
+func startManualResolverWithConfig(t *testing.T, rlsConfig *e2e.RLSConfig) *manual.Resolver {
+	t.Helper()
+
+	scJSON, err := rlsConfig.ServiceConfigJSON()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(scJSON)
+	r := manual.NewBuilderWithScheme("rls-e2e")
+	r.InitialState(resolver.State{ServiceConfig: sc})
+	t.Cleanup(r.Close)
+	return r
+}
+
+// makeTestRPCAndExpectItToReachBackend is a test helper function which makes
+// the EmptyCall RPC on the given ClientConn and verifies that it reaches a
+// backend. The latter is accomplished by listening on the provided channel
+// which gets pushed to whenever the backend in question gets an RPC.
+//
+// There are many instances where it can take a while before the attempted RPC
+// reaches the expected backend. Examples include, but are not limited to:
+//   - control channel is changed in a config update. The RLS LB policy creates a
+//     new control channel, and sends a new picker to gRPC. But it takes a while
+//     before gRPC actually starts using the new picker.
+//   - test is waiting for a cache entry to expire after which we expect a
+//     different behavior because we have configured the fake RLS server to return
+//     different backends.
+//
+// Therefore, we do not return an error when the RPC fails. Instead, we wait for
+// the context to expire before failing.
+func makeTestRPCAndExpectItToReachBackend(ctx context.Context, t *testing.T, cc *grpc.ClientConn, ch chan struct{}) {
+	t.Helper()
+
+	// Drain the backend channel before performing the RPC to remove any
+	// notifications from previous RPCs.
+	select {
+	case <-ch:
+	default:
+	}
+
+	for {
+		if err := ctx.Err(); err != nil {
+			t.Fatalf("Timeout when waiting for RPCs to be routed to the given target: %v", err)
+		}
+		sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+		client := testgrpc.NewTestServiceClient(cc)
+		client.EmptyCall(sCtx, &testpb.Empty{})
+
+		select {
+		case <-sCtx.Done():
+		case <-ch:
+			sCancel()
+			return
+		}
+	}
+}
+
+// makeTestRPCAndVerifyError is a test helper function which makes the EmptyCall
+// RPC on the given ClientConn and verifies that the RPC fails with the given
+// status code and error.
+//
+// Similar to makeTestRPCAndExpectItToReachBackend, retries until expected
+// outcome is reached or the provided context has expired.
+func makeTestRPCAndVerifyError(ctx context.Context, t *testing.T, cc *grpc.ClientConn, wantCode codes.Code, wantErr error) {
+	t.Helper()
+
+	for {
+		if err := ctx.Err(); err != nil {
+			t.Fatalf("Timeout when waiting for RPCs to fail with given error: %v", err)
+		}
+		sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+		client := testgrpc.NewTestServiceClient(cc)
+		_, err := client.EmptyCall(sCtx, &testpb.Empty{})
+
+		// If the RPC fails with the expected code and expected error message (if
+		// one was provided), we return. Else we retry after blocking for a little
+		// while to ensure that we don't keep blasting away with RPCs.
+		if code := status.Code(err); code == wantCode {
+			if wantErr == nil || strings.Contains(err.Error(), wantErr.Error()) {
+				sCancel()
+				return
+			}
+		}
+		<-sCtx.Done()
+	}
+}
+
+// verifyRLSRequest is a test helper which listens on a channel to see if an RLS
+// request was received by the fake RLS server. Based on whether the test
+// expects a request to be sent out or not, it uses a different timeout.
+func verifyRLSRequest(t *testing.T, ch chan struct{}, wantRequest bool) {
+	t.Helper()
+
+	if wantRequest {
+		select {
+		case <-time.After(defaultTestTimeout):
+			t.Fatalf("Timeout when waiting for an RLS request to be sent out")
+		case <-ch:
+		}
+	} else {
+		select {
+		case <-time.After(defaultTestShortTimeout):
+		case <-ch:
+			t.Fatalf("RLS request sent out when not expecting one")
+		}
+	}
+}

balancer/rls/internal/adaptive/adaptive.go

@@ -0,0 +1,131 @@
+/*
+ *
+ * Copyright 2020 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 adaptive provides functionality for adaptive client-side throttling.
+package adaptive
+
+import (
+	rand "math/rand/v2"
+	"sync"
+	"time"
+)
+
+// For overriding in unittests.
+var (
+	timeNowFunc = time.Now
+	randFunc    = rand.Float64
+)
+
+const (
+	defaultDuration        = 30 * time.Second
+	defaultBins            = 100
+	defaultRatioForAccepts = 2.0
+	defaultRequestsPadding = 8.0
+)
+
+// Throttler implements a client-side throttling recommendation system. All
+// methods are safe for concurrent use by multiple goroutines.
+//
+// The throttler has the following knobs for which we will use defaults for
+// now. If there is a need to make them configurable at a later point in time,
+// support for the same will be added.
+//   - Duration: amount of recent history that will be taken into account for
+//     making client-side throttling decisions. A default of 30 seconds is used.
+//   - Bins: number of bins to be used for bucketing historical data. A default
+//     of 100 is used.
+//   - RatioForAccepts: ratio by which accepts are multiplied, typically a value
+//     slightly larger than 1.0. This is used to make the throttler behave as if
+//     the backend had accepted more requests than it actually has, which lets us
+//     err on the side of sending to the backend more requests than we think it
+//     will accept for the sake of speeding up the propagation of state. A
+//     default of 2.0 is used.
+//   - RequestsPadding: is used to decrease the (client-side) throttling
+//     probability in the low QPS regime (to speed up propagation of state), as
+//     well as to safeguard against hitting a client-side throttling probability
+//     of 100%. The weight of this value decreases as the number of requests in
+//     recent history grows. A default of 8 is used.
+//
+// The adaptive throttler attempts to estimate the probability that a request
+// will be throttled using recent history. Server requests (both throttled and
+// accepted) are registered with the throttler (via the RegisterBackendResponse
+// method), which then recommends client-side throttling (via the
+// ShouldThrottle method) with probability given by:
+// (requests - RatioForAccepts * accepts) / (requests + RequestsPadding)
+type Throttler struct {
+	ratioForAccepts float64
+	requestsPadding float64
+
+	// Number of total accepts and throttles in the lookback period.
+	mu        sync.Mutex
+	accepts   *lookback
+	throttles *lookback
+}
+
+// New initializes a new adaptive throttler with the default values.
+func New() *Throttler {
+	return newWithArgs(defaultDuration, defaultBins, defaultRatioForAccepts, defaultRequestsPadding)
+}
+
+// newWithArgs initializes a new adaptive throttler with the provided values.
+// Used only in unittests.
+func newWithArgs(duration time.Duration, bins int64, ratioForAccepts, requestsPadding float64) *Throttler {
+	return &Throttler{
+		ratioForAccepts: ratioForAccepts,
+		requestsPadding: requestsPadding,
+		accepts:         newLookback(bins, duration),
+		throttles:       newLookback(bins, duration),
+	}
+}
+
+// ShouldThrottle returns a probabilistic estimate of whether the server would
+// throttle the next request. This should be called for every request before
+// allowing it to hit the network. If the returned value is true, the request
+// should be aborted immediately (as if it had been throttled by the server).
+func (t *Throttler) ShouldThrottle() bool {
+	randomProbability := randFunc()
+	now := timeNowFunc()
+
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	accepts, throttles := float64(t.accepts.sum(now)), float64(t.throttles.sum(now))
+	requests := accepts + throttles
+	throttleProbability := (requests - t.ratioForAccepts*accepts) / (requests + t.requestsPadding)
+	if throttleProbability <= randomProbability {
+		return false
+	}
+
+	t.throttles.add(now, 1)
+	return true
+}
+
+// RegisterBackendResponse registers a response received from the backend for a
+// request allowed by ShouldThrottle. This should be called for every response
+// received from the backend (i.e., once for each request for which
+// ShouldThrottle returned false).
+func (t *Throttler) RegisterBackendResponse(throttled bool) {
+	now := timeNowFunc()
+
+	t.mu.Lock()
+	if throttled {
+		t.throttles.add(now, 1)
+	} else {
+		t.accepts.add(now, 1)
+	}
+	t.mu.Unlock()
+}

balancer/rls/internal/adaptive/adaptive_test.go

@@ -0,0 +1,231 @@
+/*
+ *
+ * Copyright 2020 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 adaptive
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+// stats returns a tuple with accepts, throttles for the current time.
+func (t *Throttler) stats() (int64, int64) {
+	now := timeNowFunc()
+
+	t.mu.Lock()
+	a, th := t.accepts.sum(now), t.throttles.sum(now)
+	t.mu.Unlock()
+	return a, th
+}
+
+// Enums for responses.
+const (
+	E = iota // No response
+	A        // Accepted
+	T        // Throttled
+)
+
+func TestRegisterBackendResponse(t *testing.T) {
+	testcases := []struct {
+		desc          string
+		bins          int64
+		ticks         []int64
+		responses     []int64
+		wantAccepts   []int64
+		wantThrottled []int64
+	}{
+		{
+			"Accumulate",
+			3,
+			[]int64{0, 1, 2}, // Ticks
+			[]int64{A, T, E}, // Responses
+			[]int64{1, 1, 1}, // Accepts
+			[]int64{0, 1, 1}, // Throttled
+		},
+		{
+			"LightTimeTravel",
+			3,
+			[]int64{1, 0, 2}, // Ticks
+			[]int64{A, T, E}, // Response
+			[]int64{1, 1, 1}, // Accepts
+			[]int64{0, 1, 1}, // Throttled
+		},
+		{
+			"HeavyTimeTravel",
+			3,
+			[]int64{8, 0, 9}, // Ticks
+			[]int64{A, A, A}, // Response
+			[]int64{1, 1, 2}, // Accepts
+			[]int64{0, 0, 0}, // Throttled
+		},
+		{
+			"Rollover",
+			1,
+			[]int64{0, 1, 2}, // Ticks
+			[]int64{A, T, E}, // Responses
+			[]int64{1, 0, 0}, // Accepts
+			[]int64{0, 1, 0}, // Throttled
+		},
+	}
+
+	m := mockClock{}
+	oldTimeNowFunc := timeNowFunc
+	timeNowFunc = m.Now
+	defer func() { timeNowFunc = oldTimeNowFunc }()
+
+	for _, test := range testcases {
+		t.Run(test.desc, func(t *testing.T) {
+			th := newWithArgs(time.Duration(test.bins), test.bins, 2.0, 8)
+			for i, tick := range test.ticks {
+				m.SetNanos(tick)
+
+				if test.responses[i] != E {
+					th.RegisterBackendResponse(test.responses[i] == T)
+				}
+
+				if gotAccepts, gotThrottled := th.stats(); gotAccepts != test.wantAccepts[i] || gotThrottled != test.wantThrottled[i] {
+					t.Errorf("th.stats() = {%d, %d} for index %d, want {%d, %d}", i, gotAccepts, gotThrottled, test.wantAccepts[i], test.wantThrottled[i])
+				}
+			}
+		})
+	}
+}
+
+func TestShouldThrottleOptions(t *testing.T) {
+	// ShouldThrottle should return true iff
+	//    (requests - RatioForAccepts * accepts) / (requests + RequestsPadding) <= p
+	// where p is a random number. For the purposes of this test it's fixed
+	// to 0.5.
+	responses := []int64{T, T, T, T, T, T, T, T, T, A, A, A, A, A, A, T, T, T, T}
+
+	n := false
+	y := true
+
+	testcases := []struct {
+		desc            string
+		ratioForAccepts float64
+		requestsPadding float64
+		want            []bool
+	}{
+		{
+			"Baseline",
+			1.1,
+			8,
+			[]bool{n, n, n, n, n, n, n, n, y, y, y, y, y, n, n, n, y, y, y},
+		},
+		{
+			"ChangePadding",
+			1.1,
+			7,
+			[]bool{n, n, n, n, n, n, n, y, y, y, y, y, y, y, y, y, y, y, y},
+		},
+		{
+			"ChangeRatioForAccepts",
+			1.4,
+			8,
+			[]bool{n, n, n, n, n, n, n, n, y, y, n, n, n, n, n, n, n, n, n},
+		},
+	}
+
+	m := mockClock{}
+	oldTimeNowFunc := timeNowFunc
+	timeNowFunc = m.Now
+	oldRandFunc := randFunc
+	randFunc = func() float64 { return 0.5 }
+	defer func() {
+		timeNowFunc = oldTimeNowFunc
+		randFunc = oldRandFunc
+	}()
+
+	for _, test := range testcases {
+		t.Run(test.desc, func(t *testing.T) {
+			m.SetNanos(0)
+			th := newWithArgs(time.Nanosecond, 1, test.ratioForAccepts, test.requestsPadding)
+			for i, response := range responses {
+				if response != E {
+					th.RegisterBackendResponse(response == T)
+				}
+				if got := th.ShouldThrottle(); got != test.want[i] {
+					t.Errorf("ShouldThrottle for index %d: got %v, want %v", i, got, test.want[i])
+				}
+			}
+		})
+	}
+}
+
+func TestParallel(t *testing.T) {
+	// Uses all the defaults which comes with a 30 second duration.
+	th := New()
+
+	testDuration := 2 * time.Second
+	numRoutines := 10
+	accepts := make([]int64, numRoutines)
+	throttles := make([]int64, numRoutines)
+	var wg sync.WaitGroup
+	for i := 0; i < numRoutines; i++ {
+		wg.Add(1)
+		go func(num int) {
+			defer wg.Done()
+
+			ticker := time.NewTicker(testDuration)
+			var accept int64
+			var throttle int64
+			for i := 0; ; i++ {
+				select {
+				case <-ticker.C:
+					ticker.Stop()
+					accepts[num] = accept
+					throttles[num] = throttle
+					return
+				default:
+					if i%2 == 0 {
+						th.RegisterBackendResponse(true)
+						throttle++
+					} else {
+						th.RegisterBackendResponse(false)
+						accept++
+					}
+				}
+			}
+		}(i)
+	}
+	wg.Wait()
+
+	var wantAccepts, wantThrottles int64
+	for i := 0; i < numRoutines; i++ {
+		wantAccepts += accepts[i]
+		wantThrottles += throttles[i]
+	}
+
+	if gotAccepts, gotThrottles := th.stats(); gotAccepts != wantAccepts || gotThrottles != wantThrottles {
+		t.Errorf("th.stats() = {%d, %d}, want {%d, %d}", gotAccepts, gotThrottles, wantAccepts, wantThrottles)
+	}
+}
+
+type mockClock struct {
+	t time.Time
+}
+
+func (m *mockClock) Now() time.Time {
+	return m.t
+}
+
+func (m *mockClock) SetNanos(n int64) {
+	m.t = time.Unix(0, n)
+}

balancer/rls/internal/adaptive/lookback.go

@@ -0,0 +1,84 @@
+/*
+ *
+ * Copyright 2020 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 adaptive
+
+import "time"
+
+// lookback implements a moving sum over an int64 timeline.
+type lookback struct {
+	bins  int64         // Number of bins to use for lookback.
+	width time.Duration // Width of each bin.
+
+	head  int64   // Absolute bin index (time * bins / duration) of the current head bin.
+	total int64   // Sum over all the values in buf, within the lookback window behind head.
+	buf   []int64 // Ring buffer for keeping track of the sum elements.
+}
+
+// newLookback creates a new lookback for the given duration with a set number
+// of bins.
+func newLookback(bins int64, duration time.Duration) *lookback {
+	return &lookback{
+		bins:  bins,
+		width: duration / time.Duration(bins),
+		buf:   make([]int64, bins),
+	}
+}
+
+// add is used to increment the lookback sum.
+func (l *lookback) add(t time.Time, v int64) {
+	pos := l.advance(t)
+
+	if (l.head - pos) >= l.bins {
+		// Do not increment counters if pos is more than bins behind head.
+		return
+	}
+	l.buf[pos%l.bins] += v
+	l.total += v
+}
+
+// sum returns the sum of the lookback buffer at the given time or head,
+// whichever is greater.
+func (l *lookback) sum(t time.Time) int64 {
+	l.advance(t)
+	return l.total
+}
+
+// advance prepares the lookback buffer for calls to add() or sum() at time t.
+// If head is greater than t then the lookback buffer will be untouched. The
+// absolute bin index corresponding to t is returned. It will always be less
+// than or equal to head.
+func (l *lookback) advance(t time.Time) int64 {
+	ch := l.head                               // Current head bin index.
+	nh := t.UnixNano() / l.width.Nanoseconds() // New head bin index.
+
+	if nh <= ch {
+		// Either head unchanged or clock jitter (time has moved backwards). Do
+		// not advance.
+		return nh
+	}
+
+	jmax := min(l.bins, nh-ch)
+	for j := int64(0); j < jmax; j++ {
+		i := (ch + j + 1) % l.bins
+		l.total -= l.buf[i]
+		l.buf[i] = 0
+	}
+	l.head = nh
+	return nh
+}

balancer/rls/internal/adaptive/lookback_test.go

@@ -0,0 +1,85 @@
+/*
+ *
+ * Copyright 2020 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 adaptive
+
+import (
+	"testing"
+	"time"
+)
+
+func TestLookback(t *testing.T) {
+	makeTicks := func(offsets []int64) []time.Time {
+		var ticks []time.Time
+		now := time.Now()
+		for _, offset := range offsets {
+			ticks = append(ticks, now.Add(time.Duration(offset)))
+		}
+		return ticks
+	}
+
+	// lookback.add and lookback.sum behave correctly.
+	testcases := []struct {
+		desc   string
+		bins   int64
+		ticks  []time.Time
+		values []int64
+		want   []int64
+	}{
+		{
+			"Accumulate",
+			3,
+			makeTicks([]int64{0, 1, 2}), // Ticks
+			[]int64{1, 2, 3},            // Values
+			[]int64{1, 3, 6},            // Want
+		},
+		{
+			"LightTimeTravel",
+			3,
+			makeTicks([]int64{1, 0, 2}), // Ticks
+			[]int64{1, 2, 3},            // Values
+			[]int64{1, 3, 6},            // Want
+		},
+		{
+			"HeavyTimeTravel",
+			3,
+			makeTicks([]int64{8, 0, 9}), // Ticks
+			[]int64{1, 2, 3},            // Values
+			[]int64{1, 1, 4},            // Want
+		},
+		{
+			"Rollover",
+			1,
+			makeTicks([]int64{0, 1, 2}), // Ticks
+			[]int64{1, 2, 3},            // Values
+			[]int64{1, 2, 3},            // Want
+		},
+	}
+
+	for _, test := range testcases {
+		t.Run(test.desc, func(t *testing.T) {
+			lb := newLookback(test.bins, time.Duration(test.bins))
+			for i, tick := range test.ticks {
+				lb.add(tick, test.values[i])
+				if got := lb.sum(tick); got != test.want[i] {
+					t.Errorf("sum for index %d got %d, want %d", i, got, test.want[i])
+				}
+			}
+		})
+	}
+}

balancer/rls/internal/keys/builder.go

@@ -0,0 +1,267 @@
+/*
+ *
+ * Copyright 2020 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 keys provides functionality required to build RLS request keys.
+package keys
+
+import (
+	"errors"
+	"fmt"
+	"sort"
+	"strings"
+
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	"google.golang.org/grpc/metadata"
+)
+
+// BuilderMap maps from request path to the key builder for that path.
+type BuilderMap map[string]builder
+
+// MakeBuilderMap parses the provided RouteLookupConfig proto and returns a map
+// from paths to key builders.
+func MakeBuilderMap(cfg *rlspb.RouteLookupConfig) (BuilderMap, error) {
+	kbs := cfg.GetGrpcKeybuilders()
+	if len(kbs) == 0 {
+		return nil, errors.New("rls: RouteLookupConfig does not contain any GrpcKeyBuilder")
+	}
+
+	bm := make(map[string]builder)
+	for _, kb := range kbs {
+		// Extract keys from `headers`, `constant_keys` and `extra_keys` fields
+		// and populate appropriate values in the builder struct. Also ensure
+		// that keys are not repeated.
+		var matchers []matcher
+		seenKeys := make(map[string]bool)
+		constantKeys := kb.GetConstantKeys()
+		for k := range kb.GetConstantKeys() {
+			seenKeys[k] = true
+		}
+		for _, h := range kb.GetHeaders() {
+			if h.GetRequiredMatch() {
+				return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig has required_match field set {%+v}", kbs)
+			}
+			key := h.GetKey()
+			if seenKeys[key] {
+				return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key %q across headers, constant_keys and extra_keys {%+v}", key, kbs)
+			}
+			seenKeys[key] = true
+			matchers = append(matchers, matcher{key: h.GetKey(), names: h.GetNames()})
+		}
+		if seenKeys[kb.GetExtraKeys().GetHost()] {
+			return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key %q in extra_keys from constant_keys or headers {%+v}", kb.GetExtraKeys().GetHost(), kbs)
+		}
+		if seenKeys[kb.GetExtraKeys().GetService()] {
+			return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key %q in extra_keys from constant_keys or headers {%+v}", kb.GetExtraKeys().GetService(), kbs)
+		}
+		if seenKeys[kb.GetExtraKeys().GetMethod()] {
+			return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key %q in extra_keys from constant_keys or headers {%+v}", kb.GetExtraKeys().GetMethod(), kbs)
+		}
+		b := builder{
+			headerKeys:   matchers,
+			constantKeys: constantKeys,
+			hostKey:      kb.GetExtraKeys().GetHost(),
+			serviceKey:   kb.GetExtraKeys().GetService(),
+			methodKey:    kb.GetExtraKeys().GetMethod(),
+		}
+
+		// Store the builder created above in the BuilderMap based on the value
+		// of the `Names` field, which wraps incoming request's service and
+		// method. Also, ensure that there are no repeated `Names` field.
+		names := kb.GetNames()
+		if len(names) == 0 {
+			return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig does not contain any Name {%+v}", kbs)
+		}
+		for _, name := range names {
+			if name.GetService() == "" {
+				return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains a Name field with no Service {%+v}", kbs)
+			}
+			if strings.Contains(name.GetMethod(), `/`) {
+				return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains a method with a slash {%+v}", kbs)
+			}
+			path := "/" + name.GetService() + "/" + name.GetMethod()
+			if _, ok := bm[path]; ok {
+				return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated Name field {%+v}", kbs)
+			}
+			bm[path] = b
+		}
+	}
+	return bm, nil
+}
+
+// KeyMap represents the RLS keys to be used for a request.
+type KeyMap struct {
+	// Map is the representation of an RLS key as a Go map. This is used when
+	// an actual RLS request is to be sent out on the wire, since the
+	// RouteLookupRequest proto expects a Go map.
+	Map map[string]string
+	// Str is the representation of an RLS key as a string, sorted by keys.
+	// Since the RLS keys are part of the cache key in the request cache
+	// maintained by the RLS balancer, and Go maps cannot be used as keys for
+	// Go maps (the cache is implemented as a map), we need a stringified
+	// version of it.
+	Str string
+}
+
+// RLSKey builds the RLS keys to be used for the given request, identified by
+// the request path and the request headers stored in metadata.
+func (bm BuilderMap) RLSKey(md metadata.MD, host, path string) KeyMap {
+	// The path passed in is of the form "/service/method". The keyBuilderMap is
+	// indexed with keys of the form "/service/" or "/service/method". The service
+	// that we set in the keyMap (to be sent out in the RLS request) should not
+	// include any slashes though.
+	i := strings.LastIndex(path, "/")
+	service, method := path[:i+1], path[i+1:]
+	b, ok := bm[path]
+	if !ok {
+		b, ok = bm[service]
+		if !ok {
+			return KeyMap{}
+		}
+	}
+
+	kvMap := b.buildHeaderKeys(md)
+	if b.hostKey != "" {
+		kvMap[b.hostKey] = host
+	}
+	if b.serviceKey != "" {
+		kvMap[b.serviceKey] = strings.Trim(service, "/")
+	}
+	if b.methodKey != "" {
+		kvMap[b.methodKey] = method
+	}
+	for k, v := range b.constantKeys {
+		kvMap[k] = v
+	}
+	return KeyMap{Map: kvMap, Str: mapToString(kvMap)}
+}
+
+// Equal reports whether bm and am represent equivalent BuilderMaps.
+func (bm BuilderMap) Equal(am BuilderMap) bool {
+	if (bm == nil) != (am == nil) {
+		return false
+	}
+	if len(bm) != len(am) {
+		return false
+	}
+
+	for key, bBuilder := range bm {
+		aBuilder, ok := am[key]
+		if !ok {
+			return false
+		}
+		if !bBuilder.Equal(aBuilder) {
+			return false
+		}
+	}
+	return true
+}
+
+// builder provides the actual functionality of building RLS keys.
+type builder struct {
+	headerKeys   []matcher
+	constantKeys map[string]string
+	// The following keys mirror corresponding fields in `extra_keys`.
+	hostKey    string
+	serviceKey string
+	methodKey  string
+}
+
+// Equal reports whether b and a represent equivalent key builders.
+func (b builder) Equal(a builder) bool {
+	if len(b.headerKeys) != len(a.headerKeys) {
+		return false
+	}
+	// Protobuf serialization maintains the order of repeated fields. Matchers
+	// are specified as a repeated field inside the KeyBuilder proto. If the
+	// order changes, it means that the order in the protobuf changed. We report
+	// this case as not being equal even though the builders could possibly be
+	// functionally equal.
+	for i, bMatcher := range b.headerKeys {
+		aMatcher := a.headerKeys[i]
+		if !bMatcher.Equal(aMatcher) {
+			return false
+		}
+	}
+
+	if len(b.constantKeys) != len(a.constantKeys) {
+		return false
+	}
+	for k, v := range b.constantKeys {
+		if a.constantKeys[k] != v {
+			return false
+		}
+	}
+
+	return b.hostKey == a.hostKey && b.serviceKey == a.serviceKey && b.methodKey == a.methodKey
+}
+
+// matcher helps extract a key from request headers based on a given name.
+type matcher struct {
+	// The key used in the keyMap sent as part of the RLS request.
+	key string
+	// List of header names which can supply the value for this key.
+	names []string
+}
+
+// Equal reports if m and a are equivalent headerKeys.
+func (m matcher) Equal(a matcher) bool {
+	if m.key != a.key {
+		return false
+	}
+	if len(m.names) != len(a.names) {
+		return false
+	}
+	for i := 0; i < len(m.names); i++ {
+		if m.names[i] != a.names[i] {
+			return false
+		}
+	}
+	return true
+}
+
+func (b builder) buildHeaderKeys(md metadata.MD) map[string]string {
+	kvMap := make(map[string]string)
+	if len(md) == 0 {
+		return kvMap
+	}
+	for _, m := range b.headerKeys {
+		for _, name := range m.names {
+			if vals := md.Get(name); vals != nil {
+				kvMap[m.key] = strings.Join(vals, ",")
+				break
+			}
+		}
+	}
+	return kvMap
+}
+
+func mapToString(kv map[string]string) string {
+	keys := make([]string, 0, len(kv))
+	for k := range kv {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	var sb strings.Builder
+	for i, k := range keys {
+		if i != 0 {
+			fmt.Fprint(&sb, ",")
+		}
+		fmt.Fprintf(&sb, "%s=%s", k, kv[k])
+	}
+	return sb.String()
+}

balancer/rls/internal/keys/builder_test.go

@@ -0,0 +1,650 @@
+/*
+ *
+ * Copyright 2020 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 keys
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	"google.golang.org/grpc/metadata"
+)
+
+var (
+	goodKeyBuilder1 = &rlspb.GrpcKeyBuilder{
+		Names: []*rlspb.GrpcKeyBuilder_Name{
+			{Service: "gFoo"},
+		},
+		Headers: []*rlspb.NameMatcher{
+			{Key: "k1", Names: []string{"n1"}},
+			{Key: "k2", Names: []string{"n1"}},
+		},
+		ExtraKeys: &rlspb.GrpcKeyBuilder_ExtraKeys{
+			Host:    "host",
+			Service: "service",
+			Method:  "method",
+		},
+		ConstantKeys: map[string]string{
+			"const-key-1": "const-val-1",
+			"const-key-2": "const-val-2",
+		},
+	}
+	goodKeyBuilder2 = &rlspb.GrpcKeyBuilder{
+		Names: []*rlspb.GrpcKeyBuilder_Name{
+			{Service: "gBar", Method: "method1"},
+			{Service: "gFoobar"},
+		},
+		Headers: []*rlspb.NameMatcher{
+			{Key: "k1", Names: []string{"n1", "n2"}},
+		},
+	}
+)
+
+func TestMakeBuilderMap(t *testing.T) {
+	gFooBuilder := builder{
+		headerKeys: []matcher{{key: "k1", names: []string{"n1"}}, {key: "k2", names: []string{"n1"}}},
+		constantKeys: map[string]string{
+			"const-key-1": "const-val-1",
+			"const-key-2": "const-val-2",
+		},
+		hostKey:    "host",
+		serviceKey: "service",
+		methodKey:  "method",
+	}
+	wantBuilderMap1 := map[string]builder{"/gFoo/": gFooBuilder}
+	wantBuilderMap2 := map[string]builder{
+		"/gFoo/":        gFooBuilder,
+		"/gBar/method1": {headerKeys: []matcher{{key: "k1", names: []string{"n1", "n2"}}}},
+		"/gFoobar/":     {headerKeys: []matcher{{key: "k1", names: []string{"n1", "n2"}}}},
+	}
+
+	tests := []struct {
+		desc           string
+		cfg            *rlspb.RouteLookupConfig
+		wantBuilderMap BuilderMap
+	}{
+		{
+			desc: "One good GrpcKeyBuilder",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{goodKeyBuilder1},
+			},
+			wantBuilderMap: wantBuilderMap1,
+		},
+		{
+			desc: "Two good GrpcKeyBuilders",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{goodKeyBuilder1, goodKeyBuilder2},
+			},
+			wantBuilderMap: wantBuilderMap2,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			builderMap, err := MakeBuilderMap(test.cfg)
+			if err != nil || !builderMap.Equal(test.wantBuilderMap) {
+				t.Errorf("MakeBuilderMap(%+v) returned {%v, %v}, want: {%v, nil}", test.cfg, builderMap, err, test.wantBuilderMap)
+			}
+		})
+	}
+}
+
+func TestMakeBuilderMapErrors(t *testing.T) {
+	tests := []struct {
+		desc          string
+		cfg           *rlspb.RouteLookupConfig
+		wantErrPrefix string
+	}{
+		{
+			desc:          "No GrpcKeyBuilder",
+			cfg:           &rlspb.RouteLookupConfig{},
+			wantErrPrefix: "rls: RouteLookupConfig does not contain any GrpcKeyBuilder",
+		},
+		{
+			desc: "Two GrpcKeyBuilders with same Name",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{goodKeyBuilder1, goodKeyBuilder1},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated Name field",
+		},
+		{
+			desc: "GrpcKeyBuilder with empty Service field",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+					{
+						Names: []*rlspb.GrpcKeyBuilder_Name{
+							{Service: "bFoo", Method: "method1"},
+							{Service: "bBar"},
+							{Method: "method1"},
+						},
+						Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}},
+					},
+					goodKeyBuilder1,
+				},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains a Name field with no Service",
+		},
+		{
+			desc: "GrpcKeyBuilder with no Name",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{{}, goodKeyBuilder1},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig does not contain any Name",
+		},
+		{
+			desc: "GrpcKeyBuilder with requiredMatch field set",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+					{
+						Names:   []*rlspb.GrpcKeyBuilder_Name{{Service: "bFoo", Method: "method1"}},
+						Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}, RequiredMatch: true}},
+					},
+					goodKeyBuilder1,
+				},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig has required_match field set",
+		},
+		{
+			desc: "GrpcKeyBuilder two headers with same key",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+					{
+						Names: []*rlspb.GrpcKeyBuilder_Name{
+							{Service: "gBar", Method: "method1"},
+							{Service: "gFoobar"},
+						},
+						Headers: []*rlspb.NameMatcher{
+							{Key: "k1", Names: []string{"n1", "n2"}},
+							{Key: "k1", Names: []string{"n1", "n2"}},
+						},
+					},
+					goodKeyBuilder1,
+				},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key \"k1\" across headers, constant_keys and extra_keys",
+		},
+		{
+			desc: "GrpcKeyBuilder repeated keys across headers and constant_keys",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+					{
+						Names: []*rlspb.GrpcKeyBuilder_Name{
+							{Service: "gBar", Method: "method1"},
+							{Service: "gFoobar"},
+						},
+						Headers:      []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}},
+						ConstantKeys: map[string]string{"k1": "v1"},
+					},
+					goodKeyBuilder1,
+				},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key \"k1\" across headers, constant_keys and extra_keys",
+		},
+		{
+			desc: "GrpcKeyBuilder repeated keys across headers and extra_keys",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+					{
+						Names: []*rlspb.GrpcKeyBuilder_Name{
+							{Service: "gBar", Method: "method1"},
+							{Service: "gFoobar"},
+						},
+						Headers:   []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}},
+						ExtraKeys: &rlspb.GrpcKeyBuilder_ExtraKeys{Method: "k1"},
+					},
+					goodKeyBuilder1,
+				},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key \"k1\" in extra_keys from constant_keys or headers",
+		},
+		{
+			desc: "GrpcKeyBuilder repeated keys across constant_keys and extra_keys",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+					{
+						Names: []*rlspb.GrpcKeyBuilder_Name{
+							{Service: "gBar", Method: "method1"},
+							{Service: "gFoobar"},
+						},
+						Headers:      []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}},
+						ConstantKeys: map[string]string{"host": "v1"},
+						ExtraKeys:    &rlspb.GrpcKeyBuilder_ExtraKeys{Host: "host"},
+					},
+					goodKeyBuilder1,
+				},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated key \"host\" in extra_keys from constant_keys or headers",
+		},
+		{
+			desc: "GrpcKeyBuilder with slash in method name",
+			cfg: &rlspb.RouteLookupConfig{
+				GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+					{
+						Names:   []*rlspb.GrpcKeyBuilder_Name{{Service: "gBar", Method: "method1/foo"}},
+						Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}},
+					},
+				},
+			},
+			wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains a method with a slash",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			builderMap, err := MakeBuilderMap(test.cfg)
+			if builderMap != nil || !strings.HasPrefix(fmt.Sprint(err), test.wantErrPrefix) {
+				t.Errorf("MakeBuilderMap(%+v) returned {%v, %v}, want: {nil, %v}", test.cfg, builderMap, err, test.wantErrPrefix)
+			}
+		})
+	}
+}
+
+func TestRLSKey(t *testing.T) {
+	bm, err := MakeBuilderMap(&rlspb.RouteLookupConfig{
+		GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{goodKeyBuilder1, goodKeyBuilder2},
+	})
+	if err != nil {
+		t.Fatalf("MakeBuilderMap() failed: %v", err)
+	}
+
+	tests := []struct {
+		desc   string
+		path   string
+		md     metadata.MD
+		wantKM KeyMap
+	}{
+		{
+			// No keyBuilder is found for the provided service.
+			desc:   "service not found in key builder map",
+			path:   "/notFoundService/method",
+			md:     metadata.Pairs("n1", "v1", "n2", "v2"),
+			wantKM: KeyMap{},
+		},
+		{
+			// No keyBuilder is found for the provided method.
+			desc:   "method not found in key builder map",
+			path:   "/gBar/notFoundMethod",
+			md:     metadata.Pairs("n1", "v1", "n2", "v2"),
+			wantKM: KeyMap{},
+		},
+		{
+			// A keyBuilder is found, but none of the headers match.
+			desc:   "directPathMatch-NoMatchingKey",
+			path:   "/gBar/method1",
+			md:     metadata.Pairs("notMatchingKey", "v1"),
+			wantKM: KeyMap{Map: map[string]string{}, Str: ""},
+		},
+		{
+			// A keyBuilder is found, and a single headers matches.
+			desc:   "directPathMatch-SingleKey",
+			path:   "/gBar/method1",
+			md:     metadata.Pairs("n1", "v1"),
+			wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"},
+		},
+		{
+			// A keyBuilder is found, and multiple headers match, but the first
+			// match is chosen.
+			desc:   "directPathMatch-FirstMatchingKey",
+			path:   "/gBar/method1",
+			md:     metadata.Pairs("n2", "v2", "n1", "v1"),
+			wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"},
+		},
+		{
+			// A keyBuilder is found as a wildcard match, but none of the
+			// headers match.
+			desc:   "wildcardPathMatch-NoMatchingKey",
+			path:   "/gFoobar/method1",
+			md:     metadata.Pairs("notMatchingKey", "v1"),
+			wantKM: KeyMap{Map: map[string]string{}, Str: ""},
+		},
+		{
+			// A keyBuilder is found as a wildcard match, and a single headers
+			// matches.
+			desc:   "wildcardPathMatch-SingleKey",
+			path:   "/gFoobar/method1",
+			md:     metadata.Pairs("n1", "v1"),
+			wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"},
+		},
+		{
+			// A keyBuilder is found as a wildcard match, and multiple headers
+			// match, but the first match is chosen.
+			desc:   "wildcardPathMatch-FirstMatchingKey",
+			path:   "/gFoobar/method1",
+			md:     metadata.Pairs("n2", "v2", "n1", "v1"),
+			wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"},
+		},
+		{
+			// Multiple headerKeys find hits in the provided request headers.
+			desc: "multipleMatchers",
+			path: "/gFoo/method1",
+			md:   metadata.Pairs("n2", "v2", "n1", "v1"),
+			wantKM: KeyMap{
+				Map: map[string]string{
+					"const-key-1": "const-val-1",
+					"const-key-2": "const-val-2",
+					"host":        "dummy-host",
+					"service":     "gFoo",
+					"method":      "method1",
+					"k1":          "v1",
+					"k2":          "v1",
+				},
+				Str: "const-key-1=const-val-1,const-key-2=const-val-2,host=dummy-host,k1=v1,k2=v1,method=method1,service=gFoo",
+			},
+		},
+		{
+			// A match is found for a header which is specified multiple times.
+			// So, the values are joined with commas separating them.
+			desc:   "commaSeparated",
+			path:   "/gBar/method1",
+			md:     metadata.Pairs("n1", "v1", "n1", "v2", "n1", "v3"),
+			wantKM: KeyMap{Map: map[string]string{"k1": "v1,v2,v3"}, Str: "k1=v1,v2,v3"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			if gotKM := bm.RLSKey(test.md, "dummy-host", test.path); !cmp.Equal(gotKM, test.wantKM) {
+				t.Errorf("RLSKey(%+v, %s) = %+v, want %+v", test.md, test.path, gotKM, test.wantKM)
+			}
+		})
+	}
+}
+
+func TestMapToString(t *testing.T) {
+	tests := []struct {
+		desc    string
+		input   map[string]string
+		wantStr string
+	}{
+		{
+			desc:    "empty map",
+			input:   nil,
+			wantStr: "",
+		},
+		{
+			desc: "one key",
+			input: map[string]string{
+				"k1": "v1",
+			},
+			wantStr: "k1=v1",
+		},
+		{
+			desc: "sorted keys",
+			input: map[string]string{
+				"k1": "v1",
+				"k2": "v2",
+				"k3": "v3",
+			},
+			wantStr: "k1=v1,k2=v2,k3=v3",
+		},
+		{
+			desc: "unsorted keys",
+			input: map[string]string{
+				"k3": "v3",
+				"k1": "v1",
+				"k2": "v2",
+			},
+			wantStr: "k1=v1,k2=v2,k3=v3",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			if gotStr := mapToString(test.input); gotStr != test.wantStr {
+				t.Errorf("mapToString(%v) = %s, want %s", test.input, gotStr, test.wantStr)
+			}
+		})
+	}
+}
+
+func TestBuilderMapEqual(t *testing.T) {
+	tests := []struct {
+		desc      string
+		a         BuilderMap
+		b         BuilderMap
+		wantEqual bool
+	}{
+		{
+			desc:      "nil builder maps",
+			a:         nil,
+			b:         nil,
+			wantEqual: true,
+		},
+		{
+			desc:      "empty builder maps",
+			a:         make(map[string]builder),
+			b:         make(map[string]builder),
+			wantEqual: true,
+		},
+		{
+			desc:      "nil and non-nil builder maps",
+			a:         nil,
+			b:         map[string]builder{"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}},
+			wantEqual: false,
+		},
+		{
+			desc:      "empty and non-empty builder maps",
+			a:         make(map[string]builder),
+			b:         map[string]builder{"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}}},
+			wantEqual: false,
+		},
+		{
+			desc: "different number of map keys",
+			a: map[string]builder{
+				"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+				"/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+			},
+			b: map[string]builder{
+				"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+			},
+			wantEqual: false,
+		},
+		{
+			desc: "different map keys",
+			a: map[string]builder{
+				"/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+			},
+			b: map[string]builder{
+				"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+			},
+			wantEqual: false,
+		},
+		{
+			desc: "equal keys different values",
+			a: map[string]builder{
+				"/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+				"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1", "n2"}}}},
+			},
+			b: map[string]builder{
+				"/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+				"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+			},
+			wantEqual: false,
+		},
+		{
+			desc: "good match",
+			a: map[string]builder{
+				"/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+				"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+			},
+			b: map[string]builder{
+				"/gBar/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+				"/gFoo/": {headerKeys: []matcher{{key: "k1", names: []string{"n1"}}}},
+			},
+			wantEqual: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			if gotEqual := test.a.Equal(test.b); gotEqual != test.wantEqual {
+				t.Errorf("BuilderMap.Equal(%v, %v) = %v, want %v", test.a, test.b, gotEqual, test.wantEqual)
+			}
+		})
+	}
+}
+
+func TestBuilderEqual(t *testing.T) {
+	tests := []struct {
+		desc      string
+		a         builder
+		b         builder
+		wantEqual bool
+	}{
+		{
+			desc:      "nil builders",
+			a:         builder{headerKeys: nil},
+			b:         builder{headerKeys: nil},
+			wantEqual: true,
+		},
+		{
+			desc:      "empty builders",
+			a:         builder{headerKeys: []matcher{}},
+			b:         builder{headerKeys: []matcher{}},
+			wantEqual: true,
+		},
+		{
+			desc:      "empty and non-empty builders",
+			a:         builder{headerKeys: []matcher{}},
+			b:         builder{headerKeys: []matcher{{key: "foo"}}},
+			wantEqual: false,
+		},
+		{
+			desc:      "different number of headerKeys",
+			a:         builder{headerKeys: []matcher{{key: "foo"}, {key: "bar"}}},
+			b:         builder{headerKeys: []matcher{{key: "foo"}}},
+			wantEqual: false,
+		},
+		{
+			desc:      "equal number but differing headerKeys",
+			a:         builder{headerKeys: []matcher{{key: "bar"}}},
+			b:         builder{headerKeys: []matcher{{key: "foo"}}},
+			wantEqual: false,
+		},
+		{
+			desc:      "different number of constantKeys",
+			a:         builder{constantKeys: map[string]string{"k1": "v1"}},
+			b:         builder{constantKeys: map[string]string{"k1": "v1", "k2": "v2"}},
+			wantEqual: false,
+		},
+		{
+			desc:      "equal number but differing constantKeys",
+			a:         builder{constantKeys: map[string]string{"k1": "v1"}},
+			b:         builder{constantKeys: map[string]string{"k2": "v2"}},
+			wantEqual: false,
+		},
+		{
+			desc:      "different hostKey",
+			a:         builder{hostKey: "host1"},
+			b:         builder{hostKey: "host2"},
+			wantEqual: false,
+		},
+		{
+			desc:      "different serviceKey",
+			a:         builder{hostKey: "service1"},
+			b:         builder{hostKey: "service2"},
+			wantEqual: false,
+		},
+		{
+			desc:      "different methodKey",
+			a:         builder{hostKey: "method1"},
+			b:         builder{hostKey: "method2"},
+			wantEqual: false,
+		},
+		{
+			desc: "equal",
+			a: builder{
+				headerKeys:   []matcher{{key: "foo"}},
+				constantKeys: map[string]string{"k1": "v1"},
+				hostKey:      "host",
+				serviceKey:   "/service/",
+				methodKey:    "method",
+			},
+			b: builder{
+				headerKeys:   []matcher{{key: "foo"}},
+				constantKeys: map[string]string{"k1": "v1"},
+				hostKey:      "host",
+				serviceKey:   "/service/",
+				methodKey:    "method",
+			},
+			wantEqual: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			t.Run(test.desc, func(t *testing.T) {
+				if gotEqual := test.a.Equal(test.b); gotEqual != test.wantEqual {
+					t.Errorf("builder.Equal(%v, %v) = %v, want %v", test.a, test.b, gotEqual, test.wantEqual)
+				}
+			})
+		})
+	}
+}
+
+// matcher helps extract a key from request headers based on a given name.
+func TestMatcherEqual(t *testing.T) {
+	tests := []struct {
+		desc      string
+		a         matcher
+		b         matcher
+		wantEqual bool
+	}{
+		{
+			desc:      "different keys",
+			a:         matcher{key: "foo"},
+			b:         matcher{key: "bar"},
+			wantEqual: false,
+		},
+		{
+			desc:      "different number of names",
+			a:         matcher{key: "foo", names: []string{"v1", "v2"}},
+			b:         matcher{key: "foo", names: []string{"v1"}},
+			wantEqual: false,
+		},
+		{
+			desc:      "equal number but differing names",
+			a:         matcher{key: "foo", names: []string{"v1", "v2"}},
+			b:         matcher{key: "foo", names: []string{"v1", "v22"}},
+			wantEqual: false,
+		},
+		{
+			desc:      "same names in different order",
+			a:         matcher{key: "foo", names: []string{"v2", "v1"}},
+			b:         matcher{key: "foo", names: []string{"v1", "v3"}},
+			wantEqual: false,
+		},
+		{
+			desc:      "good match",
+			a:         matcher{key: "foo", names: []string{"v1", "v2"}},
+			b:         matcher{key: "foo", names: []string{"v1", "v2"}},
+			wantEqual: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			if gotEqual := test.a.Equal(test.b); gotEqual != test.wantEqual {
+				t.Errorf("matcher.Equal(%v, %v) = %v, want %v", test.a, test.b, gotEqual, test.wantEqual)
+			}
+		})
+	}
+}

balancer/rls/internal/test/e2e/e2e.go

@@ -0,0 +1,20 @@
+/*
+ *
+ * Copyright 2021 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 e2e contains utilities for end-to-end RouteLookupService tests.
+package e2e

balancer/rls/internal/test/e2e/rls_child_policy.go

@@ -0,0 +1,131 @@
+/*
+ *
+ * Copyright 2021 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 e2e
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+const (
+	// RLSChildPolicyTargetNameField is a top-level field name to add to the child
+	// policy's config, whose value is set to the target for the child policy.
+	RLSChildPolicyTargetNameField = "Backend"
+	// RLSChildPolicyBadTarget is a value which is considered a bad target by the
+	// child policy. This is useful to test bad child policy configuration.
+	RLSChildPolicyBadTarget = "bad-target"
+)
+
+// ErrParseConfigBadTarget is the error returned from ParseConfig when the
+// backend field is set to RLSChildPolicyBadTarget.
+var ErrParseConfigBadTarget = errors.New("backend field set to RLSChildPolicyBadTarget")
+
+// BalancerFuncs is a set of callbacks which get invoked when the corresponding
+// method on the child policy is invoked.
+type BalancerFuncs struct {
+	UpdateClientConnState func(cfg *RLSChildPolicyConfig) error
+	Close                 func()
+}
+
+// RegisterRLSChildPolicy registers a balancer builder with the given name, to
+// be used as a child policy for the RLS LB policy.
+//
+// The child policy uses a pickfirst balancer under the hood to send all traffic
+// to the single backend specified by the `RLSChildPolicyTargetNameField` field
+// in its configuration which looks like: {"Backend": "Backend-address"}.
+func RegisterRLSChildPolicy(name string, bf *BalancerFuncs) {
+	balancer.Register(bb{name: name, bf: bf})
+}
+
+type bb struct {
+	name string
+	bf   *BalancerFuncs
+}
+
+func (bb bb) Name() string { return bb.name }
+
+func (bb bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	pf := balancer.Get(pickfirst.Name)
+	b := &bal{
+		Balancer: pf.Build(cc, opts),
+		bf:       bb.bf,
+		done:     grpcsync.NewEvent(),
+	}
+	go b.run()
+	return b
+}
+
+func (bb bb) ParseConfig(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	cfg := &RLSChildPolicyConfig{}
+	if err := json.Unmarshal(c, cfg); err != nil {
+		return nil, err
+	}
+	if cfg.Backend == RLSChildPolicyBadTarget {
+		return nil, ErrParseConfigBadTarget
+	}
+	return cfg, nil
+}
+
+type bal struct {
+	balancer.Balancer
+	bf   *BalancerFuncs
+	done *grpcsync.Event
+}
+
+// RLSChildPolicyConfig is the LB config for the test child policy.
+type RLSChildPolicyConfig struct {
+	serviceconfig.LoadBalancingConfig
+	Backend string // The target for which this child policy was created.
+	Random  string // A random field to test child policy config changes.
+}
+
+func (b *bal) UpdateClientConnState(c balancer.ClientConnState) error {
+	cfg, ok := c.BalancerConfig.(*RLSChildPolicyConfig)
+	if !ok {
+		return fmt.Errorf("received balancer config of type %T, want %T", c.BalancerConfig, &RLSChildPolicyConfig{})
+	}
+	if b.bf != nil && b.bf.UpdateClientConnState != nil {
+		b.bf.UpdateClientConnState(cfg)
+	}
+	return b.Balancer.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState: resolver.State{Addresses: []resolver.Address{{Addr: cfg.Backend}}},
+	})
+}
+
+func (b *bal) Close() {
+	b.Balancer.Close()
+	if b.bf != nil && b.bf.Close != nil {
+		b.bf.Close()
+	}
+	b.done.Fire()
+}
+
+// run is a dummy goroutine to make sure that child policies are closed at the
+// end of tests. If they are not closed, these goroutines will be picked up by
+// the leak checker and tests will fail.
+func (b *bal) run() {
+	<-b.done.Done()
+}

balancer/rls/internal/test/e2e/rls_lb_config.go

@@ -0,0 +1,103 @@
+/*
+ *
+ * Copyright 2021 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 e2e
+
+import (
+	"errors"
+	"fmt"
+
+	"google.golang.org/grpc/balancer"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/serviceconfig"
+
+	"google.golang.org/protobuf/encoding/protojson"
+)
+
+// RLSConfig is a utility type to build service config for the RLS LB policy.
+type RLSConfig struct {
+	RouteLookupConfig                *rlspb.RouteLookupConfig
+	RouteLookupChannelServiceConfig  string
+	ChildPolicy                      *internalserviceconfig.BalancerConfig
+	ChildPolicyConfigTargetFieldName string
+}
+
+// ServiceConfigJSON generates service config with a load balancing config
+// corresponding to the RLS LB policy.
+func (c *RLSConfig) ServiceConfigJSON() (string, error) {
+	m := protojson.MarshalOptions{
+		Multiline:     true,
+		Indent:        "  ",
+		UseProtoNames: true,
+	}
+	routeLookupCfg, err := m.Marshal(c.RouteLookupConfig)
+	if err != nil {
+		return "", err
+	}
+	childPolicy, err := c.ChildPolicy.MarshalJSON()
+	if err != nil {
+		return "", err
+	}
+
+	return fmt.Sprintf(`
+{
+  "loadBalancingConfig": [
+    {
+      "rls_experimental": {
+        "routeLookupConfig": %s,
+				"routeLookupChannelServiceConfig": %s,
+        "childPolicy": %s,
+        "childPolicyConfigTargetFieldName": %q
+      }
+    }
+  ]
+}`, string(routeLookupCfg), c.RouteLookupChannelServiceConfig, string(childPolicy), c.ChildPolicyConfigTargetFieldName), nil
+}
+
+// LoadBalancingConfig generates load balancing config which can used as part of
+// a ClientConnState update to the RLS LB policy.
+func (c *RLSConfig) LoadBalancingConfig() (serviceconfig.LoadBalancingConfig, error) {
+	m := protojson.MarshalOptions{
+		Multiline:     true,
+		Indent:        "  ",
+		UseProtoNames: true,
+	}
+	routeLookupCfg, err := m.Marshal(c.RouteLookupConfig)
+	if err != nil {
+		return nil, err
+	}
+	childPolicy, err := c.ChildPolicy.MarshalJSON()
+	if err != nil {
+		return nil, err
+	}
+	lbConfigJSON := fmt.Sprintf(`
+{
+  "routeLookupConfig": %s,
+  "routeLookupChannelServiceConfig": %s,
+  "childPolicy": %s,
+  "childPolicyConfigTargetFieldName": %q
+}`, string(routeLookupCfg), c.RouteLookupChannelServiceConfig, string(childPolicy), c.ChildPolicyConfigTargetFieldName)
+
+	builder := balancer.Get("rls_experimental")
+	if builder == nil {
+		return nil, errors.New("balancer builder not found for RLS LB policy")
+	}
+	parser := builder.(balancer.ConfigParser)
+	return parser.ParseConfig([]byte(lbConfigJSON))
+}

balancer/rls/metrics_test.go

@@ -0,0 +1,375 @@
+/*
+ * 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 rls
+
+import (
+	"context"
+	"math/rand"
+	"testing"
+
+	"github.com/google/uuid"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/sdk/metric"
+	"go.opentelemetry.io/otel/sdk/metric/metricdata"
+	"go.opentelemetry.io/otel/sdk/metric/metricdata/metricdatatest"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	estats "google.golang.org/grpc/experimental/stats"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	"google.golang.org/grpc/internal/stubserver"
+	rlstest "google.golang.org/grpc/internal/testutils/rls"
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+	"google.golang.org/grpc/stats/opentelemetry"
+)
+
+func metricsDataFromReader(ctx context.Context, reader *metric.ManualReader) map[string]metricdata.Metrics {
+	rm := &metricdata.ResourceMetrics{}
+	reader.Collect(ctx, rm)
+	gotMetrics := map[string]metricdata.Metrics{}
+	for _, sm := range rm.ScopeMetrics {
+		for _, m := range sm.Metrics {
+			gotMetrics[m.Name] = m
+		}
+	}
+	return gotMetrics
+}
+
+// TestRLSTargetPickMetric tests RLS Metrics in the case an RLS Balancer picks a
+// target from an RLS Response for a RPC. This should emit a
+// "grpc.lb.rls.target_picks" with certain labels and cache metrics with certain
+// labels.
+func (s) TestRLSTargetPickMetric(t *testing.T) {
+	// Overwrite the uuid random number generator to be deterministic.
+	uuid.SetRand(rand.New(rand.NewSource(1)))
+	defer uuid.SetRand(nil)
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil)
+	rlsConfig := buildBasicRLSConfigWithChildPolicy(t, t.Name(), rlsServer.Address)
+	backend := &stubserver.StubServer{
+		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+			return &testpb.Empty{}, nil
+		},
+	}
+	if err := backend.StartServer(); err != nil {
+		t.Fatalf("Failed to start backend: %v", err)
+	}
+	t.Logf("Started TestService backend at: %q", backend.Address)
+	defer backend.Stop()
+
+	rlsServer.SetResponseCallback(func(context.Context, *rlspb.RouteLookupRequest) *rlstest.RouteLookupResponse {
+		return &rlstest.RouteLookupResponse{Resp: &rlspb.RouteLookupResponse{Targets: []string{backend.Address}}}
+	})
+	r := startManualResolverWithConfig(t, rlsConfig)
+	reader := metric.NewManualReader()
+	provider := metric.NewMeterProvider(metric.WithReader(reader))
+	mo := opentelemetry.MetricsOptions{
+		MeterProvider:  provider,
+		Metrics:        opentelemetry.DefaultMetrics().Add("grpc.lb.rls.cache_entries", "grpc.lb.rls.cache_size", "grpc.lb.rls.default_target_picks", "grpc.lb.rls.target_picks", "grpc.lb.rls.failed_picks"),
+		OptionalLabels: []string{"grpc.client.call.custom"},
+	}
+	grpcTarget := r.Scheme() + ":///"
+	cc, err := grpc.NewClient(grpcTarget, grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()), opentelemetry.DialOption(opentelemetry.Options{MetricsOptions: mo}))
+	if err != nil {
+		t.Fatalf("Failed to dial local test server: %v", err)
+	}
+	defer cc.Close()
+
+	wantMetrics := []metricdata.Metrics{
+		{
+			Name:        "grpc.lb.rls.target_picks",
+			Description: "EXPERIMENTAL. Number of LB picks sent to each RLS target. Note that if the default target is also returned by the RLS server, RPCs sent to that target from the cache will be counted in this metric, not in grpc.rls.default_target_picks.",
+			Unit:        "{pick}",
+			Data: metricdata.Sum[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.lb.rls.data_plane_target", backend.Address), attribute.String("grpc.lb.pick_result", "complete"), attribute.String("grpc.client.call.custom", "target-pick-custom")),
+						Value:      1,
+					},
+				},
+				Temporality: metricdata.CumulativeTemporality,
+				IsMonotonic: true,
+			},
+		},
+
+		// Receives an empty RLS Response, so a single cache entry with no size.
+		{
+			Name:        "grpc.lb.rls.cache_entries",
+			Description: "EXPERIMENTAL. Number of entries in the RLS cache.",
+			Unit:        "{entry}",
+			Data: metricdata.Gauge[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.lb.rls.instance_uuid", "52fdfc07-2182-454f-963f-5f0f9a621d72")),
+						Value:      1,
+					},
+				},
+			},
+		},
+		{
+			Name:        "grpc.lb.rls.cache_size",
+			Description: "EXPERIMENTAL. The current size of the RLS cache.",
+			Unit:        "By",
+			Data: metricdata.Gauge[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.lb.rls.instance_uuid", "52fdfc07-2182-454f-963f-5f0f9a621d72")),
+						Value:      35,
+					},
+				},
+			},
+		},
+	}
+	client := testgrpc.NewTestServiceClient(cc)
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	ctx = estats.NewContextWithCustomLabel(ctx, "target-pick-custom")
+	_, err = client.EmptyCall(ctx, &testpb.Empty{})
+	if err != nil {
+		t.Fatalf("client.EmptyCall failed with error: %v", err)
+	}
+
+	gotMetrics := metricsDataFromReader(ctx, reader)
+	for _, metric := range wantMetrics {
+		val, ok := gotMetrics[metric.Name]
+		if !ok {
+			t.Fatalf("Metric %v not present in recorded metrics", metric.Name)
+		}
+		if !metricdatatest.AssertEqual(t, metric, val, metricdatatest.IgnoreTimestamp(), metricdatatest.IgnoreExemplars()) {
+			t.Fatalf("Metrics data type not equal for metric: %v", metric.Name)
+		}
+	}
+
+	// Only one pick was made, which was a target pick, so no default target
+	// pick or failed pick metric should emit.
+	for _, metric := range []string{"grpc.lb.rls.default_target_picks", "grpc.lb.rls.failed_picks"} {
+		if _, ok := gotMetrics[metric]; ok {
+			t.Fatalf("Metric %v present in recorded metrics", metric)
+		}
+	}
+}
+
+// TestRLSDefaultTargetPickMetric tests RLS Metrics in the case an RLS Balancer
+// falls back to the default target for an RPC. This should emit a
+// "grpc.lb.rls.default_target_picks" with certain labels and cache metrics with
+// certain labels.
+func (s) TestRLSDefaultTargetPickMetric(t *testing.T) {
+	// Overwrite the uuid random number generator to be deterministic.
+	uuid.SetRand(rand.New(rand.NewSource(1)))
+	defer uuid.SetRand(nil)
+
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil)
+	// Build RLS service config with a default target.
+	rlsConfig := buildBasicRLSConfigWithChildPolicy(t, t.Name(), rlsServer.Address)
+	backend := &stubserver.StubServer{
+		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+			return &testpb.Empty{}, nil
+		},
+	}
+	if err := backend.StartServer(); err != nil {
+		t.Fatalf("Failed to start backend: %v", err)
+	}
+	t.Logf("Started TestService backend at: %q", backend.Address)
+	defer backend.Stop()
+	rlsConfig.RouteLookupConfig.DefaultTarget = backend.Address
+
+	r := startManualResolverWithConfig(t, rlsConfig)
+	reader := metric.NewManualReader()
+	provider := metric.NewMeterProvider(metric.WithReader(reader))
+	mo := opentelemetry.MetricsOptions{
+		MeterProvider:  provider,
+		Metrics:        opentelemetry.DefaultMetrics().Add("grpc.lb.rls.cache_entries", "grpc.lb.rls.cache_size", "grpc.lb.rls.default_target_picks", "grpc.lb.rls.target_picks", "grpc.lb.rls.failed_picks"),
+		OptionalLabels: []string{"grpc.client.call.custom"},
+	}
+	grpcTarget := r.Scheme() + ":///"
+	cc, err := grpc.NewClient(grpcTarget, grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()), opentelemetry.DialOption(opentelemetry.Options{MetricsOptions: mo}))
+	if err != nil {
+		t.Fatalf("Failed to dial local test server: %v", err)
+	}
+	defer cc.Close()
+
+	wantMetrics := []metricdata.Metrics{
+		{
+			Name:        "grpc.lb.rls.default_target_picks",
+			Description: "EXPERIMENTAL. Number of LB picks sent to the default target.",
+			Unit:        "{pick}",
+			Data: metricdata.Sum[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.lb.rls.data_plane_target", backend.Address), attribute.String("grpc.lb.pick_result", "complete"), attribute.String("grpc.client.call.custom", "default-target-pick-custom")),
+						Value:      1,
+					},
+				},
+				Temporality: metricdata.CumulativeTemporality,
+				IsMonotonic: true,
+			},
+		},
+		// Receives a RLS Response with target information, so a single cache
+		// entry with a certain size.
+		{
+			Name:        "grpc.lb.rls.cache_entries",
+			Description: "EXPERIMENTAL. Number of entries in the RLS cache.",
+			Unit:        "{entry}",
+			Data: metricdata.Gauge[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.lb.rls.instance_uuid", "52fdfc07-2182-454f-963f-5f0f9a621d72")),
+						Value:      1,
+					},
+				},
+			},
+		},
+		{
+			Name:        "grpc.lb.rls.cache_size",
+			Description: "EXPERIMENTAL. The current size of the RLS cache.",
+			Unit:        "By",
+			Data: metricdata.Gauge[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.lb.rls.instance_uuid", "52fdfc07-2182-454f-963f-5f0f9a621d72")),
+						Value:      0,
+					},
+				},
+			},
+		},
+	}
+	client := testgrpc.NewTestServiceClient(cc)
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	ctx = estats.NewContextWithCustomLabel(ctx, "default-target-pick-custom")
+	if _, err = client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("client.EmptyCall failed with error: %v", err)
+	}
+
+	gotMetrics := metricsDataFromReader(ctx, reader)
+	for _, metric := range wantMetrics {
+		val, ok := gotMetrics[metric.Name]
+		if !ok {
+			t.Fatalf("Metric %v not present in recorded metrics", metric.Name)
+		}
+		if !metricdatatest.AssertEqual(t, metric, val, metricdatatest.IgnoreTimestamp(), metricdatatest.IgnoreExemplars()) {
+			t.Fatalf("Metrics data type not equal for metric: %v", metric.Name)
+		}
+	}
+	// No target picks and failed pick metrics should be emitted, as the test
+	// made only one RPC which recorded as a default target pick.
+	for _, metric := range []string{"grpc.lb.rls.target_picks", "grpc.lb.rls.failed_picks"} {
+		if _, ok := gotMetrics[metric]; ok {
+			t.Fatalf("Metric %v present in recorded metrics", metric)
+		}
+	}
+}
+
+// TestRLSFailedRPCMetric tests RLS Metrics in the case an RLS Balancer fails an
+// RPC due to an RLS failure. This should emit a
+// "grpc.lb.rls.default_target_picks" with certain labels and cache metrics with
+// certain labels.
+func (s) TestRLSFailedRPCMetric(t *testing.T) {
+	// Overwrite the uuid random number generator to be deterministic.
+	uuid.SetRand(rand.New(rand.NewSource(1)))
+	defer uuid.SetRand(nil)
+
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil)
+	// Build an RLS config without a default target.
+	rlsConfig := buildBasicRLSConfigWithChildPolicy(t, t.Name(), rlsServer.Address)
+	// Register a manual resolver and push the RLS service config through it.
+	r := startManualResolverWithConfig(t, rlsConfig)
+	reader := metric.NewManualReader()
+	provider := metric.NewMeterProvider(metric.WithReader(reader))
+	mo := opentelemetry.MetricsOptions{
+		MeterProvider:  provider,
+		Metrics:        opentelemetry.DefaultMetrics().Add("grpc.lb.rls.cache_entries", "grpc.lb.rls.cache_size", "grpc.lb.rls.default_target_picks", "grpc.lb.rls.target_picks", "grpc.lb.rls.failed_picks"),
+		OptionalLabels: []string{"grpc.client.call.custom"},
+	}
+	grpcTarget := r.Scheme() + ":///"
+	cc, err := grpc.NewClient(grpcTarget, grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()), opentelemetry.DialOption(opentelemetry.Options{MetricsOptions: mo}))
+	if err != nil {
+		t.Fatalf("Failed to dial local test server: %v", err)
+	}
+	defer cc.Close()
+
+	wantMetrics := []metricdata.Metrics{
+		{
+			Name:        "grpc.lb.rls.failed_picks",
+			Description: "EXPERIMENTAL. Number of LB picks failed due to either a failed RLS request or the RLS channel being throttled.",
+			Unit:        "{pick}",
+			Data: metricdata.Sum[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.client.call.custom", "failed-pick-custom")),
+						Value:      1,
+					},
+				},
+				Temporality: metricdata.CumulativeTemporality,
+				IsMonotonic: true,
+			},
+		},
+		// Receives an empty RLS Response, so a single cache entry with no size.
+		{
+			Name:        "grpc.lb.rls.cache_entries",
+			Description: "EXPERIMENTAL. Number of entries in the RLS cache.",
+			Unit:        "{entry}",
+			Data: metricdata.Gauge[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.lb.rls.instance_uuid", "52fdfc07-2182-454f-963f-5f0f9a621d72")),
+						Value:      1,
+					},
+				},
+			},
+		},
+		{
+			Name:        "grpc.lb.rls.cache_size",
+			Description: "EXPERIMENTAL. The current size of the RLS cache.",
+			Unit:        "By",
+			Data: metricdata.Gauge[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget), attribute.String("grpc.lb.rls.server_target", rlsServer.Address), attribute.String("grpc.lb.rls.instance_uuid", "52fdfc07-2182-454f-963f-5f0f9a621d72")),
+						Value:      0,
+					},
+				},
+			},
+		},
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	ctx = estats.NewContextWithCustomLabel(ctx, "failed-pick-custom")
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err = client.EmptyCall(ctx, &testpb.Empty{}); err == nil {
+		t.Fatalf("client.EmptyCall error = %v, expected a non nil error", err)
+	}
+
+	gotMetrics := metricsDataFromReader(ctx, reader)
+	for _, metric := range wantMetrics {
+		val, ok := gotMetrics[metric.Name]
+		if !ok {
+			t.Fatalf("Metric %v not present in recorded metrics", metric.Name)
+		}
+		if !metricdatatest.AssertEqual(t, metric, val, metricdatatest.IgnoreTimestamp(), metricdatatest.IgnoreExemplars()) {
+			t.Fatalf("Metrics data type not equal for metric: %v", metric.Name)
+		}
+	}
+	// Only one RPC was made, which was a failed pick due to an RLS failure, so
+	// no metrics for target picks or default target picks should have emitted.
+	for _, metric := range []string{"grpc.lb.rls.target_picks", "grpc.lb.rls.default_target_picks"} {
+		if _, ok := gotMetrics[metric]; ok {
+			t.Fatalf("Metric %v present in recorded metrics", metric)
+		}
+	}
+}

balancer/rls/picker.go

@@ -0,0 +1,401 @@
+/*
+ *
+ * Copyright 2022 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 rls
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+	"sync/atomic"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/rls/internal/keys"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
+	estats "google.golang.org/grpc/experimental/stats"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+var (
+	errRLSThrottled = errors.New("RLS call throttled at client side")
+
+	// Function to compute data cache entry size.
+	computeDataCacheEntrySize = dcEntrySize
+)
+
+// exitIdler wraps the only method on the BalancerGroup that the picker calls.
+type exitIdler interface {
+	ExitIdleOne(id string)
+}
+
+// rlsPicker selects the subConn to be used for a particular RPC. It does not
+// manage subConns directly and delegates to pickers provided by child policies.
+type rlsPicker struct {
+	// The keyBuilder map used to generate RLS keys for the RPC. This is built
+	// by the LB policy based on the received ServiceConfig.
+	kbm keys.BuilderMap
+	// Endpoint from the user's original dial target. Used to set the `host_key`
+	// field in `extra_keys`.
+	origEndpoint string
+
+	lb *rlsBalancer
+
+	// The picker is given its own copy of the below fields from the RLS LB policy
+	// to avoid having to grab the mutex on the latter.
+	rlsServerTarget string
+	grpcTarget      string
+	metricsRecorder estats.MetricsRecorder
+	defaultPolicy   *childPolicyWrapper // Child policy for the default target.
+	ctrlCh          *controlChannel     // Control channel to the RLS server.
+	maxAge          time.Duration       // Cache max age from LB config.
+	staleAge        time.Duration       // Cache stale age from LB config.
+	bg              exitIdler
+	logger          *internalgrpclog.PrefixLogger
+}
+
+// isFullMethodNameValid return true if name is of the form `/service/method`.
+func isFullMethodNameValid(name string) bool {
+	return strings.HasPrefix(name, "/") && strings.Count(name, "/") == 2
+}
+
+// Pick makes the routing decision for every outbound RPC.
+func (p *rlsPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	if name := info.FullMethodName; !isFullMethodNameValid(name) {
+		return balancer.PickResult{}, fmt.Errorf("rls: method name %q is not of the form '/service/method", name)
+	}
+
+	// Build the request's keys using the key builders from LB config.
+	md, _ := metadata.FromOutgoingContext(info.Ctx)
+	reqKeys := p.kbm.RLSKey(md, p.origEndpoint, info.FullMethodName)
+
+	p.lb.cacheMu.Lock()
+	var pr balancer.PickResult
+	var err error
+
+	// Record metrics without the cache mutex held, to prevent lock contention
+	// between concurrent RPC's and their Pick calls. Metrics Recording can
+	// potentially be expensive.
+	metricsCallback := func() {}
+	defer func() {
+		p.lb.cacheMu.Unlock()
+		metricsCallback()
+	}()
+
+	// Lookup data cache and pending request map using request path and keys.
+	cacheKey := cacheKey{path: info.FullMethodName, keys: reqKeys.Str}
+	dcEntry := p.lb.dataCache.getEntry(cacheKey)
+	pendingEntry := p.lb.pendingMap[cacheKey]
+	now := time.Now()
+
+	switch {
+	// No data cache entry. No pending request.
+	case dcEntry == nil && pendingEntry == nil:
+		throttled := p.sendRouteLookupRequestLocked(cacheKey, &backoffState{bs: defaultBackoffStrategy}, reqKeys.Map, rlspb.RouteLookupRequest_REASON_MISS, "")
+		if throttled {
+			pr, metricsCallback, err = p.useDefaultPickIfPossible(info, errRLSThrottled)
+			return pr, err
+		}
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+
+	// No data cache entry. Pending request exits.
+	case dcEntry == nil && pendingEntry != nil:
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+
+	// Data cache hit. No pending request.
+	case dcEntry != nil && pendingEntry == nil:
+		if dcEntry.expiryTime.After(now) {
+			if !dcEntry.staleTime.IsZero() && dcEntry.staleTime.Before(now) && dcEntry.backoffTime.Before(now) {
+				p.sendRouteLookupRequestLocked(cacheKey, dcEntry.backoffState, reqKeys.Map, rlspb.RouteLookupRequest_REASON_STALE, dcEntry.headerData)
+			}
+			// Delegate to child policies.
+			pr, metricsCallback, err = p.delegateToChildPoliciesLocked(dcEntry, info)
+			return pr, err
+		}
+
+		// We get here only if the data cache entry has expired. If entry is in
+		// backoff, delegate to default target or fail the pick.
+		if dcEntry.backoffState != nil && dcEntry.backoffTime.After(now) {
+			// Avoid propagating the status code received on control plane RPCs to the
+			// data plane which can lead to unexpected outcomes as we do not control
+			// the status code sent by the control plane. Propagating the status
+			// message received from the control plane is still fine, as it could be
+			// useful for debugging purposes.
+			st := dcEntry.status
+			pr, metricsCallback, err = p.useDefaultPickIfPossible(info, status.Error(codes.Unavailable, fmt.Sprintf("most recent error from RLS server: %v", st.Error())))
+			return pr, err
+		}
+
+		// We get here only if the entry has expired and is not in backoff.
+		throttled := p.sendRouteLookupRequestLocked(cacheKey, dcEntry.backoffState, reqKeys.Map, rlspb.RouteLookupRequest_REASON_MISS, "")
+		if throttled {
+			pr, metricsCallback, err = p.useDefaultPickIfPossible(info, errRLSThrottled)
+			return pr, err
+		}
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+
+	// Data cache hit. Pending request exists.
+	default:
+		if dcEntry.expiryTime.After(now) {
+			pr, metricsCallback, err = p.delegateToChildPoliciesLocked(dcEntry, info)
+			return pr, err
+		}
+		// Data cache entry has expired and pending request exists. Queue pick.
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+	}
+}
+
+// errToPickResult is a helper function which converts the error value returned
+// by Pick() to a string that represents the pick result.
+func errToPickResult(err error) string {
+	if err == nil {
+		return "complete"
+	}
+	if errors.Is(err, balancer.ErrNoSubConnAvailable) {
+		return "queue"
+	}
+	if _, ok := status.FromError(err); ok {
+		return "drop"
+	}
+	return "fail"
+}
+
+// delegateToChildPoliciesLocked is a helper function which iterates through the
+// list of child policy wrappers in a cache entry and attempts to find a child
+// policy to which this RPC can be routed to. If all child policies are in
+// TRANSIENT_FAILURE, we delegate to the last child policy arbitrarily. Returns
+// a function to be invoked to record metrics.
+func (p *rlsPicker) delegateToChildPoliciesLocked(dcEntry *cacheEntry, info balancer.PickInfo) (balancer.PickResult, func(), error) {
+	const rlsDataHeaderName = "x-google-rls-data"
+	for i, cpw := range dcEntry.childPolicyWrappers {
+		state := (*balancer.State)(atomic.LoadPointer(&cpw.state))
+		// Delegate to the child policy if it is not in TRANSIENT_FAILURE, or if
+		// it is the last one (which handles the case of delegating to the last
+		// child picker if all child policies are in TRANSIENT_FAILURE).
+		if state.ConnectivityState != connectivity.TransientFailure || i == len(dcEntry.childPolicyWrappers)-1 {
+			// Any header data received from the RLS server is stored in the
+			// cache entry and needs to be sent to the actual backend in the
+			// X-Google-RLS-Data header.
+			res, err := state.Picker.Pick(info)
+			if err != nil {
+				pr := errToPickResult(err)
+				customLabel := estats.CustomLabelFromContext(info.Ctx)
+				return res, func() {
+					if pr == "queue" {
+						// Don't record metrics for queued Picks.
+						return
+					}
+					targetPicksMetric.Record(p.metricsRecorder, 1, p.grpcTarget, p.rlsServerTarget, cpw.target, pr, customLabel)
+				}, err
+			}
+
+			if res.Metadata == nil {
+				res.Metadata = metadata.Pairs(rlsDataHeaderName, dcEntry.headerData)
+			} else {
+				res.Metadata.Append(rlsDataHeaderName, dcEntry.headerData)
+			}
+			return res, func() {
+				customLabel := estats.CustomLabelFromContext(info.Ctx)
+				targetPicksMetric.Record(p.metricsRecorder, 1, p.grpcTarget, p.rlsServerTarget, cpw.target, "complete", customLabel)
+			}, nil
+		}
+	}
+
+	// In the unlikely event that we have a cache entry with no targets, we end up
+	// queueing the RPC.
+	return balancer.PickResult{}, func() {}, balancer.ErrNoSubConnAvailable
+}
+
+// useDefaultPickIfPossible is a helper method which delegates to the default
+// target if one is configured, or fails the pick with the given error. Returns
+// a function to be invoked to record metrics.
+func (p *rlsPicker) useDefaultPickIfPossible(info balancer.PickInfo, errOnNoDefault error) (balancer.PickResult, func(), error) {
+	customLabel := estats.CustomLabelFromContext(info.Ctx)
+
+	if p.defaultPolicy != nil {
+		state := (*balancer.State)(atomic.LoadPointer(&p.defaultPolicy.state))
+		res, err := state.Picker.Pick(info)
+		pr := errToPickResult(err)
+		return res, func() {
+			if pr == "queue" {
+				// Don't record metrics for queued Picks.
+				return
+			}
+			defaultTargetPicksMetric.Record(p.metricsRecorder, 1, p.grpcTarget, p.rlsServerTarget, p.defaultPolicy.target, pr, customLabel)
+		}, err
+	}
+
+	return balancer.PickResult{}, func() {
+		failedPicksMetric.Record(p.metricsRecorder, 1, p.grpcTarget, p.rlsServerTarget, customLabel)
+	}, errOnNoDefault
+}
+
+// sendRouteLookupRequestLocked adds an entry to the pending request map and
+// sends out an RLS request using the passed in arguments. Returns a value
+// indicating if the request was throttled by the client-side adaptive
+// throttler.
+func (p *rlsPicker) sendRouteLookupRequestLocked(cacheKey cacheKey, bs *backoffState, reqKeys map[string]string, reason rlspb.RouteLookupRequest_Reason, staleHeaders string) bool {
+	if p.lb.pendingMap[cacheKey] != nil {
+		return false
+	}
+
+	p.lb.pendingMap[cacheKey] = bs
+	throttled := p.ctrlCh.lookup(reqKeys, reason, staleHeaders, func(targets []string, headerData string, err error) {
+		p.handleRouteLookupResponse(cacheKey, targets, headerData, err)
+	})
+	if throttled {
+		delete(p.lb.pendingMap, cacheKey)
+	}
+	return throttled
+}
+
+// handleRouteLookupResponse is the callback invoked by the control channel upon
+// receipt of an RLS response. Modifies the data cache and pending requests map
+// and sends a new picker.
+//
+// Acquires the write-lock on the cache. Caller must not hold p.lb.cacheMu.
+func (p *rlsPicker) handleRouteLookupResponse(cacheKey cacheKey, targets []string, headerData string, err error) {
+	p.logger.Infof("Received RLS response for key %+v with targets %+v, headerData %q, err: %v", cacheKey, targets, headerData, err)
+
+	p.lb.cacheMu.Lock()
+	defer func() {
+		// Pending request map entry is unconditionally deleted since the request is
+		// no longer pending.
+		p.logger.Infof("Removing pending request entry for key %+v", cacheKey)
+		delete(p.lb.pendingMap, cacheKey)
+		p.lb.sendNewPicker()
+		p.lb.cacheMu.Unlock()
+	}()
+
+	// Lookup the data cache entry or create a new one.
+	dcEntry := p.lb.dataCache.getEntry(cacheKey)
+	if dcEntry == nil {
+		dcEntry = &cacheEntry{}
+		if _, ok := p.lb.dataCache.addEntry(cacheKey, dcEntry); !ok {
+			// This is a very unlikely case where we are unable to add a
+			// data cache entry. Log and leave.
+			p.logger.Warningf("Failed to add data cache entry for %+v", cacheKey)
+			return
+		}
+	}
+
+	// For failed requests, the data cache entry is modified as follows:
+	// - status is set to error returned from the control channel
+	// - current backoff state is available in the pending entry
+	//   - `retries` field is incremented and
+	//   - backoff state is moved to the data cache
+	// - backoffTime is set to the time indicated by the backoff state
+	// - backoffExpirationTime is set to twice the backoff time
+	// - backoffTimer is set to fire after backoffTime
+	//
+	// When a proactive cache refresh fails, this would leave the targets and the
+	// expiry time from the old entry unchanged. And this mean that the old valid
+	// entry would be used until expiration, and a new picker would be sent upon
+	// backoff expiry.
+	now := time.Now()
+
+	// "An RLS request is considered to have failed if it returns a non-OK
+	// status or the RLS response's targets list is non-empty." - RLS LB Policy
+	// design.
+	if len(targets) == 0 && err == nil {
+		err = fmt.Errorf("RLS response's target list does not contain any entries for key %+v", cacheKey)
+		// If err is set, rpc error from the control plane and no control plane
+		// configuration is why no targets were passed into this helper, no need
+		// to specify and tell the user this information.
+	}
+	if err != nil {
+		dcEntry.status = err
+		pendingEntry := p.lb.pendingMap[cacheKey]
+		pendingEntry.retries++
+		backoffTime := pendingEntry.bs.Backoff(pendingEntry.retries)
+		dcEntry.backoffState = pendingEntry
+		dcEntry.backoffTime = now.Add(backoffTime)
+		dcEntry.backoffExpiryTime = now.Add(2 * backoffTime)
+		if dcEntry.backoffState.timer != nil {
+			dcEntry.backoffState.timer.Stop()
+		}
+		dcEntry.backoffState.timer = time.AfterFunc(backoffTime, p.lb.sendNewPicker)
+		return
+	}
+
+	// For successful requests, the cache entry is modified as follows:
+	// - childPolicyWrappers is set to point to the child policy wrappers
+	//   associated with the targets specified in the received response
+	// - headerData is set to the value received in the response
+	// - expiryTime, stateTime and earliestEvictionTime are set
+	// - status is set to nil (OK status)
+	// - backoff state is cleared
+	p.setChildPolicyWrappersInCacheEntry(dcEntry, targets)
+	dcEntry.headerData = headerData
+	dcEntry.expiryTime = now.Add(p.maxAge)
+	if p.staleAge != 0 {
+		dcEntry.staleTime = now.Add(p.staleAge)
+	}
+	dcEntry.earliestEvictTime = now.Add(minEvictDuration)
+	dcEntry.status = nil
+	dcEntry.backoffState = &backoffState{bs: defaultBackoffStrategy}
+	dcEntry.backoffTime = time.Time{}
+	dcEntry.backoffExpiryTime = time.Time{}
+	p.lb.dataCache.updateEntrySize(dcEntry, computeDataCacheEntrySize(cacheKey, dcEntry))
+}
+
+// setChildPolicyWrappersInCacheEntry sets up the childPolicyWrappers field in
+// the cache entry to point to the child policy wrappers for the targets
+// specified in the RLS response.
+//
+// Caller must hold a write-lock on p.lb.cacheMu.
+func (p *rlsPicker) setChildPolicyWrappersInCacheEntry(dcEntry *cacheEntry, newTargets []string) {
+	// If the childPolicyWrappers field is already pointing to the right targets,
+	// then the field's value does not need to change.
+	targetsChanged := true
+	func() {
+		if cpws := dcEntry.childPolicyWrappers; cpws != nil {
+			if len(newTargets) != len(cpws) {
+				return
+			}
+			for i, target := range newTargets {
+				if cpws[i].target != target {
+					return
+				}
+			}
+			targetsChanged = false
+		}
+	}()
+	if !targetsChanged {
+		return
+	}
+
+	// If the childPolicyWrappers field is not already set to the right targets,
+	// then it must be reset. We construct a new list of child policies and
+	// then swap out the old list for the new one.
+	newChildPolicies := p.lb.acquireChildPolicyReferences(newTargets)
+	oldChildPolicyTargets := make([]string, len(dcEntry.childPolicyWrappers))
+	for i, cpw := range dcEntry.childPolicyWrappers {
+		oldChildPolicyTargets[i] = cpw.target
+	}
+	p.lb.releaseChildPolicyReferences(oldChildPolicyTargets)
+	dcEntry.childPolicyWrappers = newChildPolicies
+}
+
+func dcEntrySize(key cacheKey, entry *cacheEntry) int64 {
+	return int64(len(key.path) + len(key.keys) + len(entry.headerData))
+}

balancer/roundrobin/roundrobin.go

@@ -0,0 +1,72 @@
+/*
+ *
+ * Copyright 2017 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 roundrobin defines a roundrobin balancer. Roundrobin balancer is
+// installed as one of the default balancers in gRPC, users don't need to
+// explicitly install this balancer.
+package roundrobin
+
+import (
+	"fmt"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/grpclog"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+)
+
+// Name is the name of round_robin balancer.
+const Name = "round_robin"
+
+var logger = grpclog.Component("roundrobin")
+
+func init() {
+	balancer.Register(builder{})
+}
+
+type builder struct{}
+
+func (bb builder) Name() string {
+	return Name
+}
+
+func (bb builder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	childBuilder := balancer.Get(pickfirst.Name).Build
+	bal := &rrBalancer{
+		cc:       cc,
+		Balancer: endpointsharding.NewBalancer(cc, opts, childBuilder, endpointsharding.Options{}),
+	}
+	bal.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[%p] ", bal))
+	bal.logger.Infof("Created")
+	return bal
+}
+
+type rrBalancer struct {
+	balancer.Balancer
+	cc     balancer.ClientConn
+	logger *internalgrpclog.PrefixLogger
+}
+
+func (b *rrBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	return b.Balancer.UpdateClientConnState(balancer.ClientConnState{
+		// Enable the health listener in pickfirst children for client side health
+		// checks and outlier detection, if configured.
+		ResolverState: pickfirst.EnableHealthListener(ccs.ResolverState),
+	})
+}

balancer/subconn.go

@@ -0,0 +1,120 @@
+/*
+ *
+ * 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 balancer
+
+import (
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/resolver"
+)
+
+// A SubConn represents a single connection to a gRPC backend service.
+//
+// All SubConns start in IDLE, and will not try to connect. To trigger a
+// connection attempt, Balancers must call Connect.
+//
+// If the connection attempt fails, the SubConn will transition to
+// TRANSIENT_FAILURE for a backoff period, and then return to IDLE.  If the
+// connection attempt succeeds, it will transition to READY.
+//
+// If a READY SubConn becomes disconnected, the SubConn will transition to IDLE.
+//
+// If a connection re-enters IDLE, Balancers must call Connect again to trigger
+// a new connection attempt.
+//
+// Each SubConn contains a list of addresses.  gRPC will try to connect to the
+// addresses in sequence, and stop trying the remainder once the first
+// connection is successful.  However, this behavior is deprecated.  SubConns
+// should only use a single address.
+//
+// NOTICE: This interface is intended to be implemented by gRPC, or intercepted
+// by custom load balancing polices.  Users should not need their own complete
+// implementation of this interface -- they should always delegate to a SubConn
+// returned by ClientConn.NewSubConn() by embedding it in their implementations.
+// An embedded SubConn must never be nil, or runtime panics will occur.
+type SubConn interface {
+	// UpdateAddresses updates the addresses used in this SubConn.
+	// gRPC checks if currently-connected address is still in the new list.
+	// If it's in the list, the connection will be kept.
+	// If it's not in the list, the connection will gracefully close, and
+	// a new connection will be created.
+	//
+	// This will trigger a state transition for the SubConn.
+	//
+	// Deprecated: this method will be removed.  Create new SubConns for new
+	// addresses instead.
+	UpdateAddresses([]resolver.Address)
+	// Connect starts the connecting for this SubConn.
+	Connect()
+	// GetOrBuildProducer returns a reference to the existing Producer for this
+	// ProducerBuilder in this SubConn, or, if one does not currently exist,
+	// creates a new one and returns it.  Returns a close function which may be
+	// called when the Producer is no longer needed.  Otherwise the producer
+	// will automatically be closed upon connection loss or subchannel close.
+	// Should only be called on a SubConn in state Ready.  Otherwise the
+	// producer will be unable to create streams.
+	GetOrBuildProducer(ProducerBuilder) (p Producer, close func())
+	// Shutdown shuts down the SubConn gracefully.  Any started RPCs will be
+	// allowed to complete.  No future calls should be made on the SubConn.
+	// One final state update will be delivered to the StateListener (or
+	// UpdateSubConnState; deprecated) with ConnectivityState of Shutdown to
+	// indicate the shutdown operation.  This may be delivered before
+	// in-progress RPCs are complete and the actual connection is closed.
+	Shutdown()
+	// RegisterHealthListener registers a health listener that receives health
+	// updates for a Ready SubConn. Only one health listener can be registered
+	// at a time. A health listener should be registered each time the SubConn's
+	// connectivity state changes to READY. Registering a health listener when
+	// the connectivity state is not READY may result in undefined behaviour.
+	// This method must not be called synchronously while handling an update
+	// from a previously registered health listener.
+	RegisterHealthListener(func(SubConnState))
+	// EnforceSubConnEmbedding is included to force implementers to embed
+	// another implementation of this interface, allowing gRPC to add methods
+	// without breaking users.
+	internal.EnforceSubConnEmbedding
+}
+
+// A ProducerBuilder is a simple constructor for a Producer.  It is used by the
+// SubConn to create producers when needed.
+type ProducerBuilder interface {
+	// Build creates a Producer.  The first parameter is always a
+	// grpc.ClientConnInterface (a type to allow creating RPCs/streams on the
+	// associated SubConn), but is declared as `any` to avoid a dependency
+	// cycle.  Build also returns a close function that will be called when all
+	// references to the Producer have been given up for a SubConn, or when a
+	// connectivity state change occurs on the SubConn.  The close function
+	// should always block until all asynchronous cleanup work is completed.
+	Build(grpcClientConnInterface any) (p Producer, close func())
+}
+
+// SubConnState describes the state of a SubConn.
+type SubConnState struct {
+	// ConnectivityState is the connectivity state of the SubConn.
+	ConnectivityState connectivity.State
+	// ConnectionError is set if the ConnectivityState is TransientFailure,
+	// describing the reason the SubConn failed.  Otherwise, it is nil.
+	ConnectionError error
+}
+
+// A Producer is a type shared among potentially many consumers.  It is
+// associated with a SubConn, and an implementation will typically contain
+// other methods to provide additional functionality, e.g. configuration or
+// subscription registration.
+type Producer any

balancer/weightedroundrobin/balancer.go

@@ -0,0 +1,674 @@
+/*
+ *
+ * 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 weightedroundrobin provides an implementation of the weighted round
+// robin LB policy, as defined in [gRFC A58].
+//
+// # Experimental
+//
+// Notice: This package is EXPERIMENTAL and may be changed or removed in a
+// later release.
+//
+// [gRFC A58]: https://github.com/grpc/proposal/blob/master/A58-client-side-weighted-round-robin-lb-policy.md
+package weightedroundrobin
+
+import (
+	"encoding/json"
+	"fmt"
+	rand "math/rand/v2"
+	"sync"
+	"sync/atomic"
+	"time"
+	"unsafe"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/balancer/weightedroundrobin/internal"
+	"google.golang.org/grpc/balancer/weightedtarget"
+	"google.golang.org/grpc/connectivity"
+	estats "google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/grpcsync"
+	iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/orca"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+
+	v3orcapb "github.com/cncf/xds/go/xds/data/orca/v3"
+)
+
+// Name is the name of the weighted round robin balancer.
+const Name = "weighted_round_robin"
+
+var (
+	rrFallbackMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:           "grpc.lb.wrr.rr_fallback",
+		Description:    "EXPERIMENTAL. Number of scheduler updates in which there were not enough endpoints with valid weight, which caused the WRR policy to fall back to RR behavior.",
+		Unit:           "{update}",
+		Labels:         []string{"grpc.target"},
+		OptionalLabels: []string{"grpc.lb.locality", "grpc.lb.backend_service"},
+		Default:        false,
+	})
+
+	endpointWeightNotYetUsableMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:           "grpc.lb.wrr.endpoint_weight_not_yet_usable",
+		Description:    "EXPERIMENTAL. Number of endpoints from each scheduler update that don't yet have usable weight information (i.e., either the load report has not yet been received, or it is within the blackout period).",
+		Unit:           "{endpoint}",
+		Labels:         []string{"grpc.target"},
+		OptionalLabels: []string{"grpc.lb.locality", "grpc.lb.backend_service"},
+		Default:        false,
+	})
+
+	endpointWeightStaleMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:           "grpc.lb.wrr.endpoint_weight_stale",
+		Description:    "EXPERIMENTAL. Number of endpoints from each scheduler update whose latest weight is older than the expiration period.",
+		Unit:           "{endpoint}",
+		Labels:         []string{"grpc.target"},
+		OptionalLabels: []string{"grpc.lb.locality", "grpc.lb.backend_service"},
+		Default:        false,
+	})
+	endpointWeightsMetric = estats.RegisterFloat64Histo(estats.MetricDescriptor{
+		Name:           "grpc.lb.wrr.endpoint_weights",
+		Description:    "EXPERIMENTAL. Weight of each endpoint, recorded on every scheduler update. Endpoints without usable weights will be recorded as weight 0.",
+		Unit:           "{endpoint}",
+		Labels:         []string{"grpc.target"},
+		OptionalLabels: []string{"grpc.lb.locality", "grpc.lb.backend_service"},
+		Default:        false,
+	})
+)
+
+func init() {
+	balancer.Register(bb{})
+}
+
+type bb struct{}
+
+func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+	b := &wrrBalancer{
+		ClientConn:       cc,
+		target:           bOpts.Target.String(),
+		metricsRecorder:  cc.MetricsRecorder(),
+		addressWeights:   resolver.NewAddressMapV2[*endpointWeight](),
+		endpointToWeight: resolver.NewEndpointMap[*endpointWeight](),
+		scToWeight:       make(map[balancer.SubConn]*endpointWeight),
+	}
+
+	b.child = endpointsharding.NewBalancer(b, bOpts, balancer.Get(pickfirst.Name).Build, endpointsharding.Options{})
+	b.logger = prefixLogger(b)
+	if b.logger.V(2) {
+		b.logger.Infof("Created")
+	}
+	return b
+}
+
+func (bb) ParseConfig(js json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	lbCfg := &lbConfig{
+		// Default values as documented in A58.
+		OOBReportingPeriod:      iserviceconfig.Duration(10 * time.Second),
+		BlackoutPeriod:          iserviceconfig.Duration(10 * time.Second),
+		WeightExpirationPeriod:  iserviceconfig.Duration(3 * time.Minute),
+		WeightUpdatePeriod:      iserviceconfig.Duration(time.Second),
+		ErrorUtilizationPenalty: 1,
+	}
+	if err := json.Unmarshal(js, lbCfg); err != nil {
+		return nil, fmt.Errorf("wrr: unable to unmarshal LB policy config: %s, error: %v", string(js), err)
+	}
+
+	if lbCfg.ErrorUtilizationPenalty < 0 {
+		return nil, fmt.Errorf("wrr: errorUtilizationPenalty must be non-negative")
+	}
+
+	// For easier comparisons later, ensure the OOB reporting period is unset
+	// (0s) when OOB reports are disabled.
+	if !lbCfg.EnableOOBLoadReport {
+		lbCfg.OOBReportingPeriod = 0
+	}
+
+	// Impose lower bound of 100ms on weightUpdatePeriod.
+	if !internal.AllowAnyWeightUpdatePeriod && lbCfg.WeightUpdatePeriod < iserviceconfig.Duration(100*time.Millisecond) {
+		lbCfg.WeightUpdatePeriod = iserviceconfig.Duration(100 * time.Millisecond)
+	}
+
+	return lbCfg, nil
+}
+
+func (bb) Name() string {
+	return Name
+}
+
+// updateEndpointsLocked updates endpoint weight state based off new update, by
+// starting and clearing any endpoint weights needed.
+//
+// Caller must hold b.mu.
+func (b *wrrBalancer) updateEndpointsLocked(endpoints []resolver.Endpoint) {
+	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)
+		}
+		ew, ok := b.endpointToWeight.Get(endpoint)
+		if !ok {
+			ew = &endpointWeight{
+				logger:            b.logger,
+				connectivityState: connectivity.Connecting,
+				// Initially, we set load reports to off, because they are not
+				// running upon initial endpointWeight creation.
+				cfg:             &lbConfig{EnableOOBLoadReport: false},
+				metricsRecorder: b.metricsRecorder,
+				target:          b.target,
+				locality:        b.locality,
+				clusterName:     b.clusterName,
+			}
+			for _, addr := range endpoint.Addresses {
+				b.addressWeights.Set(addr, ew)
+			}
+			b.endpointToWeight.Set(endpoint, ew)
+		}
+		ew.updateConfig(b.cfg)
+	}
+
+	for endpoint := range b.endpointToWeight.All() {
+		if _, ok := endpointSet.Get(endpoint); ok {
+			// Existing endpoint also in new endpoint list; skip.
+			continue
+		}
+		b.endpointToWeight.Delete(endpoint)
+		for _, addr := range endpoint.Addresses {
+			if _, ok := addressSet.Get(addr); !ok { // old endpoints to be deleted can share addresses with new endpoints, so only delete if necessary
+				b.addressWeights.Delete(addr)
+			}
+		}
+		// SubConn map will get handled in updateSubConnState
+		// when receives SHUTDOWN signal.
+	}
+}
+
+// wrrBalancer implements the weighted round robin LB policy.
+type wrrBalancer struct {
+	// The following fields are set at initialization time and read only after that,
+	// so they do not need to be protected by a mutex.
+	child               balancer.Balancer
+	balancer.ClientConn // Embed to intercept NewSubConn operation
+	logger              *grpclog.PrefixLogger
+	target              string
+	metricsRecorder     estats.MetricsRecorder
+
+	mu               sync.Mutex
+	cfg              *lbConfig // active config
+	locality         string
+	clusterName      string
+	stopPicker       *grpcsync.Event
+	addressWeights   *resolver.AddressMapV2[*endpointWeight]
+	endpointToWeight *resolver.EndpointMap[*endpointWeight]
+	scToWeight       map[balancer.SubConn]*endpointWeight
+}
+
+func (b *wrrBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	if b.logger.V(2) {
+		b.logger.Infof("Received update from resolver with state: %+v", ccs)
+	}
+	cfg, ok := ccs.BalancerConfig.(*lbConfig)
+	if !ok {
+		return fmt.Errorf("wrr: received nil or illegal BalancerConfig (type %T): %v", ccs.BalancerConfig, ccs.BalancerConfig)
+	}
+
+	// Note: empty endpoints and duplicate addresses across endpoints won't
+	// explicitly error but will have undefined behavior.
+	b.mu.Lock()
+	b.cfg = cfg
+	b.locality = weightedtarget.LocalityFromResolverState(ccs.ResolverState)
+	b.clusterName = backendServiceFromState(ccs.ResolverState)
+	b.updateEndpointsLocked(ccs.ResolverState.Endpoints)
+	b.mu.Unlock()
+
+	// This causes child to update picker inline and will thus cause inline
+	// picker update.
+	return b.child.UpdateClientConnState(balancer.ClientConnState{
+		// Make pickfirst children use health listeners for outlier detection to
+		// work.
+		ResolverState: pickfirst.EnableHealthListener(ccs.ResolverState),
+	})
+}
+
+func (b *wrrBalancer) UpdateState(state balancer.State) {
+	if b.logger.V(2) {
+		b.logger.Infof("Received update from child policy with state: %+v", state)
+	}
+
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if b.stopPicker != nil {
+		b.stopPicker.Fire()
+		b.stopPicker = nil
+	}
+
+	childStates := endpointsharding.ChildStatesFromPicker(state.Picker)
+
+	var readyPickersWeight []pickerWeightedEndpoint
+
+	for _, childState := range childStates {
+		if childState.State.ConnectivityState == connectivity.Ready {
+			ew, ok := b.endpointToWeight.Get(childState.Endpoint)
+			if !ok {
+				// Should never happen, simply continue and ignore this endpoint
+				// for READY pickers.
+				continue
+			}
+			readyPickersWeight = append(readyPickersWeight, pickerWeightedEndpoint{
+				picker:           childState.State.Picker,
+				weightedEndpoint: ew,
+			})
+		}
+	}
+	// If no ready pickers are present, simply defer to the round robin picker
+	// from endpoint sharding, which will round robin across the most relevant
+	// pick first children in the highest precedence connectivity state.
+	if len(readyPickersWeight) == 0 {
+		b.ClientConn.UpdateState(balancer.State{
+			ConnectivityState: state.ConnectivityState,
+			Picker:            state.Picker,
+		})
+		return
+	}
+
+	p := &picker{
+		cfg:             b.cfg,
+		weightedPickers: readyPickersWeight,
+		metricsRecorder: b.metricsRecorder,
+		locality:        b.locality,
+		target:          b.target,
+		clusterName:     b.clusterName,
+	}
+	p.idx.Store(rand.Uint32()) // start the scheduler at a random point
+
+	b.stopPicker = grpcsync.NewEvent()
+	p.start(b.stopPicker)
+
+	b.ClientConn.UpdateState(balancer.State{
+		ConnectivityState: state.ConnectivityState,
+		Picker:            p,
+	})
+}
+
+type pickerWeightedEndpoint struct {
+	picker           balancer.Picker
+	weightedEndpoint *endpointWeight
+}
+
+func (b *wrrBalancer) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	addr := addrs[0] // The new pick first policy for DualStack will only ever create a SubConn with one address.
+	var sc balancer.SubConn
+
+	oldListener := opts.StateListener
+	opts.StateListener = func(state balancer.SubConnState) {
+		b.updateSubConnState(sc, state)
+		oldListener(state)
+	}
+
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	ewi, ok := b.addressWeights.Get(addr)
+	if !ok {
+		// SubConn state updates can come in for a no longer relevant endpoint
+		// weight (from the old system after a new config update is applied).
+		return nil, fmt.Errorf("wrr: balancer is being closed; no new SubConns allowed")
+	}
+	sc, err := b.ClientConn.NewSubConn([]resolver.Address{addr}, opts)
+	if err != nil {
+		return nil, err
+	}
+	b.scToWeight[sc] = ewi
+	return sc, nil
+}
+
+func (b *wrrBalancer) ResolverError(err error) {
+	if b.logger.V(2) {
+		b.logger.Infof("Received error from resolver: %v", err)
+	}
+	// Will cause inline picker update from endpoint sharding.
+	b.child.ResolverError(err)
+}
+
+func (b *wrrBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+func (b *wrrBalancer) updateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	if b.logger.V(2) {
+		b.logger.Infof("Received update from SubConn %v with state: %+v", sc, state)
+	}
+
+	b.mu.Lock()
+	ew := b.scToWeight[sc]
+	// updates from a no longer relevant SubConn update, nothing to do here but
+	// forward state to state listener, which happens in wrapped listener. Will
+	// eventually get cleared from scMap once receives Shutdown signal.
+	if ew == nil {
+		b.mu.Unlock()
+		return
+	}
+	if state.ConnectivityState == connectivity.Shutdown {
+		delete(b.scToWeight, sc)
+	}
+	b.mu.Unlock()
+
+	// On the first READY SubConn/Transition for an endpoint, set pickedSC,
+	// clear endpoint tracking weight state, and potentially start an OOB watch.
+	if state.ConnectivityState == connectivity.Ready && ew.pickedSC == nil {
+		ew.pickedSC = sc
+		ew.mu.Lock()
+		ew.nonEmptySince = time.Time{}
+		ew.lastUpdated = time.Time{}
+		cfg := ew.cfg
+		ew.mu.Unlock()
+		ew.updateORCAListener(cfg)
+		return
+	}
+
+	// If the pickedSC (the one pick first uses for an endpoint) transitions out
+	// of READY, stop OOB listener if needed and clear pickedSC so the next
+	// created SubConn for the endpoint that goes READY will be chosen for
+	// endpoint as the active SubConn.
+	if state.ConnectivityState != connectivity.Ready && ew.pickedSC == sc {
+		// The first SubConn that goes READY for an endpoint is what pick first
+		// will pick. Only once that SubConn goes not ready will pick first
+		// restart this cycle of creating SubConns and using the first READY
+		// one. The lower level endpoint sharding will ping the Pick First once
+		// this occurs to ExitIdle which will trigger a connection attempt.
+		if ew.stopORCAListener != nil {
+			ew.stopORCAListener()
+		}
+		ew.pickedSC = nil
+	}
+}
+
+// Close stops the balancer.  It cancels any ongoing scheduler updates and
+// stops any ORCA listeners.
+func (b *wrrBalancer) Close() {
+	b.mu.Lock()
+	if b.stopPicker != nil {
+		b.stopPicker.Fire()
+		b.stopPicker = nil
+	}
+	b.mu.Unlock()
+
+	// Ensure any lingering OOB watchers are stopped.
+	for _, ew := range b.endpointToWeight.All() {
+		if ew.stopORCAListener != nil {
+			ew.stopORCAListener()
+		}
+	}
+	b.child.Close()
+	if b.logger.V(2) {
+		b.logger.Infof("Shutdown")
+	}
+}
+
+func (b *wrrBalancer) ExitIdle() {
+	b.child.ExitIdle()
+}
+
+// picker is the WRR policy's picker.  It uses live-updating backend weights to
+// update the scheduler periodically and ensure picks are routed proportional
+// to those weights.
+type picker struct {
+	scheduler unsafe.Pointer // *scheduler; accessed atomically
+	idx       atomic.Uint32  // incrementing value used by the scheduler
+	cfg       *lbConfig      // active config when picker created
+
+	weightedPickers []pickerWeightedEndpoint // all READY pickers
+
+	// The following fields are immutable.
+	target          string
+	locality        string
+	clusterName     string
+	metricsRecorder estats.MetricsRecorder
+}
+
+func (p *picker) endpointWeights(recordMetrics bool) []float64 {
+	wp := make([]float64, len(p.weightedPickers))
+	now := internal.TimeNow()
+	for i, wpi := range p.weightedPickers {
+		wp[i] = wpi.weightedEndpoint.weight(now, time.Duration(p.cfg.WeightExpirationPeriod), time.Duration(p.cfg.BlackoutPeriod), recordMetrics)
+	}
+	return wp
+}
+
+func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	// Read the scheduler atomically.  All scheduler operations are threadsafe,
+	// and if the scheduler is replaced during this usage, we want to use the
+	// scheduler that was live when the pick started.
+	sched := *(*scheduler)(atomic.LoadPointer(&p.scheduler))
+
+	pickedPicker := p.weightedPickers[sched.nextIndex()]
+	pr, err := pickedPicker.picker.Pick(info)
+	if err != nil {
+		return balancer.PickResult{}, err
+	}
+	if !p.cfg.EnableOOBLoadReport {
+		oldDone := pr.Done
+		pr.Done = func(info balancer.DoneInfo) {
+			if load, ok := info.ServerLoad.(*v3orcapb.OrcaLoadReport); ok && load != nil {
+				pickedPicker.weightedEndpoint.OnLoadReport(load)
+			}
+			if oldDone != nil {
+				oldDone(info)
+			}
+		}
+	}
+	return pr, nil
+}
+
+func (p *picker) inc() uint32 {
+	return p.idx.Add(1)
+}
+
+func (p *picker) regenerateScheduler() {
+	s := p.newScheduler(true)
+	atomic.StorePointer(&p.scheduler, unsafe.Pointer(&s))
+}
+
+func (p *picker) start(stopPicker *grpcsync.Event) {
+	p.regenerateScheduler()
+	if len(p.weightedPickers) == 1 {
+		// No need to regenerate weights with only one backend.
+		return
+	}
+
+	go func() {
+		ticker := time.NewTicker(time.Duration(p.cfg.WeightUpdatePeriod))
+		defer ticker.Stop()
+		for {
+			select {
+			case <-stopPicker.Done():
+				return
+			case <-ticker.C:
+				p.regenerateScheduler()
+			}
+		}
+	}()
+}
+
+// endpointWeight is the weight for an endpoint. It tracks the SubConn that will
+// be picked for the endpoint, and other parameters relevant to computing the
+// effective weight. When needed, it also tracks connectivity state, listens for
+// metrics updates by implementing the orca.OOBListener interface and manages
+// that listener.
+type endpointWeight struct {
+	// The following fields are immutable.
+	logger          *grpclog.PrefixLogger
+	target          string
+	metricsRecorder estats.MetricsRecorder
+	locality        string
+	clusterName     string
+
+	// The following fields are only accessed on calls into the LB policy, and
+	// do not need a mutex.
+	connectivityState connectivity.State
+	stopORCAListener  func()
+	// The first SubConn for the endpoint that goes READY when endpoint has no
+	// READY SubConns yet, cleared on that sc disconnecting (i.e. going out of
+	// READY). Represents what pick first will use as it's picked SubConn for
+	// this endpoint.
+	pickedSC balancer.SubConn
+
+	// The following fields are accessed asynchronously and are protected by
+	// mu.  Note that mu may not be held when calling into the stopORCAListener
+	// or when registering a new listener, as those calls require the ORCA
+	// producer mu which is held when calling the listener, and the listener
+	// holds mu.
+	mu            sync.Mutex
+	weightVal     float64
+	nonEmptySince time.Time
+	lastUpdated   time.Time
+	cfg           *lbConfig
+}
+
+func (w *endpointWeight) OnLoadReport(load *v3orcapb.OrcaLoadReport) {
+	if w.logger.V(2) {
+		w.logger.Infof("Received load report for subchannel %v: %v", w.pickedSC, load)
+	}
+	// Update weights of this endpoint according to the reported load.
+	utilization := load.ApplicationUtilization
+	if utilization == 0 {
+		utilization = load.CpuUtilization
+	}
+	if utilization == 0 || load.RpsFractional == 0 {
+		if w.logger.V(2) {
+			w.logger.Infof("Ignoring empty load report for subchannel %v", w.pickedSC)
+		}
+		return
+	}
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	errorRate := load.Eps / load.RpsFractional
+	w.weightVal = load.RpsFractional / (utilization + errorRate*w.cfg.ErrorUtilizationPenalty)
+	if w.logger.V(2) {
+		w.logger.Infof("New weight for subchannel %v: %v", w.pickedSC, w.weightVal)
+	}
+
+	w.lastUpdated = internal.TimeNow()
+	if w.nonEmptySince.Equal(time.Time{}) {
+		w.nonEmptySince = w.lastUpdated
+	}
+}
+
+// updateConfig updates the parameters of the WRR policy and
+// stops/starts/restarts the ORCA OOB listener.
+func (w *endpointWeight) updateConfig(cfg *lbConfig) {
+	w.mu.Lock()
+	oldCfg := w.cfg
+	w.cfg = cfg
+	w.mu.Unlock()
+
+	if cfg.EnableOOBLoadReport == oldCfg.EnableOOBLoadReport &&
+		cfg.OOBReportingPeriod == oldCfg.OOBReportingPeriod {
+		// Load reporting wasn't enabled before or after, or load reporting was
+		// enabled before and after, and had the same period.  (Note that with
+		// load reporting disabled, OOBReportingPeriod is always 0.)
+		return
+	}
+	// (Re)start the listener to use the new config's settings for OOB
+	// reporting.
+	w.updateORCAListener(cfg)
+}
+
+func (w *endpointWeight) updateORCAListener(cfg *lbConfig) {
+	if w.stopORCAListener != nil {
+		w.stopORCAListener()
+	}
+	if !cfg.EnableOOBLoadReport {
+		w.stopORCAListener = nil
+		return
+	}
+	if w.pickedSC == nil { // No picked SC for this endpoint yet, nothing to listen on.
+		return
+	}
+	if w.logger.V(2) {
+		w.logger.Infof("Registering ORCA listener for %v with interval %v", w.pickedSC, cfg.OOBReportingPeriod)
+	}
+	opts := orca.OOBListenerOptions{ReportInterval: time.Duration(cfg.OOBReportingPeriod)}
+	w.stopORCAListener = orca.RegisterOOBListener(w.pickedSC, w, opts)
+}
+
+// weight returns the current effective weight of the endpoint, taking into
+// account the parameters.  Returns 0 for blacked out or expired data, which
+// will cause the backend weight to be treated as the mean of the weights of the
+// other backends. If forScheduler is set to true, this function will emit
+// metrics through the metrics registry.
+func (w *endpointWeight) weight(now time.Time, weightExpirationPeriod, blackoutPeriod time.Duration, recordMetrics bool) (weight float64) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	if recordMetrics {
+		defer func() {
+			endpointWeightsMetric.Record(w.metricsRecorder, weight, w.target, w.locality, w.clusterName)
+		}()
+	}
+
+	// The endpoint has not received a load report (i.e. just turned READY with
+	// no load report).
+	if w.lastUpdated.Equal(time.Time{}) {
+		endpointWeightNotYetUsableMetric.Record(w.metricsRecorder, 1, w.target, w.locality, w.clusterName)
+		return 0
+	}
+
+	// If the most recent update was longer ago than the expiration period,
+	// reset nonEmptySince so that we apply the blackout period again if we
+	// start getting data again in the future, and return 0.
+	if now.Sub(w.lastUpdated) >= weightExpirationPeriod {
+		if recordMetrics {
+			endpointWeightStaleMetric.Record(w.metricsRecorder, 1, w.target, w.locality, w.clusterName)
+		}
+		w.nonEmptySince = time.Time{}
+		return 0
+	}
+
+	// If we don't have at least blackoutPeriod worth of data, return 0.
+	if blackoutPeriod != 0 && (w.nonEmptySince.Equal(time.Time{}) || now.Sub(w.nonEmptySince) < blackoutPeriod) {
+		if recordMetrics {
+			endpointWeightNotYetUsableMetric.Record(w.metricsRecorder, 1, w.target, w.locality, w.clusterName)
+		}
+		return 0
+	}
+
+	return w.weightVal
+}
+
+type backendServiceKey struct{}
+
+// SetBackendService stores the backendService on the resolver state so
+// that it can be used later as a label in wrr metrics.
+func SetBackendService(state resolver.State, backendService string) resolver.State {
+	state.Attributes = state.Attributes.WithValue(backendServiceKey{}, backendService)
+	return state
+}
+
+// getBackendServiceFromState retrieves the cluster name stored as an attribute
+// in the resolver state.
+func backendServiceFromState(state resolver.State) string {
+	v := state.Attributes.Value(backendServiceKey{})
+	name, _ := v.(string)
+	return name
+}

balancer/weightedroundrobin/balancer_test.go

@@ -0,0 +1,869 @@
+/*
+ *
+ * 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 weightedroundrobin_test
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/stubserver"
+	"google.golang.org/grpc/internal/testutils/roundrobin"
+	"google.golang.org/grpc/internal/testutils/stats"
+	"google.golang.org/grpc/orca"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/resolver"
+
+	wrr "google.golang.org/grpc/balancer/weightedroundrobin"
+	iwrr "google.golang.org/grpc/balancer/weightedroundrobin/internal"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+const defaultTestTimeout = 10 * time.Second
+const weightUpdatePeriod = 50 * time.Millisecond
+const weightExpirationPeriod = time.Minute
+const oobReportingInterval = 10 * time.Millisecond
+
+func init() {
+	iwrr.AllowAnyWeightUpdatePeriod = true
+}
+
+func boolp(b bool) *bool          { return &b }
+func float64p(f float64) *float64 { return &f }
+func stringp(s string) *string    { return &s }
+
+var (
+	perCallConfig = iwrr.LBConfig{
+		EnableOOBLoadReport:     boolp(false),
+		OOBReportingPeriod:      stringp("0.005s"),
+		BlackoutPeriod:          stringp("0s"),
+		WeightExpirationPeriod:  stringp("60s"),
+		WeightUpdatePeriod:      stringp(".050s"),
+		ErrorUtilizationPenalty: float64p(0),
+	}
+	oobConfig = iwrr.LBConfig{
+		EnableOOBLoadReport:     boolp(true),
+		OOBReportingPeriod:      stringp("0.005s"),
+		BlackoutPeriod:          stringp("0s"),
+		WeightExpirationPeriod:  stringp("60s"),
+		WeightUpdatePeriod:      stringp(".050s"),
+		ErrorUtilizationPenalty: float64p(0),
+	}
+	testMetricsConfig = iwrr.LBConfig{
+		EnableOOBLoadReport:     boolp(false),
+		OOBReportingPeriod:      stringp("0.005s"),
+		BlackoutPeriod:          stringp("0s"),
+		WeightExpirationPeriod:  stringp("60s"),
+		WeightUpdatePeriod:      stringp("30s"),
+		ErrorUtilizationPenalty: float64p(0),
+	}
+)
+
+type testServer struct {
+	*stubserver.StubServer
+
+	oobMetrics  orca.ServerMetricsRecorder // Attached to the OOB stream.
+	callMetrics orca.CallMetricsRecorder   // Attached to per-call metrics.
+}
+
+type reportType int
+
+const (
+	reportNone reportType = iota
+	reportOOB
+	reportCall
+	reportBoth
+)
+
+func startServer(t *testing.T, r reportType) *testServer {
+	t.Helper()
+
+	smr := orca.NewServerMetricsRecorder()
+	cmr := orca.NewServerMetricsRecorder().(orca.CallMetricsRecorder)
+
+	ss := &stubserver.StubServer{
+		EmptyCallF: func(ctx context.Context, _ *testpb.Empty) (*testpb.Empty, error) {
+			if r := orca.CallMetricsRecorderFromContext(ctx); r != nil {
+				// Copy metrics from what the test set in cmr into r.
+				sm := cmr.(orca.ServerMetricsProvider).ServerMetrics()
+				r.SetApplicationUtilization(sm.AppUtilization)
+				r.SetQPS(sm.QPS)
+				r.SetEPS(sm.EPS)
+			}
+			return &testpb.Empty{}, nil
+		},
+	}
+
+	var sopts []grpc.ServerOption
+	if r == reportCall || r == reportBoth {
+		sopts = append(sopts, orca.CallMetricsServerOption(nil))
+	}
+
+	if r == reportOOB || r == reportBoth {
+		oso := orca.ServiceOptions{
+			ServerMetricsProvider: smr,
+			MinReportingInterval:  10 * time.Millisecond,
+		}
+		internal.ORCAAllowAnyMinReportingInterval.(func(so *orca.ServiceOptions))(&oso)
+		sopts = append(sopts, stubserver.RegisterServiceServerOption(func(s grpc.ServiceRegistrar) {
+			if err := orca.Register(s, oso); err != nil {
+				t.Fatalf("Failed to register orca service: %v", err)
+			}
+		}))
+	}
+
+	if err := ss.StartServer(sopts...); err != nil {
+		t.Fatalf("Error starting server: %v", err)
+	}
+	t.Cleanup(ss.Stop)
+
+	return &testServer{
+		StubServer:  ss,
+		oobMetrics:  smr,
+		callMetrics: cmr,
+	}
+}
+
+func svcConfig(t *testing.T, wrrCfg iwrr.LBConfig) string {
+	t.Helper()
+	m, err := json.Marshal(wrrCfg)
+	if err != nil {
+		t.Fatalf("Error marshaling JSON %v: %v", wrrCfg, err)
+	}
+	sc := fmt.Sprintf(`{"loadBalancingConfig": [ {%q:%v} ] }`, wrr.Name, string(m))
+	t.Logf("Marshaled service config: %v", sc)
+	return sc
+}
+
+// Tests basic functionality with one address.  With only one address, load
+// reporting doesn't affect routing at all.
+func (s) TestBalancer_OneAddress(t *testing.T) {
+	testCases := []struct {
+		rt  reportType
+		cfg iwrr.LBConfig
+	}{
+		{rt: reportNone, cfg: perCallConfig},
+		{rt: reportCall, cfg: perCallConfig},
+		{rt: reportOOB, cfg: oobConfig},
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("reportType:%v", tc.rt), func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+			defer cancel()
+
+			srv := startServer(t, tc.rt)
+
+			sc := svcConfig(t, tc.cfg)
+			if err := srv.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+				t.Fatalf("Error starting client: %v", err)
+			}
+
+			// Perform many RPCs to ensure the LB policy works with 1 address.
+			for i := 0; i < 100; i++ {
+				srv.callMetrics.SetQPS(float64(i))
+				srv.oobMetrics.SetQPS(float64(i))
+				if _, err := srv.Client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+					t.Fatalf("Error from EmptyCall: %v", err)
+				}
+				time.Sleep(time.Millisecond) // Delay; test will run 100ms and should perform ~10 weight updates
+			}
+		})
+	}
+}
+
+// TestWRRMetricsBasic tests metrics emitted from the WRR balancer. It
+// configures a weighted round robin balancer as the top level balancer of a
+// ClientConn, and configures a fake stats handler on the ClientConn to receive
+// metrics. It verifies stats emitted from the Weighted Round Robin Balancer on
+// balancer startup case which triggers the first picker and scheduler update
+// before any load reports are received.
+//
+// Note that this test and others, metrics emission assertions are a snapshot
+// of the most recently emitted metrics. This is due to the nondeterminism of
+// scheduler updates with respect to test bodies, so the assertions made are
+// from the most recently synced state of the system (picker/scheduler) from the
+// test body.
+func (s) TestWRRMetricsBasic(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv := startServer(t, reportCall)
+	sc := svcConfig(t, testMetricsConfig)
+
+	tmr := stats.NewTestMetricsRecorder()
+	if err := srv.StartClient(grpc.WithDefaultServiceConfig(sc), grpc.WithStatsHandler(tmr)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	srv.callMetrics.SetQPS(float64(1))
+
+	if _, err := srv.Client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("Error from EmptyCall: %v", err)
+	}
+
+	if got, _ := tmr.Metric("grpc.lb.wrr.rr_fallback"); got != 1 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.rr_fallback", got, 1)
+	}
+	if got, _ := tmr.Metric("grpc.lb.wrr.endpoint_weight_stale"); got != 0 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.endpoint_weight_stale", got, 0)
+	}
+	if got, _ := tmr.Metric("grpc.lb.wrr.endpoint_weight_not_yet_usable"); got != 1 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.endpoint_weight_not_yet_usable", got, 1)
+	}
+	// Unusable, so no endpoint weight. Due to only one SubConn, this will never
+	// update the weight. Thus, this will stay 0.
+	if got, _ := tmr.Metric("grpc.lb.wrr.endpoint_weight_stale"); got != 0 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.endpoint_weight_stale", got, 0)
+	}
+}
+
+// Tests two addresses with ORCA reporting disabled (should fall back to pure
+// RR).
+func (s) TestBalancer_TwoAddresses_ReportingDisabled(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv1 := startServer(t, reportNone)
+	srv2 := startServer(t, reportNone)
+
+	sc := svcConfig(t, perCallConfig)
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+	// Perform many RPCs to ensure the LB policy works with 2 addresses.
+	for i := 0; i < 20; i++ {
+		roundrobin.CheckRoundRobinRPCs(ctx, srv1.Client, addrs)
+	}
+}
+
+// Tests two addresses with per-call ORCA reporting enabled.  Checks the
+// backends are called in the appropriate ratios.
+func (s) TestBalancer_TwoAddresses_ReportingEnabledPerCall(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv1 := startServer(t, reportCall)
+	srv2 := startServer(t, reportCall)
+
+	// srv1 starts loaded and srv2 starts without load; ensure RPCs are routed
+	// disproportionately to srv2 (10:1).
+	srv1.callMetrics.SetQPS(10.0)
+	srv1.callMetrics.SetApplicationUtilization(1.0)
+
+	srv2.callMetrics.SetQPS(10.0)
+	srv2.callMetrics.SetApplicationUtilization(.1)
+
+	sc := svcConfig(t, perCallConfig)
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+	// Call each backend once to ensure the weights have been received.
+	ensureReached(ctx, t, srv1.Client, 2)
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+}
+
+// Tests two addresses with OOB ORCA reporting enabled.  Checks the backends
+// are called in the appropriate ratios.
+func (s) TestBalancer_TwoAddresses_ReportingEnabledOOB(t *testing.T) {
+	testCases := []struct {
+		name       string
+		utilSetter func(orca.ServerMetricsRecorder, float64)
+	}{{
+		name: "application_utilization",
+		utilSetter: func(smr orca.ServerMetricsRecorder, val float64) {
+			smr.SetApplicationUtilization(val)
+		},
+	}, {
+		name: "cpu_utilization",
+		utilSetter: func(smr orca.ServerMetricsRecorder, val float64) {
+			smr.SetCPUUtilization(val)
+		},
+	}, {
+		name: "application over cpu",
+		utilSetter: func(smr orca.ServerMetricsRecorder, val float64) {
+			smr.SetApplicationUtilization(val)
+			smr.SetCPUUtilization(2.0) // ignored because ApplicationUtilization is set
+		},
+	}}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+			defer cancel()
+
+			srv1 := startServer(t, reportOOB)
+			srv2 := startServer(t, reportOOB)
+
+			// srv1 starts loaded and srv2 starts without load; ensure RPCs are routed
+			// disproportionately to srv2 (10:1).
+			srv1.oobMetrics.SetQPS(10.0)
+			tc.utilSetter(srv1.oobMetrics, 1.0)
+
+			srv2.oobMetrics.SetQPS(10.0)
+			tc.utilSetter(srv2.oobMetrics, 0.1)
+
+			sc := svcConfig(t, oobConfig)
+			if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+				t.Fatalf("Error starting client: %v", err)
+			}
+			addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
+			srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+			// Call each backend once to ensure the weights have been received.
+			ensureReached(ctx, t, srv1.Client, 2)
+
+			// Wait for the weight update period to allow the new weights to be processed.
+			time.Sleep(weightUpdatePeriod)
+			checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+		})
+	}
+}
+
+// Tests two addresses with OOB ORCA reporting enabled, where the reports
+// change over time.  Checks the backends are called in the appropriate ratios
+// before and after modifying the reports.
+func (s) TestBalancer_TwoAddresses_UpdateLoads(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv1 := startServer(t, reportOOB)
+	srv2 := startServer(t, reportOOB)
+
+	// srv1 starts loaded and srv2 starts without load; ensure RPCs are routed
+	// disproportionately to srv2 (10:1).
+	srv1.oobMetrics.SetQPS(10.0)
+	srv1.oobMetrics.SetApplicationUtilization(1.0)
+
+	srv2.oobMetrics.SetQPS(10.0)
+	srv2.oobMetrics.SetApplicationUtilization(.1)
+
+	sc := svcConfig(t, oobConfig)
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+	// Call each backend once to ensure the weights have been received.
+	ensureReached(ctx, t, srv1.Client, 2)
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+
+	// Update the loads so srv2 is loaded and srv1 is not; ensure RPCs are
+	// routed disproportionately to srv1.
+	srv1.oobMetrics.SetQPS(10.0)
+	srv1.oobMetrics.SetApplicationUtilization(.1)
+
+	srv2.oobMetrics.SetQPS(10.0)
+	srv2.oobMetrics.SetApplicationUtilization(1.0)
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod + oobReportingInterval)
+	checkWeights(ctx, t, srvWeight{srv1, 10}, srvWeight{srv2, 1})
+}
+
+// Tests two addresses with OOB ORCA reporting enabled, then with switching to
+// per-call reporting.  Checks the backends are called in the appropriate
+// ratios before and after the change.
+func (s) TestBalancer_TwoAddresses_OOBThenPerCall(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv1 := startServer(t, reportBoth)
+	srv2 := startServer(t, reportBoth)
+
+	// srv1 starts loaded and srv2 starts without load; ensure RPCs are routed
+	// disproportionately to srv2 (10:1).
+	srv1.oobMetrics.SetQPS(10.0)
+	srv1.oobMetrics.SetApplicationUtilization(1.0)
+
+	srv2.oobMetrics.SetQPS(10.0)
+	srv2.oobMetrics.SetApplicationUtilization(.1)
+
+	// For per-call metrics (not used initially), srv2 reports that it is
+	// loaded and srv1 reports low load.  After confirming OOB works, switch to
+	// per-call and confirm the new routing weights are applied.
+	srv1.callMetrics.SetQPS(10.0)
+	srv1.callMetrics.SetApplicationUtilization(.1)
+
+	srv2.callMetrics.SetQPS(10.0)
+	srv2.callMetrics.SetApplicationUtilization(1.0)
+
+	sc := svcConfig(t, oobConfig)
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+	// Call each backend once to ensure the weights have been received.
+	ensureReached(ctx, t, srv1.Client, 2)
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+
+	// Update to per-call weights.
+	c := svcConfig(t, perCallConfig)
+	parsedCfg := srv1.R.CC().ParseServiceConfig(c)
+	if parsedCfg.Err != nil {
+		panic(fmt.Sprintf("Error parsing config %q: %v", c, parsedCfg.Err))
+	}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs, ServiceConfig: parsedCfg})
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 10}, srvWeight{srv2, 1})
+}
+
+// TestEndpoints_SharedAddress tests the case where two endpoints have the same
+// address. The expected behavior is undefined, however the program should not
+// crash.
+func (s) TestEndpoints_SharedAddress(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv := startServer(t, reportCall)
+	sc := svcConfig(t, perCallConfig)
+	if err := srv.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+
+	endpointsSharedAddress := []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: srv.Address}}}, {Addresses: []resolver.Address{{Addr: srv.Address}}}}
+	srv.R.UpdateState(resolver.State{Endpoints: endpointsSharedAddress})
+
+	// Make some RPC's and make sure doesn't crash. It should go to one of the
+	// endpoints addresses, it's undefined which one it will choose and the load
+	// reporting might not work, but it should be able to make an RPC.
+	for i := 0; i < 10; i++ {
+		if _, err := srv.Client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+			t.Fatalf("EmptyCall failed with err: %v", err)
+		}
+	}
+}
+
+// TestEndpoints_MultipleAddresses tests WRR on endpoints with numerous
+// addresses. It configures WRR with two endpoints with one bad address followed
+// by a good address. It configures two backends that each report per call
+// metrics, each corresponding to the two endpoints good address. It then
+// asserts load is distributed as expected corresponding to the call metrics
+// received.
+func (s) TestEndpoints_MultipleAddresses(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	srv1 := startServer(t, reportCall)
+	srv2 := startServer(t, reportCall)
+
+	srv1.callMetrics.SetQPS(10.0)
+	srv1.callMetrics.SetApplicationUtilization(.1)
+
+	srv2.callMetrics.SetQPS(10.0)
+	srv2.callMetrics.SetApplicationUtilization(1.0)
+
+	sc := svcConfig(t, perCallConfig)
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+
+	twoEndpoints := []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: "bad-address-1"}, {Addr: srv1.Address}}}, {Addresses: []resolver.Address{{Addr: "bad-address-2"}, {Addr: srv2.Address}}}}
+	srv1.R.UpdateState(resolver.State{Endpoints: twoEndpoints})
+
+	// Call each backend once to ensure the weights have been received.
+	ensureReached(ctx, t, srv1.Client, 2)
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 10}, srvWeight{srv2, 1})
+}
+
+// Tests two addresses with OOB ORCA reporting enabled and a non-zero error
+// penalty applied.
+func (s) TestBalancer_TwoAddresses_ErrorPenalty(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv1 := startServer(t, reportOOB)
+	srv2 := startServer(t, reportOOB)
+
+	// srv1 starts loaded and srv2 starts without load; ensure RPCs are routed
+	// disproportionately to srv2 (10:1).  EPS values are set (but ignored
+	// initially due to ErrorUtilizationPenalty=0).  Later EUP will be updated
+	// to 0.9 which will cause the weights to be equal and RPCs to be routed
+	// 50/50.
+	srv1.oobMetrics.SetQPS(10.0)
+	srv1.oobMetrics.SetApplicationUtilization(1.0)
+	srv1.oobMetrics.SetEPS(0)
+	// srv1 weight before: 10.0 / 1.0 = 10.0
+	// srv1 weight after:  10.0 / 1.0 = 10.0
+
+	srv2.oobMetrics.SetQPS(10.0)
+	srv2.oobMetrics.SetApplicationUtilization(.1)
+	srv2.oobMetrics.SetEPS(10.0)
+	// srv2 weight before: 10.0 / 0.1 = 100.0
+	// srv2 weight after:  10.0 / 1.0 = 10.0
+
+	sc := svcConfig(t, oobConfig)
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+	// Call each backend once to ensure the weights have been received.
+	ensureReached(ctx, t, srv1.Client, 2)
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+
+	// Update to include an error penalty in the weights.
+	newCfg := oobConfig
+	newCfg.ErrorUtilizationPenalty = float64p(0.9)
+	c := svcConfig(t, newCfg)
+	parsedCfg := srv1.R.CC().ParseServiceConfig(c)
+	if parsedCfg.Err != nil {
+		panic(fmt.Sprintf("Error parsing config %q: %v", c, parsedCfg.Err))
+	}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs, ServiceConfig: parsedCfg})
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod + oobReportingInterval)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 1})
+}
+
+// Tests that the blackout period causes backends to use 0 as their weight
+// (meaning to use the average weight) until the blackout period elapses.
+func (s) TestBalancer_TwoAddresses_BlackoutPeriod(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	var mu sync.Mutex
+	start := time.Now()
+	now := start
+	setNow := func(t time.Time) {
+		mu.Lock()
+		defer mu.Unlock()
+		now = t
+	}
+
+	setTimeNow(func() time.Time {
+		mu.Lock()
+		defer mu.Unlock()
+		return now
+	})
+	t.Cleanup(func() { setTimeNow(time.Now) })
+
+	testCases := []struct {
+		blackoutPeriodCfg *string
+		blackoutPeriod    time.Duration
+	}{{
+		blackoutPeriodCfg: stringp("1s"),
+		blackoutPeriod:    time.Second,
+	}, {
+		blackoutPeriodCfg: nil,
+		blackoutPeriod:    10 * time.Second, // the default
+	}}
+	for _, tc := range testCases {
+		setNow(start)
+		srv1 := startServer(t, reportOOB)
+		srv2 := startServer(t, reportOOB)
+
+		// srv1 starts loaded and srv2 starts without load; ensure RPCs are routed
+		// disproportionately to srv2 (10:1).
+		srv1.oobMetrics.SetQPS(10.0)
+		srv1.oobMetrics.SetApplicationUtilization(1.0)
+
+		srv2.oobMetrics.SetQPS(10.0)
+		srv2.oobMetrics.SetApplicationUtilization(.1)
+
+		cfg := oobConfig
+		cfg.BlackoutPeriod = tc.blackoutPeriodCfg
+		sc := svcConfig(t, cfg)
+		if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+			t.Fatalf("Error starting client: %v", err)
+		}
+		addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
+		srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+		// Call each backend once to ensure the weights have been received.
+		ensureReached(ctx, t, srv1.Client, 2)
+
+		// Wait for the weight update period to allow the new weights to be processed.
+		time.Sleep(weightUpdatePeriod)
+		// During the blackout period (1s) we should route roughly 50/50.
+		checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 1})
+
+		// Advance time to right before the blackout period ends and the weights
+		// should still be zero.
+		setNow(start.Add(tc.blackoutPeriod - time.Nanosecond))
+		// Wait for the weight update period to allow the new weights to be processed.
+		time.Sleep(weightUpdatePeriod)
+		checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 1})
+
+		// Advance time to right after the blackout period ends and the weights
+		// should now activate.
+		setNow(start.Add(tc.blackoutPeriod))
+		// Wait for the weight update period to allow the new weights to be processed.
+		time.Sleep(weightUpdatePeriod)
+		checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+	}
+}
+
+// Tests that the weight expiration period causes backends to use 0 as their
+// weight (meaning to use the average weight) once the expiration period
+// elapses.
+func (s) TestBalancer_TwoAddresses_WeightExpiration(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	var mu sync.Mutex
+	start := time.Now()
+	now := start
+	setNow := func(t time.Time) {
+		mu.Lock()
+		defer mu.Unlock()
+		now = t
+	}
+	setTimeNow(func() time.Time {
+		mu.Lock()
+		defer mu.Unlock()
+		return now
+	})
+	t.Cleanup(func() { setTimeNow(time.Now) })
+
+	srv1 := startServer(t, reportBoth)
+	srv2 := startServer(t, reportBoth)
+
+	// srv1 starts loaded and srv2 starts without load; ensure RPCs are routed
+	// disproportionately to srv2 (10:1).  Because the OOB reporting interval
+	// is 1 minute but the weights expire in 1 second, routing will go to 50/50
+	// after the weights expire.
+	srv1.oobMetrics.SetQPS(10.0)
+	srv1.oobMetrics.SetApplicationUtilization(1.0)
+
+	srv2.oobMetrics.SetQPS(10.0)
+	srv2.oobMetrics.SetApplicationUtilization(.1)
+
+	cfg := oobConfig
+	cfg.OOBReportingPeriod = stringp("60s")
+	sc := svcConfig(t, cfg)
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+	// Call each backend once to ensure the weights have been received.
+	ensureReached(ctx, t, srv1.Client, 2)
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+
+	// Advance what time.Now returns to the weight expiration time minus 1s to
+	// ensure all weights are still honored.
+	setNow(start.Add(weightExpirationPeriod - time.Second))
+
+	// Wait for the weight update period to allow the new weights to be processed.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+
+	// Advance what time.Now returns to the weight expiration time plus 1s to
+	// ensure all weights expired and addresses are routed evenly.
+	setNow(start.Add(weightExpirationPeriod + time.Second))
+
+	// Wait for the weight expiration period so the weights have expired.
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 1})
+}
+
+// Tests logic surrounding subchannel management.
+func (s) TestBalancer_AddressesChanging(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv1 := startServer(t, reportBoth)
+	srv2 := startServer(t, reportBoth)
+	srv3 := startServer(t, reportBoth)
+	srv4 := startServer(t, reportBoth)
+
+	// srv1: weight 10
+	srv1.oobMetrics.SetQPS(10.0)
+	srv1.oobMetrics.SetApplicationUtilization(1.0)
+	// srv2: weight 100
+	srv2.oobMetrics.SetQPS(10.0)
+	srv2.oobMetrics.SetApplicationUtilization(.1)
+	// srv3: weight 20
+	srv3.oobMetrics.SetQPS(20.0)
+	srv3.oobMetrics.SetApplicationUtilization(1.0)
+	// srv4: weight 200
+	srv4.oobMetrics.SetQPS(20.0)
+	srv4.oobMetrics.SetApplicationUtilization(.1)
+
+	sc := svcConfig(t, oobConfig)
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	srv2.Client = srv1.Client
+	addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}, {Addr: srv3.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+
+	// Call each backend once to ensure the weights have been received.
+	ensureReached(ctx, t, srv1.Client, 3)
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10}, srvWeight{srv3, 2})
+
+	// Add backend 4
+	addrs = append(addrs, resolver.Address{Addr: srv4.Address})
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10}, srvWeight{srv3, 2}, srvWeight{srv4, 20})
+
+	// Shutdown backend 3.  RPCs will no longer be routed to it.
+	srv3.Stop()
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10}, srvWeight{srv4, 20})
+
+	// Remove addresses 2 and 3.  RPCs will no longer be routed to 2 either.
+	addrs = []resolver.Address{{Addr: srv1.Address}, {Addr: srv4.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv4, 20})
+
+	// Re-add 2 and remove the rest.
+	addrs = []resolver.Address{{Addr: srv2.Address}}
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv2, 10})
+
+	// Re-add 4.
+	addrs = append(addrs, resolver.Address{Addr: srv4.Address})
+	srv1.R.UpdateState(resolver.State{Addresses: addrs})
+	time.Sleep(weightUpdatePeriod)
+	checkWeights(ctx, t, srvWeight{srv2, 10}, srvWeight{srv4, 20})
+}
+
+func ensureReached(ctx context.Context, t *testing.T, c testgrpc.TestServiceClient, n int) {
+	t.Helper()
+	reached := make(map[string]struct{})
+	for len(reached) != n {
+		var peer peer.Peer
+		if _, err := c.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
+			t.Fatalf("Error from EmptyCall: %v", err)
+		}
+		reached[peer.Addr.String()] = struct{}{}
+	}
+}
+
+type srvWeight struct {
+	srv *testServer
+	w   int
+}
+
+const rrIterations = 100
+
+// checkWeights does rrIterations RPCs and expects the different backends to be
+// routed in a ratio as determined by the srvWeights passed in.  Allows for
+// some variance (+/- 2 RPCs per backend).
+func checkWeights(ctx context.Context, t *testing.T, sws ...srvWeight) {
+	t.Helper()
+
+	c := sws[0].srv.Client
+
+	// Replace the weights with approximate counts of RPCs wanted given the
+	// iterations performed.
+	weightSum := 0
+	for _, sw := range sws {
+		weightSum += sw.w
+	}
+	for i := range sws {
+		sws[i].w = rrIterations * sws[i].w / weightSum
+	}
+
+	for attempts := 0; attempts < 10; attempts++ {
+		serverCounts := make(map[string]int)
+		for i := 0; i < rrIterations; i++ {
+			var peer peer.Peer
+			if _, err := c.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
+				t.Fatalf("Error from EmptyCall: %v; timed out waiting for weighted RR behavior?", err)
+			}
+			serverCounts[peer.Addr.String()]++
+		}
+		if len(serverCounts) != len(sws) {
+			continue
+		}
+		success := true
+		for _, sw := range sws {
+			c := serverCounts[sw.srv.Address]
+			if c < sw.w-2 || c > sw.w+2 {
+				success = false
+				break
+			}
+		}
+		if success {
+			t.Logf("Passed iteration %v; counts: %v", attempts, serverCounts)
+			return
+		}
+		t.Logf("Failed iteration %v; counts: %v; want %+v", attempts, serverCounts, sws)
+		time.Sleep(5 * time.Millisecond)
+	}
+	t.Fatalf("Failed to route RPCs with proper ratio")
+}
+
+func init() {
+	setTimeNow(time.Now)
+	iwrr.TimeNow = timeNow
+}
+
+var timeNowFunc atomic.Value // func() time.Time
+
+func timeNow() time.Time {
+	return timeNowFunc.Load().(func() time.Time)()
+}
+
+func setTimeNow(f func() time.Time) {
+	timeNowFunc.Store(f)
+}

balancer/weightedroundrobin/config.go

@@ -0,0 +1,59 @@
+/*
+ *
+ * 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 weightedroundrobin
+
+import (
+	iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+type lbConfig struct {
+	serviceconfig.LoadBalancingConfig `json:"-"`
+
+	// Whether to enable out-of-band utilization reporting collection from the
+	// endpoints.  By default, per-request utilization reporting is used.
+	EnableOOBLoadReport bool `json:"enableOobLoadReport,omitempty"`
+
+	// Load reporting interval to request from the server.  Note that the
+	// server may not provide reports as frequently as the client requests.
+	// Used only when enable_oob_load_report is true.  Default is 10 seconds.
+	OOBReportingPeriod iserviceconfig.Duration `json:"oobReportingPeriod,omitempty"`
+
+	// A given endpoint must report load metrics continuously for at least this
+	// long before the endpoint weight will be used.  This avoids churn when
+	// the set of endpoint addresses changes.  Takes effect both immediately
+	// after we establish a connection to an endpoint and after
+	// weight_expiration_period has caused us to stop using the most recent
+	// load metrics.  Default is 10 seconds.
+	BlackoutPeriod iserviceconfig.Duration `json:"blackoutPeriod,omitempty"`
+
+	// If a given endpoint has not reported load metrics in this long,
+	// then we stop using the reported weight.  This ensures that we do
+	// not continue to use very stale weights.  Once we stop using a stale
+	// value, if we later start seeing fresh reports again, the
+	// blackout_period applies.  Defaults to 3 minutes.
+	WeightExpirationPeriod iserviceconfig.Duration `json:"weightExpirationPeriod,omitempty"`
+
+	// How often endpoint weights are recalculated.  Default is 1 second.
+	WeightUpdatePeriod iserviceconfig.Duration `json:"weightUpdatePeriod,omitempty"`
+
+	// The multiplier used to adjust endpoint weights with the error rate
+	// calculated as eps/qps. Default is 1.0.
+	ErrorUtilizationPenalty float64 `json:"errorUtilizationPenalty,omitempty"`
+}

balancer/weightedroundrobin/internal/internal.go

@@ -0,0 +1,44 @@
+/*
+ *
+ * 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 internal allows for easier testing of the weightedroundrobin
+// package.
+package internal
+
+import (
+	"time"
+)
+
+// AllowAnyWeightUpdatePeriod permits any setting of WeightUpdatePeriod for
+// testing.  Normally a minimum of 100ms is applied.
+var AllowAnyWeightUpdatePeriod bool
+
+// LBConfig allows tests to produce a JSON form of the config from the struct
+// instead of using a string.
+type LBConfig struct {
+	EnableOOBLoadReport     *bool    `json:"enableOobLoadReport,omitempty"`
+	OOBReportingPeriod      *string  `json:"oobReportingPeriod,omitempty"`
+	BlackoutPeriod          *string  `json:"blackoutPeriod,omitempty"`
+	WeightExpirationPeriod  *string  `json:"weightExpirationPeriod,omitempty"`
+	WeightUpdatePeriod      *string  `json:"weightUpdatePeriod,omitempty"`
+	ErrorUtilizationPenalty *float64 `json:"errorUtilizationPenalty,omitempty"`
+}
+
+// TimeNow can be overridden by tests to return a different value for the
+// current iserviceconfig.
+var TimeNow = time.Now

balancer/weightedroundrobin/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 weightedroundrobin
+
+import (
+	"fmt"
+
+	"google.golang.org/grpc/grpclog"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+)
+
+const prefix = "[%p] "
+
+var logger = grpclog.Component("weighted-round-robin")
+
+func prefixLogger(p *wrrBalancer) *internalgrpclog.PrefixLogger {
+	return internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(prefix, p))
+}

balancer/weightedroundrobin/metrics_test.go

@@ -0,0 +1,173 @@
+/*
+ *
+ * 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 weightedroundrobin
+
+import (
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/grpctest"
+	iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/internal/testutils/stats"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// TestWRR_Metrics_SubConnWeight tests different scenarios for the weight call
+// on a weighted SubConn, and expects certain metrics for each of these
+// scenarios.
+func (s) TestWRR_Metrics_SubConnWeight(t *testing.T) {
+	tests := []struct {
+		name                           string
+		weightExpirationPeriod         time.Duration
+		blackoutPeriod                 time.Duration
+		lastUpdated                    time.Time
+		nonEmpty                       time.Time
+		nowTime                        time.Time
+		endpointWeightStaleWant        float64
+		endpointWeightNotYetUsableWant float64
+		endpointWeightWant             float64
+	}{
+		// The weighted SubConn's lastUpdated field hasn't been set, so this
+		// SubConn's weight is not yet usable. Thus, should emit that endpoint
+		// weight is not yet usable, and 0 for weight.
+		{
+			name:                           "no weight set",
+			weightExpirationPeriod:         time.Second,
+			blackoutPeriod:                 time.Second,
+			nowTime:                        time.Now(),
+			endpointWeightStaleWant:        0,
+			endpointWeightNotYetUsableWant: 1,
+			endpointWeightWant:             0,
+		},
+		{
+			name:                           "weight expiration",
+			lastUpdated:                    time.Now(),
+			weightExpirationPeriod:         2 * time.Second,
+			blackoutPeriod:                 time.Second,
+			nowTime:                        time.Now().Add(100 * time.Second),
+			endpointWeightStaleWant:        1,
+			endpointWeightNotYetUsableWant: 0,
+			endpointWeightWant:             0,
+		},
+		{
+			name:                           "in blackout period",
+			lastUpdated:                    time.Now(),
+			weightExpirationPeriod:         time.Minute,
+			blackoutPeriod:                 10 * time.Second,
+			nowTime:                        time.Now(),
+			endpointWeightStaleWant:        0,
+			endpointWeightNotYetUsableWant: 1,
+			endpointWeightWant:             0,
+		},
+		{
+			name:                           "normal weight",
+			lastUpdated:                    time.Now(),
+			nonEmpty:                       time.Now(),
+			weightExpirationPeriod:         time.Minute,
+			blackoutPeriod:                 time.Second,
+			nowTime:                        time.Now().Add(10 * time.Second),
+			endpointWeightStaleWant:        0,
+			endpointWeightNotYetUsableWant: 0,
+			endpointWeightWant:             3,
+		},
+		{
+			name:                           "weight expiration takes precdedence over blackout",
+			lastUpdated:                    time.Now(),
+			nonEmpty:                       time.Now(),
+			weightExpirationPeriod:         time.Second,
+			blackoutPeriod:                 time.Minute,
+			nowTime:                        time.Now().Add(10 * time.Second),
+			endpointWeightStaleWant:        1,
+			endpointWeightNotYetUsableWant: 0,
+			endpointWeightWant:             0,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			tmr := stats.NewTestMetricsRecorder()
+			wsc := &endpointWeight{
+				metricsRecorder: tmr,
+				weightVal:       3,
+				lastUpdated:     test.lastUpdated,
+				nonEmptySince:   test.nonEmpty,
+			}
+			wsc.weight(test.nowTime, test.weightExpirationPeriod, test.blackoutPeriod, true)
+
+			if got, _ := tmr.Metric("grpc.lb.wrr.endpoint_weight_stale"); got != test.endpointWeightStaleWant {
+				t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.endpoint_weight_stale", got, test.endpointWeightStaleWant)
+			}
+			if got, _ := tmr.Metric("grpc.lb.wrr.endpoint_weight_not_yet_usable"); got != test.endpointWeightNotYetUsableWant {
+				t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.endpoint_weight_not_yet_usable", got, test.endpointWeightNotYetUsableWant)
+			}
+			if got, _ := tmr.Metric("grpc.lb.wrr.endpoint_weight_stale"); got != test.endpointWeightStaleWant {
+				t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.endpoint_weight_stale", got, test.endpointWeightStaleWant)
+			}
+		})
+	}
+
+}
+
+// TestWRR_Metrics_Scheduler_RR_Fallback tests the round robin fallback metric
+// for scheduler updates. It tests the case with one SubConn, and two SubConns
+// with no weights. Both of these should emit a count metric for round robin
+// fallback.
+func (s) TestWRR_Metrics_Scheduler_RR_Fallback(t *testing.T) {
+	tmr := stats.NewTestMetricsRecorder()
+	ew := &endpointWeight{
+		metricsRecorder: tmr,
+		weightVal:       0,
+	}
+
+	p := &picker{
+		cfg: &lbConfig{
+			BlackoutPeriod:         iserviceconfig.Duration(10 * time.Second),
+			WeightExpirationPeriod: iserviceconfig.Duration(3 * time.Minute),
+		},
+		weightedPickers: []pickerWeightedEndpoint{{weightedEndpoint: ew}},
+		metricsRecorder: tmr,
+	}
+	// There is only one SubConn, so no matter if the SubConn has a weight or
+	// not will fallback to round robin.
+	p.regenerateScheduler()
+	if got, _ := tmr.Metric("grpc.lb.wrr.rr_fallback"); got != 1 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.rr_fallback", got, 1)
+	}
+	tmr.ClearMetrics()
+
+	// With two SubConns, if neither of them have weights, it will also fallback
+	// to round robin.
+	ew2 := &endpointWeight{
+		target:          "target",
+		metricsRecorder: tmr,
+		weightVal:       0,
+	}
+	p.weightedPickers = append(p.weightedPickers, pickerWeightedEndpoint{weightedEndpoint: ew2})
+	p.regenerateScheduler()
+	if got, _ := tmr.Metric("grpc.lb.wrr.rr_fallback"); got != 1 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.wrr.rr_fallback", got, 1)
+	}
+}

balancer/weightedroundrobin/scheduler.go

@@ -0,0 +1,146 @@
+/*
+ *
+ * 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 weightedroundrobin
+
+import (
+	"math"
+)
+
+type scheduler interface {
+	nextIndex() int
+}
+
+// newScheduler uses scWeights to create a new scheduler for selecting endpoints
+// in a picker.  It will return a round robin implementation if at least
+// len(scWeights)-1 are zero or there is only a single endpoint, otherwise it
+// will return an Earliest Deadline First (EDF) scheduler implementation that
+// selects the endpoints according to their weights.
+func (p *picker) newScheduler(recordMetrics bool) scheduler {
+	epWeights := p.endpointWeights(recordMetrics)
+	n := len(epWeights)
+	if n == 0 {
+		return nil
+	}
+	if n == 1 {
+		if recordMetrics {
+			rrFallbackMetric.Record(p.metricsRecorder, 1, p.target, p.locality, p.clusterName)
+		}
+		return &rrScheduler{numSCs: 1, inc: p.inc}
+	}
+	sum := float64(0)
+	numZero := 0
+	max := float64(0)
+	for _, w := range epWeights {
+		sum += w
+		if w > max {
+			max = w
+		}
+		if w == 0 {
+			numZero++
+		}
+	}
+
+	if numZero >= n-1 {
+		if recordMetrics {
+			rrFallbackMetric.Record(p.metricsRecorder, 1, p.target, p.locality, p.clusterName)
+		}
+		return &rrScheduler{numSCs: uint32(n), inc: p.inc}
+	}
+	unscaledMean := sum / float64(n-numZero)
+	scalingFactor := maxWeight / max
+	mean := uint16(math.Round(scalingFactor * unscaledMean))
+
+	weights := make([]uint16, n)
+	allEqual := true
+	for i, w := range epWeights {
+		if w == 0 {
+			// Backends with weight = 0 use the mean.
+			weights[i] = mean
+		} else {
+			scaledWeight := uint16(math.Round(scalingFactor * w))
+			weights[i] = scaledWeight
+			if scaledWeight != mean {
+				allEqual = false
+			}
+		}
+	}
+
+	if allEqual {
+		return &rrScheduler{numSCs: uint32(n), inc: p.inc}
+	}
+
+	logger.Infof("using edf scheduler with weights: %v", weights)
+	return &edfScheduler{weights: weights, inc: p.inc}
+}
+
+const maxWeight = math.MaxUint16
+
+// edfScheduler implements EDF using the same algorithm as grpc-c++ here:
+//
+// https://github.com/grpc/grpc/blob/master/src/core/ext/filters/client_channel/lb_policy/weighted_round_robin/static_stride_scheduler.cc
+type edfScheduler struct {
+	inc     func() uint32
+	weights []uint16
+}
+
+// Returns the index in s.weights for the picker to choose.
+func (s *edfScheduler) nextIndex() int {
+	const offset = maxWeight / 2
+
+	for {
+		idx := uint64(s.inc())
+
+		// The sequence number (idx) is split in two: the lower %n gives the
+		// index of the backend, and the rest gives the number of times we've
+		// iterated through all backends. `generation` is used to
+		// deterministically decide whether we pick or skip the backend on this
+		// iteration, in proportion to the backend's weight.
+
+		backendIndex := idx % uint64(len(s.weights))
+		generation := idx / uint64(len(s.weights))
+		weight := uint64(s.weights[backendIndex])
+
+		// We pick a backend `weight` times per `maxWeight` generations. The
+		// multiply and modulus ~evenly spread out the picks for a given
+		// backend between different generations. The offset by `backendIndex`
+		// helps to reduce the chance of multiple consecutive non-picks: if we
+		// have two consecutive backends with an equal, say, 80% weight of the
+		// max, with no offset we would see 1/5 generations that skipped both.
+		// TODO(b/190488683): add test for offset efficacy.
+		mod := uint64(weight*generation+backendIndex*offset) % maxWeight
+
+		if mod < maxWeight-weight {
+			continue
+		}
+		return int(backendIndex)
+	}
+}
+
+// A simple RR scheduler to use for fallback when fewer than two backends have
+// non-zero weights, or all backends have the same weight, or when only one
+// subconn exists.
+type rrScheduler struct {
+	inc    func() uint32
+	numSCs uint32
+}
+
+func (s *rrScheduler) nextIndex() int {
+	idx := s.inc()
+	return int(idx % s.numSCs)
+}

balancer/weightedtarget/logging.go

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

balancer/weightedtarget/weightedaggregator/aggregator.go

@@ -0,0 +1,318 @@
+/*
+ *
+ * Copyright 2020 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 weightedaggregator implements state aggregator for weighted_target
+// balancer.
+//
+// This is a separate package so it can be shared by weighted_target and eds.
+// The eds balancer will be refactored to use weighted_target directly. After
+// that, all functions and structs in this package can be moved to package
+// weightedtarget and unexported.
+package weightedaggregator
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/wrr"
+)
+
+type weightedPickerState struct {
+	weight uint32
+	state  balancer.State
+	// stateToAggregate is the connectivity state used only for state
+	// aggregation. It could be different from state.ConnectivityState. For
+	// example when a sub-balancer transitions from TransientFailure to
+	// connecting, state.ConnectivityState is Connecting, but stateToAggregate
+	// is still TransientFailure.
+	stateToAggregate connectivity.State
+}
+
+func (s *weightedPickerState) String() string {
+	return fmt.Sprintf("weight:%v,picker:%p,state:%v,stateToAggregate:%v", s.weight, s.state.Picker, s.state.ConnectivityState, s.stateToAggregate)
+}
+
+// Aggregator is the weighted balancer state aggregator.
+type Aggregator struct {
+	cc     balancer.ClientConn
+	logger *grpclog.PrefixLogger
+	newWRR func() wrr.WRR
+
+	csEvltr *balancer.ConnectivityStateEvaluator
+
+	mu sync.Mutex
+	// If started is false, no updates should be sent to the parent cc. A closed
+	// sub-balancer could still send pickers to this aggregator. This makes sure
+	// that no updates will be forwarded to parent when the whole balancer group
+	// and states aggregator is closed.
+	started bool
+	// All balancer IDs exist as keys in this map, even if balancer group is not
+	// started.
+	//
+	// If an ID is not in map, it's either removed or never added.
+	idToPickerState map[string]*weightedPickerState
+	// Set when UpdateState call propagation is paused.
+	pauseUpdateState bool
+	// Set when UpdateState call propagation is paused and an UpdateState call
+	// is suppressed.
+	needUpdateStateOnResume bool
+}
+
+// New creates a new weighted balancer state aggregator.
+func New(cc balancer.ClientConn, logger *grpclog.PrefixLogger, newWRR func() wrr.WRR) *Aggregator {
+	return &Aggregator{
+		cc:              cc,
+		logger:          logger,
+		newWRR:          newWRR,
+		csEvltr:         &balancer.ConnectivityStateEvaluator{},
+		idToPickerState: make(map[string]*weightedPickerState),
+	}
+}
+
+// Start starts the aggregator. It can be called after Stop to restart the
+// aggregator.
+func (wbsa *Aggregator) Start() {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	wbsa.started = true
+}
+
+// Stop stops the aggregator. When the aggregator is stopped, it won't call
+// parent ClientConn to update balancer state.
+func (wbsa *Aggregator) Stop() {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	wbsa.started = false
+	wbsa.clearStates()
+}
+
+// Add adds a sub-balancer state with weight. It adds a place holder, and waits for
+// the real sub-balancer to update state.
+func (wbsa *Aggregator) Add(id string, weight uint32) {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	wbsa.idToPickerState[id] = &weightedPickerState{
+		weight: weight,
+		// Start everything in CONNECTING, so if one of the sub-balancers
+		// reports TransientFailure, the RPCs will still wait for the other
+		// sub-balancers.
+		state: balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            base.NewErrPicker(balancer.ErrNoSubConnAvailable),
+		},
+		stateToAggregate: connectivity.Connecting,
+	}
+	wbsa.csEvltr.RecordTransition(connectivity.Shutdown, connectivity.Connecting)
+
+	wbsa.buildAndUpdateLocked()
+}
+
+// Remove removes the sub-balancer state. Future updates from this sub-balancer,
+// if any, will be ignored.
+func (wbsa *Aggregator) Remove(id string) {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	if _, ok := wbsa.idToPickerState[id]; !ok {
+		return
+	}
+	// Setting the state of the deleted sub-balancer to Shutdown will get csEvltr
+	// to remove the previous state for any aggregated state evaluations.
+	// transitions to and from connectivity.Shutdown are ignored by csEvltr.
+	wbsa.csEvltr.RecordTransition(wbsa.idToPickerState[id].stateToAggregate, connectivity.Shutdown)
+	// Remove id and picker from picker map. This also results in future updates
+	// for this ID to be ignored.
+	delete(wbsa.idToPickerState, id)
+	wbsa.buildAndUpdateLocked()
+}
+
+// UpdateWeight updates the weight for the given id. Note that this doesn't
+// trigger an update to the parent ClientConn. The caller should decide when
+// it's necessary, and call BuildAndUpdate.
+func (wbsa *Aggregator) UpdateWeight(id string, newWeight uint32) {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	pState, ok := wbsa.idToPickerState[id]
+	if !ok {
+		return
+	}
+	pState.weight = newWeight
+}
+
+// PauseStateUpdates causes UpdateState calls to not propagate to the parent
+// ClientConn.  The last state will be remembered and propagated when
+// ResumeStateUpdates is called.
+func (wbsa *Aggregator) PauseStateUpdates() {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	wbsa.pauseUpdateState = true
+	wbsa.needUpdateStateOnResume = false
+}
+
+// ResumeStateUpdates will resume propagating UpdateState calls to the parent,
+// and call UpdateState on the parent if any UpdateState call was suppressed.
+func (wbsa *Aggregator) ResumeStateUpdates() {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	wbsa.pauseUpdateState = false
+	if wbsa.needUpdateStateOnResume {
+		wbsa.cc.UpdateState(wbsa.build())
+	}
+}
+
+// NeedUpdateStateOnResume sets the UpdateStateOnResume bool to true, letting a
+// picker update be sent once ResumeStateUpdates is called.
+func (wbsa *Aggregator) NeedUpdateStateOnResume() {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	wbsa.needUpdateStateOnResume = true
+}
+
+// UpdateState is called to report a balancer state change from sub-balancer.
+// It's usually called by the balancer group.
+//
+// It calls parent ClientConn's UpdateState with the new aggregated state.
+func (wbsa *Aggregator) UpdateState(id string, newState balancer.State) {
+	wbsa.mu.Lock()
+	defer wbsa.mu.Unlock()
+	state, ok := wbsa.idToPickerState[id]
+	if !ok {
+		// All state starts with an entry in pickStateMap. If ID is not in map,
+		// it's either removed, or never existed.
+		return
+	}
+
+	if !(state.state.ConnectivityState == connectivity.TransientFailure && newState.ConnectivityState == connectivity.Connecting) {
+		// If old state is TransientFailure, and new state is Connecting, don't
+		// update the state, to prevent the aggregated state from being always
+		// CONNECTING. Otherwise, stateToAggregate is the same as
+		// state.ConnectivityState.
+		wbsa.csEvltr.RecordTransition(state.stateToAggregate, newState.ConnectivityState)
+		state.stateToAggregate = newState.ConnectivityState
+	}
+	state.state = newState
+
+	wbsa.buildAndUpdateLocked()
+}
+
+// clearState Reset everything to init state (Connecting) but keep the entry in
+// map (to keep the weight).
+//
+// Caller must hold wbsa.mu.
+func (wbsa *Aggregator) clearStates() {
+	for _, pState := range wbsa.idToPickerState {
+		pState.state = balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            base.NewErrPicker(balancer.ErrNoSubConnAvailable),
+		}
+		pState.stateToAggregate = connectivity.Connecting
+	}
+}
+
+// buildAndUpdateLocked aggregates the connectivity states of the sub-balancers,
+// builds a new picker and sends an update to the parent ClientConn.
+//
+// Caller must hold wbsa.mu.
+func (wbsa *Aggregator) buildAndUpdateLocked() {
+	if !wbsa.started {
+		return
+	}
+	if wbsa.pauseUpdateState {
+		// If updates are paused, do not call UpdateState, but remember that we
+		// need to call it when they are resumed.
+		wbsa.needUpdateStateOnResume = true
+		return
+	}
+
+	wbsa.cc.UpdateState(wbsa.build())
+}
+
+// build combines sub-states into one.
+//
+// Caller must hold wbsa.mu.
+func (wbsa *Aggregator) build() balancer.State {
+	wbsa.logger.Infof("Child pickers with config: %+v", wbsa.idToPickerState)
+
+	if len(wbsa.idToPickerState) == 0 {
+		// This is the case when all sub-balancers are removed.
+		return balancer.State{
+			ConnectivityState: connectivity.TransientFailure,
+			Picker:            base.NewErrPicker(errors.New("weighted-target: no targets to pick from")),
+		}
+	}
+
+	// Make sure picker's return error is consistent with the aggregatedState.
+	pickers := make([]weightedPickerState, 0, len(wbsa.idToPickerState))
+
+	switch aggState := wbsa.csEvltr.CurrentState(); aggState {
+	case connectivity.Connecting:
+		return balancer.State{
+			ConnectivityState: aggState,
+			Picker:            base.NewErrPicker(balancer.ErrNoSubConnAvailable)}
+	case connectivity.TransientFailure:
+		// this means that all sub-balancers are now in TransientFailure.
+		for _, ps := range wbsa.idToPickerState {
+			pickers = append(pickers, *ps)
+		}
+		return balancer.State{
+			ConnectivityState: aggState,
+			Picker:            newWeightedPickerGroup(pickers, wbsa.newWRR)}
+	default:
+		for _, ps := range wbsa.idToPickerState {
+			if ps.stateToAggregate == connectivity.Ready {
+				pickers = append(pickers, *ps)
+			}
+		}
+		return balancer.State{
+			ConnectivityState: aggState,
+			Picker:            newWeightedPickerGroup(pickers, wbsa.newWRR)}
+	}
+
+}
+
+type weightedPickerGroup struct {
+	w wrr.WRR
+}
+
+// newWeightedPickerGroup takes pickers with weights, and groups them into one
+// picker.
+//
+// Note it only takes ready pickers. The map shouldn't contain non-ready
+// pickers.
+func newWeightedPickerGroup(readyWeightedPickers []weightedPickerState, newWRR func() wrr.WRR) *weightedPickerGroup {
+	w := newWRR()
+	for _, ps := range readyWeightedPickers {
+		w.Add(ps.state.Picker, int64(ps.weight))
+	}
+
+	return &weightedPickerGroup{
+		w: w,
+	}
+}
+
+func (pg *weightedPickerGroup) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	p, ok := pg.w.Next().(balancer.Picker)
+	if !ok {
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+	}
+	return p.Pick(info)
+}

balancer/weightedtarget/weightedtarget.go

@@ -0,0 +1,196 @@
+/*
+ *
+ * Copyright 2020 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 weightedtarget implements the weighted_target balancer.
+//
+// All APIs in this package are experimental.
+package weightedtarget
+
+import (
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/weightedtarget/weightedaggregator"
+	"google.golang.org/grpc/internal/balancergroup"
+	"google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/hierarchy"
+	"google.golang.org/grpc/internal/pretty"
+	"google.golang.org/grpc/internal/wrr"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// Name is the name of the weighted_target balancer.
+const Name = "weighted_target_experimental"
+
+// NewRandomWRR is the WRR constructor used to pick sub-pickers from
+// sub-balancers. It's to be modified in tests.
+var NewRandomWRR = wrr.NewRandom
+
+func init() {
+	balancer.Register(bb{})
+}
+
+type bb struct{}
+
+func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+	b := &weightedTargetBalancer{}
+	b.logger = prefixLogger(b)
+	b.stateAggregator = weightedaggregator.New(cc, b.logger, NewRandomWRR)
+	b.stateAggregator.Start()
+	b.bg = balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               bOpts,
+		StateAggregator:         b.stateAggregator,
+		Logger:                  b.logger,
+		SubBalancerCloseTimeout: time.Duration(0), // Disable caching of removed child policies
+	})
+	b.logger.Infof("Created")
+	return b
+}
+
+func (bb) Name() string {
+	return Name
+}
+
+func (bb) ParseConfig(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	return parseConfig(c)
+}
+
+type weightedTargetBalancer struct {
+	logger *grpclog.PrefixLogger
+
+	bg              *balancergroup.BalancerGroup
+	stateAggregator *weightedaggregator.Aggregator
+
+	targets map[string]Target
+}
+
+type localityKeyType string
+
+const localityKey = localityKeyType("locality")
+
+// LocalityFromResolverState returns the locality from the resolver.State
+// provided, or an empty string if not present.
+func LocalityFromResolverState(state resolver.State) string {
+	locality, _ := state.Attributes.Value(localityKey).(string)
+	return locality
+}
+
+// UpdateClientConnState takes the new targets in balancer group,
+// creates/deletes sub-balancers and sends them update. addresses are split into
+// groups based on hierarchy path.
+func (b *weightedTargetBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
+	if b.logger.V(2) {
+		b.logger.Infof("Received update from resolver, balancer config: %+v", pretty.ToJSON(s.BalancerConfig))
+	}
+
+	newConfig, ok := s.BalancerConfig.(*LBConfig)
+	if !ok {
+		return fmt.Errorf("unexpected balancer config with type: %T", s.BalancerConfig)
+	}
+	endpointsSplit := hierarchy.Group(s.ResolverState.Endpoints)
+
+	b.stateAggregator.PauseStateUpdates()
+	defer b.stateAggregator.ResumeStateUpdates()
+
+	// Remove sub-pickers and sub-balancers that are not in the new config.
+	for name := range b.targets {
+		if _, ok := newConfig.Targets[name]; !ok {
+			b.stateAggregator.Remove(name)
+			b.bg.Remove(name)
+		}
+	}
+
+	// For sub-balancers in the new config
+	// - if it's new. add to balancer group,
+	// - if it's old, but has a new weight, update weight in balancer group.
+	//
+	// For all sub-balancers, forward the address/balancer config update.
+	for name, newT := range newConfig.Targets {
+		oldT, ok := b.targets[name]
+		if !ok {
+			// If this is a new sub-balancer, add weights to the picker map.
+			b.stateAggregator.Add(name, newT.Weight)
+			// Then add to the balancer group.
+			b.bg.Add(name, balancer.Get(newT.ChildPolicy.Name))
+			// Not trigger a state/picker update. Wait for the new sub-balancer
+			// to send its updates.
+		} else if newT.ChildPolicy.Name != oldT.ChildPolicy.Name {
+			// If the child policy name is different, remove from balancer group
+			// and re-add.
+			b.stateAggregator.Remove(name)
+			b.bg.Remove(name)
+			b.stateAggregator.Add(name, newT.Weight)
+			b.bg.Add(name, balancer.Get(newT.ChildPolicy.Name))
+		} else if newT.Weight != oldT.Weight {
+			// If this is an existing sub-balancer, update weight if necessary.
+			b.stateAggregator.UpdateWeight(name, newT.Weight)
+		}
+
+		// Forwards all the update:
+		// - addresses are from the map after splitting with hierarchy path,
+		// - Top level service config and attributes are the same,
+		// - Balancer config comes from the targets map.
+		//
+		// TODO: handle error? How to aggregate errors and return?
+		_ = b.bg.UpdateClientConnState(name, balancer.ClientConnState{
+			ResolverState: resolver.State{
+				Endpoints:     endpointsSplit[name],
+				ServiceConfig: s.ResolverState.ServiceConfig,
+				Attributes:    s.ResolverState.Attributes.WithValue(localityKey, name),
+			},
+			BalancerConfig: newT.ChildPolicy.Config,
+		})
+	}
+
+	b.targets = newConfig.Targets
+
+	// If the targets length is zero, it means we have removed all child
+	// policies from the balancer group and aggregator.
+	// At the start of this UpdateClientConnState() operation, a call to
+	// b.stateAggregator.ResumeStateUpdates() is deferred. Thus, setting the
+	// needUpdateStateOnResume bool to true here will ensure a new picker is
+	// built as part of that deferred function. Since there are now no child
+	// policies, the aggregated connectivity state reported form the Aggregator
+	// will be TRANSIENT_FAILURE.
+	if len(b.targets) == 0 {
+		b.stateAggregator.NeedUpdateStateOnResume()
+	}
+
+	return nil
+}
+
+func (b *weightedTargetBalancer) ResolverError(err error) {
+	b.bg.ResolverError(err)
+}
+
+func (b *weightedTargetBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+func (b *weightedTargetBalancer) Close() {
+	b.stateAggregator.Stop()
+	b.bg.Close()
+}
+
+func (b *weightedTargetBalancer) ExitIdle() {
+	b.bg.ExitIdle()
+}

balancer/weightedtarget/weightedtarget_config.go

@@ -0,0 +1,49 @@
+/*
+ *
+ * Copyright 2020 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 weightedtarget
+
+import (
+	"encoding/json"
+
+	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// Target represents one target with the weight and the child policy.
+type Target struct {
+	// Weight is the weight of the child policy.
+	Weight uint32 `json:"weight,omitempty"`
+	// ChildPolicy is the child policy and it's config.
+	ChildPolicy *internalserviceconfig.BalancerConfig `json:"childPolicy,omitempty"`
+}
+
+// LBConfig is the balancer config for weighted_target.
+type LBConfig struct {
+	serviceconfig.LoadBalancingConfig `json:"-"`
+
+	Targets map[string]Target `json:"targets,omitempty"`
+}
+
+func parseConfig(c json.RawMessage) (*LBConfig, error) {
+	var cfg LBConfig
+	if err := json.Unmarshal(c, &cfg); err != nil {
+		return nil, err
+	}
+	return &cfg, nil
+}

balancer/weightedtarget/weightedtarget_config_test.go

@@ -0,0 +1,104 @@
+/*
+ *
+ * Copyright 2020 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 weightedtarget
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/balancer"
+	_ "google.golang.org/grpc/balancer/grpclb"
+	"google.golang.org/grpc/balancer/roundrobin"
+	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+)
+
+const (
+	testJSONConfig = `{
+  "targets": {
+	"cluster_1": {
+	  "weight": 75,
+	  "childPolicy": [{
+        "grpclb": {
+          "childPolicy": [{"pick_first":{}}],
+          "targetName": "foo-service"
+        }
+      }]
+	},
+	"cluster_2": {
+	  "weight": 25,
+	  "childPolicy": [{"round_robin": ""}]
+	}
+  }
+}`
+)
+
+var (
+	grpclbConfigParser = balancer.Get("grpclb").(balancer.ConfigParser)
+	grpclbConfigJSON   = `{"childPolicy": [{"pick_first":{}}], "targetName": "foo-service"}`
+	grpclbConfig, _    = grpclbConfigParser.ParseConfig([]byte(grpclbConfigJSON))
+)
+
+func (s) TestParseConfig(t *testing.T) {
+	tests := []struct {
+		name    string
+		js      string
+		want    *LBConfig
+		wantErr bool
+	}{
+		{
+			name:    "empty json",
+			js:      "",
+			want:    nil,
+			wantErr: true,
+		},
+		{
+			name: "OK",
+			js:   testJSONConfig,
+			want: &LBConfig{
+				Targets: map[string]Target{
+					"cluster_1": {
+						Weight: 75,
+						ChildPolicy: &internalserviceconfig.BalancerConfig{
+							Name:   "grpclb",
+							Config: grpclbConfig,
+						},
+					},
+					"cluster_2": {
+						Weight: 25,
+						ChildPolicy: &internalserviceconfig.BalancerConfig{
+							Name: roundrobin.Name,
+						},
+					},
+				},
+			},
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := parseConfig([]byte(tt.js))
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("parseConfig() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if !cmp.Equal(got, tt.want) {
+				t.Errorf("parseConfig() got unexpected result, diff: %v", cmp.Diff(got, tt.want))
+			}
+		})
+	}
+}