Implement A72: OpenTelemetry Tracing

balancer/endpointsharding/endpointsharding.go

@@ -28,19 +28,33 @@ package endpointsharding
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
+	rand "math/rand/v2"
 	"sync"
 	"sync/atomic"
 
-	rand "math/rand/v2"
-
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/internal/balancer/gracefulswitch"
+	"google.golang.org/grpc/internal/envconfig"
 	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/serviceconfig"
 )
 
+// PickFirstConfig is a pick first config without shuffling enabled.
+var PickFirstConfig string
+
+func init() {
+	name := pickfirst.Name
+	if !envconfig.NewPickFirstEnabled {
+		name = pickfirstleaf.Name
+	}
+	PickFirstConfig = fmt.Sprintf("[{%q: {}}]", name)
+}
+
 // ChildState is the balancer state of a child along with the endpoint which
 // identifies the child balancer.
 type ChildState struct {
@@ -100,9 +114,6 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 
 	// Update/Create new children.
 	for _, endpoint := range state.ResolverState.Endpoints {
-		if len(endpoint.Addresses) == 0 {
-			continue
-		}
 		if _, ok := newChildren.Get(endpoint); ok {
 			// Endpoint child was already created, continue to avoid duplicate
 			// update.
@@ -143,6 +154,9 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 		}
 	}
 	es.children.Store(newChildren)
+	if newChildren.Len() == 0 {
+		return balancer.ErrBadResolverState
+	}
 	return ret
 }
 
@@ -306,6 +320,3 @@ func (bw *balancerWrapper) UpdateState(state balancer.State) {
 func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 	return gracefulswitch.ParseConfig(cfg)
 }
-
-// PickFirstConfig is a pick first config without shuffling enabled.
-const PickFirstConfig = "[{\"pick_first\": {}}]"

balancer/grpclb/grpclb_test.go

@@ -460,7 +460,7 @@ func (s) TestGRPCLB_Basic(t *testing.T) {
 	}
 	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
 	defer cc.Close()
 
@@ -517,7 +517,7 @@ func (s) TestGRPCLB_Weighted(t *testing.T) {
 	}
 	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
 	defer cc.Close()
 
@@ -597,7 +597,7 @@ func (s) TestGRPCLB_DropRequest(t *testing.T) {
 	}
 	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
 	defer cc.Close()
 	testC := testgrpc.NewTestServiceClient(cc)
@@ -769,7 +769,7 @@ func (s) TestGRPCLB_BalancerDisconnects(t *testing.T) {
 	}
 	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
 	defer cc.Close()
 	testC := testgrpc.NewTestServiceClient(cc)
@@ -940,7 +940,7 @@ func (s) TestGRPCLB_ExplicitFallback(t *testing.T) {
 	}
 	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
 	defer cc.Close()
 	testC := testgrpc.NewTestServiceClient(cc)
@@ -1008,11 +1008,12 @@ func (s) TestGRPCLB_FallBackWithNoServerAddress(t *testing.T) {
 		grpc.WithTransportCredentials(&serverNameCheckCreds{}),
 		grpc.WithContextDialer(fakeNameDialer),
 	}
-	cc, err := grpc.Dial(r.Scheme()+":///"+beServerName, dopts...)
+	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
 	defer cc.Close()
+	cc.Connect()
 	testC := testgrpc.NewTestServiceClient(cc)
 
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
@@ -1102,10 +1103,11 @@ func (s) TestGRPCLB_PickFirst(t *testing.T) {
 		grpc.WithTransportCredentials(&serverNameCheckCreds{}),
 		grpc.WithContextDialer(fakeNameDialer),
 	}
-	cc, err := grpc.Dial(r.Scheme()+":///"+beServerName, dopts...)
+	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend: %v", err)
 	}
+	cc.Connect()
 	defer cc.Close()
 
 	// Push a service config with grpclb as the load balancing policy and
@@ -1198,7 +1200,7 @@ func (s) TestGRPCLB_BackendConnectionErrorPropagation(t *testing.T) {
 		grpc.WithTransportCredentials(&serverNameCheckCreds{}),
 		grpc.WithContextDialer(fakeNameDialer))
 	if err != nil {
-		t.Fatalf("Failed to create new client to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend: %v", err)
 	}
 	defer cc.Close()
 	testC := testgrpc.NewTestServiceClient(cc)
@@ -1241,10 +1243,11 @@ func testGRPCLBEmptyServerList(t *testing.T, svcfg string) {
 		grpc.WithTransportCredentials(&serverNameCheckCreds{}),
 		grpc.WithContextDialer(fakeNameDialer),
 	}
-	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName, dopts...)
+	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
+	cc.Connect()
 	defer cc.Close()
 	testC := testgrpc.NewTestServiceClient(cc)
 
