rework for A68

Author: marek-szews

balancer/randomsubsetting/randomsubsetting.go

@@ -24,10 +24,10 @@
 package randomsubsetting
 
 import (
+	"cmp"
 	"encoding/json"
-	"errors"
 	"fmt"
-	"sort"
+	"slices"
 	"time"
 
 	"github.com/cespare/xxhash/v2"
@@ -72,28 +72,38 @@ func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Ba
 	return b
 }
 
-// LBConfig is the config for the outlier detection balancer.
+// LBConfig is the config for the random subsetting balancer.
 type LBConfig struct {
 	serviceconfig.LoadBalancingConfig `json:"-"`
 
-	SubsetSize uint64 `json:"subset_size,omitempty"`
-
+	SubsetSize  uint64                         `json:"subset_size,omitempty"`
 	ChildPolicy *iserviceconfig.BalancerConfig `json:"child_policy,omitempty"`
 }
 
 func (bb) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 	lbCfg := &LBConfig{
-		// Default top layer values.
-		SubsetSize: 10,
+		SubsetSize:  2, // default value
+		ChildPolicy: &iserviceconfig.BalancerConfig{Name: "round_robin"},
 	}
 
 	if err := json.Unmarshal(s, lbCfg); err != nil { // Validates child config if present as well.
-		return nil, fmt.Errorf("subsetting: unable to unmarshal LBconfig: %s, error: %v", string(s), err)
+		return nil, fmt.Errorf("randomsubsetting: unable to unmarshal LBConfig: %s, error: %v", string(s), err)
 	}
 
-	// if someonw needs subsetSize == 1, he should use pick_first instead
+	// if someone needs SubsetSize == 1, he should use pick_first instead
 	if lbCfg.SubsetSize < 2 {
-		return nil, errors.New("subsetting: subsetSize must be >= 2")
+		return nil, fmt.Errorf("randomsubsetting: SubsetSize must be >= 2")
+	}
+
+	if lbCfg.ChildPolicy == nil {
+		return nil, fmt.Errorf("randomsubsetting: child policy field must be set")
+	}
+
+	// Reject whole config if child policy doesn't exist, don't persist it for
+	// later.
+	bb := balancer.Get(lbCfg.ChildPolicy.Name)
+	if bb == nil {
+		return nil, fmt.Errorf("randomsubsetting: child balancer %q not registered", lbCfg.ChildPolicy.Name)
 	}
 
 	return lbCfg, nil
