examples/features/loadbalancing: Add custom lb example

examples/features/load_balancing/README.md

@@ -81,3 +81,34 @@ Note that it's possible to see two continues RPC sent to the same backend.
 That's because `round_robin` only picks the connections ready for RPCs. So if
 one of the two connections is not ready for some reason, all RPCs will be sent
 to the ready connection.
+
+# custom_round_robin
+
+```
+this is examples/load_balancing (from :50052)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50052)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50052)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50052)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50052)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50052)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50051)
+this is examples/load_balancing (from :50052)
+this is examples/load_balancing (from :50051)
+```
+
+It is also possible to configure a custom load balancing policy through the
+service config. Specify the balancer type through the service config which you
+register in the registry. In this example, a custom balancer which routes to the
+first SubConn it receives except every n times, where n is configurable, in
+which is chooses the second SubConn it receives.

examples/features/load_balancing/client/main.go

@@ -21,14 +21,19 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log"
+	"sync/atomic"
 	"time"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
 	"google.golang.org/grpc/credentials/insecure"
 	ecpb "google.golang.org/grpc/examples/features/proto/echo"
 	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
 )
 
 const (
@@ -84,11 +89,24 @@ func main() {
 
 	fmt.Println("--- calling helloworld.Greeter/SayHello with round_robin ---")
 	makeRPCs(roundrobinConn, 10)
+	// You can also plug in your own custom lb policy, which needs to be
+	// configurable. This n is configurable. Try changing it and see how the
+	// behavior changes.
+	customroundrobinConn, err := grpc.Dial(
+		fmt.Sprintf("%s:///%s", exampleScheme, exampleServiceName),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"custom_round_robin":{"n": 3}}]}`),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
+	if err != nil {
+		log.Fatalf("did not connect: %v", err)
+	}
+	defer customroundrobinConn.Close()
+	fmt.Println("--- calling helloworld.Greeter/SayHello with custom_round_robin ---")
+	makeRPCs(customroundrobinConn, 20)
 }
 
 // Following is an example name resolver implementation. Read the name
 // resolution example to learn more about it.
-
 type exampleResolverBuilder struct{}
 
 func (*exampleResolverBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
@@ -123,4 +141,87 @@ func (*exampleResolver) Close()                                  {}
 
 func init() {
 	resolver.Register(&exampleResolverBuilder{})
+	balancer.Register(&customRoundRobinBuilder{})
+}
+
+const customRRName = "custom_round_robin"
+
+type customRRConfig struct {
+	serviceconfig.LoadBalancingConfig `json:"-"`
+
+	// N represents how often pick iterations chose the second SubConn
+	// in the list. Defaults to 3.
+	N uint32 `json:"n,omitempty"`
+}
+
+func (customRoundRobinBuilder) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	lbConfig := &customRRConfig{
+		N: 2,
+	}
+	if err := json.Unmarshal(s, lbConfig); err != nil {
+		return nil, fmt.Errorf("custom-round-robin: unable to unmarshal customRRConfig: %v", err)
+	}
+	return lbConfig, nil
+}
+
+type customRoundRobinBuilder struct{}
+
+func (customRoundRobinBuilder) Name() string {
+	return customRRName
+}
+
+func (customRoundRobinBuilder) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+	crr := &customRoundRobin{}
+	baseBuilder := base.NewBalancerBuilder(customRRName, crr, base.Config{HealthCheck: true})
+	crr.Balancer = baseBuilder.Build(cc, bOpts)
+	return crr
+}
+
+type customRoundRobin struct {
+	// Embeds balancer.Balancer because needs to intercept UpdateClientConnState
+	// to learn about N.
+	balancer.Balancer
+	n uint32
+}
+
+func (crr *customRoundRobin) UpdateClientConnState(state balancer.ClientConnState) error {
+	crrCfg, ok := state.BalancerConfig.(*customRRConfig)
+	if !ok {
+		return balancer.ErrBadResolverState
+	}
+	crr.n = crrCfg.N
+	return crr.Balancer.UpdateClientConnState(state)
+}
+
+func (crr *customRoundRobin) Build(info base.PickerBuildInfo) balancer.Picker {
+	if len(info.ReadySCs) == 0 {
+		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
+	}
+
+	scs := make([]balancer.SubConn, 0, len(info.ReadySCs))
+	for sc := range info.ReadySCs {
+		scs = append(scs, sc)
+	}
+	return &customRoundRobinPicker{
+		subConns: scs,
+		n:        crr.n,
+		next:     0,
+	}
+}
+
+type customRoundRobinPicker struct {
+	subConns []balancer.SubConn
+	n        uint32
+
+	next uint32
+}
+
+func (crrp *customRoundRobinPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	next := atomic.AddUint32(&crrp.next, 1)
+	index := 0
+	if next%crrp.n == 0 {
+		index = 1
+	}
+	sc := crrp.subConns[index%len(crrp.subConns)]
+	return balancer.PickResult{SubConn: sc}, nil
 }