@@ -1311,15 +1314,16 @@ func (s) TestGRPCLBWithTargetNameFieldInConfig(t *testing.T) {
 	// Push the backend address to the remote balancer.
 	tss.ls.sls <- sl
 
-	cc, err := grpc.Dial(r.Scheme()+":///"+beServerName,
+	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName,
 		grpc.WithResolvers(r),
 		grpc.WithTransportCredentials(&serverNameCheckCreds{}),
 		grpc.WithContextDialer(fakeNameDialer),
 		grpc.WithUserAgent(testUserAgent))
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
 	defer cc.Close()
+	cc.Connect()
 	testC := testgrpc.NewTestServiceClient(cc)
 
 	// Push a resolver update with grpclb configuration which does not contain the
@@ -1418,15 +1422,14 @@ func runAndCheckStats(t *testing.T, drop bool, statsChan chan *lbpb.ClientStats,
 	tss.ls.statsDura = 100 * time.Millisecond
 	creds := serverNameCheckCreds{}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
-	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName, grpc.WithResolvers(r),
+	cc, err := grpc.NewClient(r.Scheme()+":///"+beServerName, grpc.WithResolvers(r),
 		grpc.WithTransportCredentials(&creds),
 		grpc.WithPerRPCCredentials(failPreRPCCred{}),
 		grpc.WithContextDialer(fakeNameDialer))
 	if err != nil {
-		t.Fatalf("Failed to dial to the backend %v", err)
+		t.Fatalf("Failed to create a client for the backend %v", err)
 	}
+	cc.Connect()
 	defer cc.Close()
 
 	rstate := resolver.State{ServiceConfig: r.CC.ParseServiceConfig(grpclbConfig)}

balancer/pickfirst/pickfirst_ext_test.go

@@ -376,12 +376,12 @@ func (s) TestPickFirst_StickyTransientFailure(t *testing.T) {
 			},
 		}),
 	}
-	cc, err := grpc.Dial(lis.Addr().String(), dopts...)
+	cc, err := grpc.NewClient(lis.Addr().String(), dopts...)
 	if err != nil {
-		t.Fatalf("Failed to dial server at %q: %v", lis.Addr(), err)
+		t.Fatalf("Failed to create new client: %v", err)
 	}
 	t.Cleanup(func() { cc.Close() })
-
+	cc.Connect()
 	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
 
 	// Spawn a goroutine to ensure that the channel stays in TransientFailure.
@@ -837,12 +837,12 @@ func (s) TestPickFirst_ResolverError_WithPreviousUpdate_Connecting(t *testing.T)
 		grpc.WithResolvers(r),
 		grpc.WithDefaultServiceConfig(pickFirstServiceConfig),
 	}
-	cc, err := grpc.Dial(r.Scheme()+":///test.server", dopts...)
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
 	if err != nil {
-		t.Fatalf("grpc.Dial() failed: %v", err)
+		t.Fatalf("grpc.NewClient() failed: %v", err)
 	}
 	t.Cleanup(func() { cc.Close() })
-
+	cc.Connect()
 	addrs := []resolver.Address{{Addr: lis.Addr().String()}}
 	r.UpdateState(resolver.State{Addresses: addrs})
 	testutils.AwaitState(ctx, t, cc, connectivity.Connecting)
@@ -892,12 +892,12 @@ func (s) TestPickFirst_ResolverError_WithPreviousUpdate_TransientFailure(t *test
 		grpc.WithResolvers(r),
 		grpc.WithDefaultServiceConfig(pickFirstServiceConfig),
 	}
-	cc, err := grpc.Dial(r.Scheme()+":///test.server", dopts...)
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
 	if err != nil {
-		t.Fatalf("grpc.Dial() failed: %v", err)
+		t.Fatalf("grpc.NewClient() failed: %v", err)
 	}
 	t.Cleanup(func() { cc.Close() })
-
+	cc.Connect()
 	addrs := []resolver.Address{{Addr: lis.Addr().String()}}
 	r.UpdateState(resolver.State{Addresses: addrs})
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)

balancer/pickfirst/pickfirst_test.go

@@ -29,6 +29,7 @@ import (
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/internal/grpctest"
 	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/internal/testutils/stats"
 	"google.golang.org/grpc/resolver"
 )
 
@@ -55,7 +56,7 @@ 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{})
+	bal := balancer.Get(Name).Build(cc, balancer.BuildOptions{MetricsRecorder: &stats.NoopMetricsRecorder{}})
 	defer bal.Close()
 	bal.ResolverError(errors.New("resolution failed: test error"))
 
@@ -88,7 +89,7 @@ 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{})
+	bal := balancer.Get(Name).Build(cc, balancer.BuildOptions{MetricsRecorder: &stats.NoopMetricsRecorder{}})
 	defer bal.Close()
 
 	// After sending a valid update, the LB policy should report CONNECTING.