@@ -114,73 +124,68 @@ type subsettingBalancer struct {
 func (b *subsettingBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 	lbCfg, ok := s.BalancerConfig.(*LBConfig)
 	if !ok {
-		b.logger.Errorf("received config with unexpected type %T: %v", s.BalancerConfig, s.BalancerConfig)
+		b.logger.Errorf("randomsubsetting: received config with unexpected type %T: %v", s.BalancerConfig, s.BalancerConfig)
 		return balancer.ErrBadResolverState
 	}
 
-	// Reject whole config if child policy doesn't exist, don't persist it for
-	// later.
-	bb := balancer.Get(lbCfg.ChildPolicy.Name)
-	if bb == nil {
-		return fmt.Errorf("subsetting: child balancer %q not registered", lbCfg.ChildPolicy.Name)
-	}
-
 	if b.cfg == nil || b.cfg.ChildPolicy.Name != lbCfg.ChildPolicy.Name {
-		err := b.child.SwitchTo(bb)
-		if err != nil {
-			return fmt.Errorf("subsetting: error switching to child of type %q: %v", lbCfg.ChildPolicy.Name, err)
+
+		if err := b.child.SwitchTo(balancer.Get(lbCfg.ChildPolicy.Name)); err != nil {
+			return fmt.Errorf("randomsubsetting: error switching to child of type %q: %v", lbCfg.ChildPolicy.Name, err)
 		}
 	}
 	b.cfg = lbCfg
 
-	err := b.child.UpdateClientConnState(balancer.ClientConnState{
-		ResolverState:  b.prepareChildResolverState(s.ResolverState),
+	return b.child.UpdateClientConnState(balancer.ClientConnState{
+		ResolverState:  b.prepareChildResolverState(s),
 		BalancerConfig: b.cfg.ChildPolicy.Config,
 	})
-
-	return err
 }
 
-type AddressWithHash struct {
+type endpointWithHash struct {
 	hash uint64
-	addr resolver.Address
+	ep   resolver.Endpoint
 }
 
-// implements the subsetting algorithm, as described in A68: https://github.com/grpc/proposal/pull/423
-func (b *subsettingBalancer) prepareChildResolverState(s resolver.State) resolver.State {
-	addresses := s.Addresses
-	backendCount := len(addresses)
-	if backendCount <= int(b.cfg.SubsetSize) {
-		return s
+// implements the subsetting algorithm,
+// as described in A68: https://github.com/grpc/proposal/blob/master/A68-random-subsetting.md
+func (b *subsettingBalancer) prepareChildResolverState(s balancer.ClientConnState) resolver.State {
+	subsetSize := b.cfg.SubsetSize
+	endPoints := s.ResolverState.Endpoints
+	backendCount := len(endPoints)
+	if backendCount <= int(subsetSize) || subsetSize < 2 {
+		return s.ResolverState
 	}
 
-	addressesSet := make([]AddressWithHash, backendCount)
 	// calculate hash for each endpoint
-	for i, endpoint := range addresses {
-
-		b.hashf.Write([]byte(s.Addresses[0].String()))
-		addressesSet[i] = AddressWithHash{
+	endpointSet := make([]endpointWithHash, backendCount)
+	for i, endpoint := range endPoints {
+		b.hashf.Write([]byte(endpoint.Addresses[0].String()))
+		endpointSet[i] = endpointWithHash{
 			hash: b.hashf.Sum64(),
-			addr: endpoint,
+			ep:   endpoint,
 		}
 	}
-	// sort addresses by hash
-	sort.Slice(addressesSet, func(i, j int) bool {
-		return addressesSet[i].hash < addressesSet[j].hash
+
+	// sort endpoint by hash
+	slices.SortFunc(endpointSet, func(a, b endpointWithHash) int {
+		return cmp.Compare(a.hash, b.hash)
 	})
 
-	b.logger.Infof("resulting subset: %v", addressesSet[:b.cfg.SubsetSize])
+	if b.logger.V(2) {
+		b.logger.Infof("randomsubsetting: resulting subset: %v", endpointSet[:subsetSize])
+	}
 
-	// Convert back to resolver.addresses
-	addressesSubset := make([]resolver.Address, b.cfg.SubsetSize)
-	for _, eh := range addressesSet[:b.cfg.SubsetSize] {
-		addressesSubset = append(addressesSubset, eh.addr)
+	// Convert back to resolver.Endpoints
+	endpointSubset := make([]resolver.Endpoint, subsetSize)
+	for i, endpoint := range endpointSet[:subsetSize] {
+		endpointSubset[i] = endpoint.ep
 	}
 
 	return resolver.State{
-		Addresses:     addressesSubset,
-		ServiceConfig: s.ServiceConfig,
-		Attributes:    s.Attributes,
+		Endpoints:     endpointSubset,
+		ServiceConfig: s.ResolverState.ServiceConfig,
+		Attributes:    s.ResolverState.Attributes,
 	}
 }
 

balancer/randomsubsetting/randomsubsetting_test.go

@@ -17,15 +17,18 @@
  */
 
 // Package e2e_test contains e2e test cases for the Subsetting LB Policy.
-package randomsubsetting_test
+package randomsubsetting
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"math"
+	"strings"
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal"
@@ -36,13 +39,12 @@ import (
 	"google.golang.org/grpc/resolver/manual"
 	"google.golang.org/grpc/serviceconfig"
 
+	iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
 	testgrpc "google.golang.org/grpc/interop/grpc_testing"
 	testpb "google.golang.org/grpc/interop/grpc_testing"
-
-	_ "google.golang.org/grpc/balancer/randomsubsetting"
 )
 
-var defaultTestTimeout = 5 * time.Second
+var defaultTestTimeout = 120 * time.Second
 
 type s struct {
 	grpctest.Tester
@@ -52,6 +54,66 @@ 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{
+				SubsetSize:  2,
+				ChildPolicy: &iserviceconfig.BalancerConfig{Name: "round_robin"},
+			},
+		},
+		{
+			name:  "happy-case-subset_size-set",
+			input: `{ "subset_size": 3 }`,
+			wantCfg: &LBConfig{
+				SubsetSize:  3,
+				ChildPolicy: &iserviceconfig.BalancerConfig{Name: "round_robin"},
+			},
+		},
+		{
+			name: "subset_size-less-than-2",
+			input: `{ "subset_size": 1,
+					  "child_policy": [{"round_robin": {}}]}`,
+			wantErr: "randomsubsetting: SubsetSize 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 setupBackends(t *testing.T, backendsCount int) ([]resolver.Address, func()) {
 	t.Helper()
 
@@ -123,7 +185,7 @@ func setupClients(t *testing.T, clientsCount int, subsetSize int, addresses []re
 	return clients, cancel
 }
 
-func checkRoundRobinRPCs(t *testing.T, ctx context.Context, clients []testgrpc.TestServiceClient, subsetSize int, maxDiff int) {
+func checkRoundRobinRPCs(ctx context.Context, t *testing.T, clients []testgrpc.TestServiceClient, subsetSize int, maxDiff int) {
 	clientsPerBackend := map[string]map[int]struct{}{}
 
 	for clientIdx, client := range clients {
@@ -176,33 +238,40 @@ func (s) TestSubsettingE2E(t *testing.T) {
 		backends   int
 		maxDiff    int
 	}{
+		{
+			name:       "backends could be evenly distributed between small number of clients",
+			backends:   3,
+			clients:    2,
+			subsetSize: 2,
+			maxDiff:    1,
+		},
 		{
 			name:       "backends could be evenly distributed between clients",
 			backends:   12,
 			clients:    8,
 			subsetSize: 3,
-			maxDiff:    0,
+			maxDiff:    3,
 		},
 		{
 			name:       "backends could NOT be evenly distributed between clients",
 			backends:   37,
 			clients:    22,
 			subsetSize: 5,
-			maxDiff:    2,
+			maxDiff:    15,
 		},
 		{
 			name:       "Nbackends %% subsetSize == 0, but there are not enough clients to fill the last round",
 			backends:   20,
 			clients:    7,
 			subsetSize: 5,
-			maxDiff:    1,
+			maxDiff:    20,
 		},
 		{
 			name:       "last round is completely filled, but there are some excluded backends on every round",
 			backends:   21,
 			clients:    8,
 			subsetSize: 5,
-			maxDiff:    1,
+			maxDiff:    3,
 		},
 	}
 	for _, test := range tests {
@@ -216,7 +285,7 @@ func (s) TestSubsettingE2E(t *testing.T) {
 			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 			defer cancel()
 
-			checkRoundRobinRPCs(t, ctx, clients, test.subsetSize, test.maxDiff)
+			checkRoundRobinRPCs(ctx, t, clients, test.subsetSize, test.maxDiff)
 		})
 	}
 }