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

examples/features/customloadbalancer/client/customroundrobin/customroundrobin.go

@@ -0,0 +1,148 @@
+/*
+ *
+ * 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 customroundrobin provides an example for the custom roundrobin balancer.
+package customroundrobin
+
+import (
+	"encoding/json"
+	"fmt"
+	"sync/atomic"
+
+	_ "google.golang.org/grpc" // to register pick_first
+	"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/serviceconfig"
+)
+
+func init() {
+	balancer.Register(customRoundRobinBuilder{})
+}
+
+const customRRName = "custom_round_robin"
+
+type customRRConfig struct {
+	serviceconfig.LoadBalancingConfig `json:"-"`
+
+	// ChooseSecond represents how often pick iterations choose the second
+	// SubConn in the list. Defaults to 3. If 0 never choose the second SubConn.
+	ChooseSecond uint32 `json:"chooseSecond,omitempty"`
+}
+
+type customRoundRobinBuilder struct{}
+
+func (customRoundRobinBuilder) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	lbConfig := &customRRConfig{
+		ChooseSecond: 3,
+	}
+
+	if err := json.Unmarshal(s, lbConfig); err != nil {
+		return nil, fmt.Errorf("custom-round-robin: unable to unmarshal customRRConfig: %v", err)
+	}
+	return lbConfig, nil
+}
+
+func (customRoundRobinBuilder) Name() string {
+	return customRRName
+}
+
+func (customRoundRobinBuilder) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+	crr := &customRoundRobin{
+		ClientConn: cc,
+		bOpts:      bOpts,
+	}
+	crr.Balancer = endpointsharding.NewBalancer(crr, bOpts, balancer.Get(pickfirst.Name).Build, endpointsharding.Options{})
+	return crr
+}
+
+type customRoundRobin struct {
+	// All state and operations on this balancer are either initialized at build
+	// time and read only after, or are only accessed as part of its
+	// balancer.Balancer API (UpdateState from children only comes in from
+	// balancer.Balancer calls as well, and children are called one at a time),
+	// in which calls are guaranteed to come synchronously. Thus, no extra
+	// synchronization is required in this balancer.
+	balancer.Balancer
+	balancer.ClientConn
+	bOpts balancer.BuildOptions
+
+	cfg atomic.Pointer[customRRConfig]
+}
+
+func (crr *customRoundRobin) UpdateClientConnState(state balancer.ClientConnState) error {
+	crrCfg, ok := state.BalancerConfig.(*customRRConfig)
+	if !ok {
+		return balancer.ErrBadResolverState
+	}
+	if el := state.ResolverState.Endpoints; len(el) != 2 {
+		return fmt.Errorf("UpdateClientConnState wants two endpoints, got: %v", el)
+	}
+	crr.cfg.Store(crrCfg)
+	// A call to UpdateClientConnState should always produce a new Picker.  That
+	// is guaranteed to happen since the aggregator will always call
+	// UpdateChildState in its UpdateClientConnState.
+	return crr.Balancer.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState: state.ResolverState,
+	})
+}
+
+func (crr *customRoundRobin) UpdateState(state balancer.State) {
+	if state.ConnectivityState == connectivity.Ready {
+		childStates := endpointsharding.ChildStatesFromPicker(state.Picker)
+		var readyPickers []balancer.Picker
+		for _, childState := range childStates {
+			if childState.State.ConnectivityState == connectivity.Ready {
+				readyPickers = append(readyPickers, childState.State.Picker)
+			}
+		}
+		// If both children are ready, pick using the custom round robin
+		// algorithm.
+		if len(readyPickers) == 2 {
+			picker := &customRoundRobinPicker{
+				pickers:      readyPickers,
+				chooseSecond: crr.cfg.Load().ChooseSecond,
+				next:         0,
+			}
+			crr.ClientConn.UpdateState(balancer.State{
+				ConnectivityState: connectivity.Ready,
+				Picker:            picker,
+			})
+			return
+		}
+	}
+	// Delegate to default behavior/picker from below.
+	crr.ClientConn.UpdateState(state)
+}
+
+type customRoundRobinPicker struct {
+	pickers      []balancer.Picker
+	chooseSecond uint32
+	next         uint32
+}
+
+func (crrp *customRoundRobinPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	next := atomic.AddUint32(&crrp.next, 1)
+	index := 0
+	if next != 0 && next%crrp.chooseSecond == 0 {
+		index = 1
+	}
+	childPicker := crrp.pickers[index%len(crrp.pickers)]
+	return childPicker.Pick(info)
+}

examples/features/customloadbalancer/client/main.go

@@ -0,0 +1,155 @@
+/*
+ *
+ * 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.
+ *
+ */
+
+// Binary client is a client for the custom load balancer example.
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"strings"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	_ "google.golang.org/grpc/examples/features/customloadbalancer/client/customroundrobin" // To register custom_round_robin.
+	pb "google.golang.org/grpc/examples/features/proto/echo"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+var (
+	addr1 = "localhost:50050"
+	addr2 = "localhost:50051"
+)
+
+func main() {
+	mr := manual.NewBuilderWithScheme("example")
+	defer mr.Close()
+
+	// You can also plug in your own custom lb policy, which needs to be
+	// configurable. This n is configurable. Try changing n and see how the
+	// behavior changes.
+	json := `{"loadBalancingConfig": [{"custom_round_robin":{"chooseSecond": 3}}]}`
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(json)
+	mr.InitialState(resolver.State{
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{{Addr: addr1}}},
+			{Addresses: []resolver.Address{{Addr: addr2}}},
+		},
+		ServiceConfig: sc,
+	})
+
+	cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		log.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
+	defer cancel()
+	ec := pb.NewEchoClient(cc)
+	if err := waitForDistribution(ctx, ec); err != nil {
+		log.Fatal(err)
+	}
+	fmt.Println("Successful multiple iterations of 1:2 ratio")
+}
+
+// waitForDistribution makes RPC's on the echo client until 3 RPC's follow the
+// same 1:2 address ratio for the peer. Returns an error if fails to do so
+// before context timeout.
+func waitForDistribution(ctx context.Context, ec pb.EchoClient) error {
+	for {
+		results := make(map[string]uint32)
+	InnerLoop:
+		for {
+			if ctx.Err() != nil {
+				return fmt.Errorf("timeout waiting for 1:2 distribution between addresses %v and %v", addr1, addr2)
+			}
+
+			for i := 0; i < 3; i++ {
+				res := make(map[string]uint32)
+				for j := 0; j < 3; j++ {
+					var peer peer.Peer
+					r, err := ec.UnaryEcho(ctx, &pb.EchoRequest{Message: "this is examples/customloadbalancing"}, grpc.Peer(&peer))
+					if err != nil {
+						return fmt.Errorf("UnaryEcho failed: %v", err)
+					}
+					fmt.Println(r)
+					peerAddr := peer.Addr.String()
+					if !strings.HasSuffix(peerAddr, "50050") && !strings.HasSuffix(peerAddr, "50051") {
+						return fmt.Errorf("peer address was not one of %v or %v, got: %v", addr1, addr2, peerAddr)
+					}
+					res[peerAddr]++
+					time.Sleep(time.Millisecond)
+				}
+				// Make sure the addresses come in a 1:2 ratio for this
+				// iteration.
+				var seen1, seen2 bool
+				for addr, count := range res {
+					if count != 1 && count != 2 {
+						break InnerLoop
+					}
+					if count == 1 {
+						if seen1 {
+							break InnerLoop
+						}
+						seen1 = true
+					}
+					if count == 2 {
+						if seen2 {
+							break InnerLoop
+						}
+						seen2 = true
+					}
+					results[addr] = results[addr] + count
+				}
+				if !seen1 || !seen2 {
+					break InnerLoop
+				}
+			}
+			// Make sure iteration is 3 and 6 for addresses seen. This makes
+			// sure the distribution is the same 1:2 ratio for each iteration.
+			var seen3, seen6 bool
+			for _, count := range results {
+				if count != 3 && count != 6 {
+					break InnerLoop
+				}
+				if count == 3 {
+					if seen3 {
+						break InnerLoop
+					}
+					seen3 = true
+				}
+				if count == 6 {
+					if seen6 {
+						break InnerLoop
+					}
+					seen6 = true
+				}
+				return nil
+			}
+			if !seen3 || !seen6 {
+				break InnerLoop
+			}
+		}
+	}
+}