feat(bigtable): add load balancing penalty for channel (#14149)

Author: sushanb

bigtable/internal/transport/connpool.go

@@ -35,9 +35,11 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
 	gtransport "google.golang.org/api/transport/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/alts"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
 
 	btopt "cloud.google.com/go/bigtable/internal/option"
 
@@ -48,6 +50,8 @@ import (
 // We cap the max draining timeout to 30mins as there might be a long running stream (such as full table scan).
 var maxDrainingTimeout = 30 * time.Minute
 
+const artificialLoadIfError = 10
+const artificialLoadPenalizedTimer = 5 * time.Second
 const requestParamsHeader = "x-goog-request-params"
 
 // ipProtocol represents the type of IP protocol used.
@@ -250,6 +254,7 @@ type connEntry struct {
 	streamingLoad atomic.Int32 // In-flight streaming requests
 	errorCount    atomic.Int64 // Errors since the last metric report
 	drainingState atomic.Bool  // True if the connection is being gracefully drained.
+	penaltyExpiry atomic.Int64 // penaltyExpiry stores the UnixNano timestamp of when the penalty ends
 
 }
 
@@ -271,6 +276,25 @@ func (e *connEntry) createdAt() int64 {
 	return e.conn.creationTime()
 }
 
+// applyErrorPenalty checks if the error warrants a load balancing penalty,
+// and if so, sets an expiration time for the artificial load.
+func (e *connEntry) applyErrorPenalty(err error) {
+	if err == nil {
+		return
+	}
+
+	code := status.Code(err)
+
+	// Penalize errors that typically indicate target-specific health or capacity issues.
+	if code == codes.Unavailable ||
+		code == codes.ResourceExhausted ||
+		code == codes.Internal {
+		// A simple Store is safe here; concurrent updates is fine here.
+		newExpiry := time.Now().Add(artificialLoadPenalizedTimer).UnixNano()
+		e.penaltyExpiry.Store(newExpiry)
+	}
+}
+
 // isDraining atomically checks if the connection is in the draining state.
 func (e *connEntry) isDraining() bool {
 	return e.drainingState.Load()
@@ -313,7 +337,18 @@ func (p *BigtableChannelPool) waitForDrainAndClose(entry *connEntry) {
 func (e *connEntry) calculateConnLoad() int32 {
 	unary := e.unaryLoad.Load()
 	streaming := e.streamingLoad.Load()
-	return unary + streaming
+	load := unary + streaming
+
+	expiry := e.penaltyExpiry.Load()
+	if expiry > 0 {
+		if time.Now().UnixNano() < expiry {
+			load += artificialLoadIfError // Apply the artificial penalty weight
+		} else {
+			// restore to zero
+			e.penaltyExpiry.CompareAndSwap(expiry, 0)
+		}
+	}
+	return load
 }
 
 // BigtableChannelPool implements ConnPool and routes requests to the connection
@@ -673,6 +708,7 @@ func (p *BigtableChannelPool) Invoke(ctx context.Context, method string, args in
 	err = entry.conn.Invoke(ctx, method, args, reply, opts...)
 	if err != nil {
 		entry.errorCount.Add(1)
+		entry.applyErrorPenalty(err) // Apply penalty on error
 	}
 	return err
 
@@ -823,6 +859,7 @@ func (s *refCountedStream) SendMsg(m interface{}) error {
 	err := s.ClientStream.SendMsg(m)
 	if err != nil {
 		s.entry.errorCount.Add(1)
+		s.entry.applyErrorPenalty(err)
 		s.decrementLoad()
 	}
 	return err
@@ -835,6 +872,7 @@ func (s *refCountedStream) RecvMsg(m interface{}) error {
 		// io.EOF is a normal stream termination, not an error to be counted.
 		if !errors.Is(err, io.EOF) {
 			s.entry.errorCount.Add(1)
+			s.entry.applyErrorPenalty(err)
 		}
 		s.decrementLoad()
 	}

bigtable/internal/transport/connpool_test.go

@@ -689,6 +689,53 @@ func TestNewBigtableChannelPool(t *testing.T) {
 	})
 }
 
+func TestConnEntryCalculateConnLoadWithPenalty(t *testing.T) {
+	entry := &connEntry{}
+
+	// 5 load
+	entry.unaryLoad.Store(2)
+	entry.streamingLoad.Store(3)
+
+	t.Run("NoPenalty", func(t *testing.T) {
+		entry.penaltyExpiry.Store(0)
+		if load := entry.calculateConnLoad(); load != 5 {
+			t.Errorf("Expected load of 5, got %d", load)
+		}
+	})
+
+	t.Run("ActivePenalty", func(t *testing.T) {
+		// Set expiration 5 seconds into the future
+		future := time.Now().Add(5 * time.Second).UnixNano()
+		entry.penaltyExpiry.Store(future)
+
+		// Expected: 5 (base) + 10 (penalty) = 15
+		if load := entry.calculateConnLoad(); load != 15 {
+			t.Errorf("Expected load of 15 due to active penalty, got %d", load)
+		}
+
+		// Verify the timestamp wasn't accidentally cleared
+		if entry.penaltyExpiry.Load() == 0 {
+			t.Errorf("Active penalty timestamp was incorrectly cleared")
+		}
+	})
+
+	t.Run("ExpiredPenaltySelfCleans", func(t *testing.T) {
+		// Set expiration 1 second in the past
+		past := time.Now().Add(-1 * time.Second).UnixNano()
+		entry.penaltyExpiry.Store(past)
+
+		// Expected: 5 (penalty is ignored because it expired)
+		if load := entry.calculateConnLoad(); load != 5 {
+			t.Errorf("Expected load of 5 due to expired penalty, got %d", load)
+		}
+
+		// Verify the lazy-evaluation reset the timestamp to 0
+		if expiry := entry.penaltyExpiry.Load(); expiry != 0 {
+			t.Errorf("Expected expired penalty timestamp to be reset to 0, but it remained %d", expiry)
+		}
+	})
+}
+
 func TestSelectLeastLoaded(t *testing.T) {
 	pool := &BigtableChannelPool{}
 	pool.conns.Store((&[]*connEntry{}))
@@ -983,6 +1030,44 @@ func TestPoolClose(t *testing.T) {
 	}
 }
 
+func TestConnEntryApplyErrorPenalty(t *testing.T) {
+	tests := []struct {
+		name        string
+		err         error
+		wantPenalty bool
+	}{
+		{"Nil error", nil, false},
+		{"Non-gRPC error", errors.New("standard network error"), false}, // Treated as codes.Unknown
+		{"OK", status.Error(codes.OK, "ok"), false},
+		{"NotFound", status.Error(codes.NotFound, "key not found"), false},
+		{"Canceled", status.Error(codes.Canceled, "client canceled"), false},
+		{"Unavailable", status.Error(codes.Unavailable, "server down"), true},
+		{"ResourceExhausted", status.Error(codes.ResourceExhausted, "too many requests"), true},
+		{"Internal", status.Error(codes.Internal, "proxy error"), true},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			entry := &connEntry{}
+
+			// Apply the error
+			entry.applyErrorPenalty(tc.err)
+
+			expiry := entry.penaltyExpiry.Load()
+
+			if tc.wantPenalty {
+				if expiry == 0 {
+					t.Errorf("Expected a penalty to be applied for %v, but expiry was 0", tc.err)
+				}
+			} else {
+				if expiry != 0 {
+					t.Errorf("Expected NO penalty for %v, but got expiry timestamp %d", tc.err, expiry)
+				}
+			}
+		})
+	}
+}
+
 func TestGracefulDraining(t *testing.T) {
 	ctx := context.Background()
 	fake := &fakeService{}