rls: only reset backoff on recovery from TRANSIENT_FAILURE (#9137)

Fixes #8693

This PR is a continuation and finalization of the stale/closed PR #8720.

It addresses the remaining feedback from the maintainers and the code
review tools to resolve control channel state monitoring issues in the
RLS balancer.

RELEASE NOTES: N/A

---------

Co-authored-by: ulascansenturk <ulascansenturk@protonmail.com>

Author: mswierq

balancer/rls/balancer_test.go

@@ -39,6 +39,8 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/buffer"
+	"google.golang.org/grpc/internal/grpcsync"
 	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
 	"google.golang.org/grpc/internal/testutils"
 	rlstest "google.golang.org/grpc/internal/testutils/rls"
@@ -975,9 +977,46 @@ func (s) TestDataCachePurging(t *testing.T) {
 	verifyRLSRequest(t, rlsReqCh, true)
 }
 
+// wrappingConnectivityStateSubscriber wraps a grpcsync.Subscriber and forwards
+// connectivity state updates to both the delegate and a channel for testing.
+type wrappingConnectivityStateSubscriber struct {
+	delegate    grpcsync.Subscriber
+	connStateCh *buffer.Unbounded
+}
+
+func (w *wrappingConnectivityStateSubscriber) OnMessage(msg any) {
+	w.delegate.OnMessage(msg)
+	w.connStateCh.Put(msg.(connectivity.State))
+}
+
+// waitForConnectivityState waits for one of the specified connectivity states
+// to appear on the channel, returning the state that matched first. It skips
+// intermediate states that do not match any of the wanted states, and fails
+// the test if the context expires before any of the desired states are reached.
+func waitForConnectivityState(ctx context.Context, t *testing.T, ch *buffer.Unbounded, wants ...connectivity.State) connectivity.State {
+	t.Helper()
+	for {
+		select {
+		case gotState := <-ch.Get():
+			ch.Load()
+			got := gotState.(connectivity.State)
+			for _, want := range wants {
+				if got == want {
+					return got
+				}
+			}
+		case <-ctx.Done():
+			t.Fatalf("Timeout waiting for RLS control channel to become one of %v", wants)
+		}
+	}
+}
+
 // TestControlChannelConnectivityStateMonitoring tests the scenario where the
 // control channel goes down and comes back up again and verifies that backoff
-// state is reset for cache entries in this scenario.
+// state is reset for cache entries in this scenario. It also verifies that:
+//   - Backoff is NOT reset when the control channel first becomes READY (i.e.,
+//     the initial CONNECTING → READY transition should not trigger a backoff reset)
+//   - Backoff is reset for READY → TRANSIENT_FAILURE → READY transitions
 func (s) TestControlChannelConnectivityStateMonitoring(t *testing.T) {
 	// Create a restartable listener which can close existing connections.
 	l, err := testutils.LocalTCPListener()
@@ -1004,6 +1043,16 @@ func (s) TestControlChannelConnectivityStateMonitoring(t *testing.T) {
 	defaultBackoffStrategy = &fakeBackoffStrategy{backoff: defaultTestTimeout}
 	defer func() { defaultBackoffStrategy = origBackoffStrategy }()
 
+	// Override the connectivity state subscriber to wrap the original and
+	// make connectivity state changes visible to the test.
+	wrappedSubscriber := &wrappingConnectivityStateSubscriber{connStateCh: buffer.NewUnbounded()}
+	origConnectivityStateSubscriber := newConnectivityStateSubscriber
+	newConnectivityStateSubscriber = func(delegate grpcsync.Subscriber) grpcsync.Subscriber {
+		wrappedSubscriber.delegate = delegate
+		return wrappedSubscriber
+	}
+	defer func() { newConnectivityStateSubscriber = origConnectivityStateSubscriber }()
+
 	// Register an LB policy to act as the child policy for RLS LB policy.
 	childPolicyName := "test-child-policy" + t.Name()
 	e2e.RegisterRLSChildPolicy(childPolicyName, nil)
@@ -1038,16 +1087,39 @@ func (s) TestControlChannelConnectivityStateMonitoring(t *testing.T) {
 	// Make sure an RLS request is sent out.
 	verifyRLSRequest(t, rlsReqCh, true)
 
+	// Wait for the control channel to move to CONNECTING and then to READY.
+	waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.Connecting)
+	waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.Ready)
+
+	// Verify that the initial READY state of the control channel did NOT trigger
+	// a backoff reset. The resetBackoffHook should only be called when
+	// transitioning from TRANSIENT_FAILURE to READY, not for the initial
+	// CONNECTING → READY transition.
+	select {
+	case <-resetBackoffDone:
+		t.Fatal("Backoff reset was triggered for initial READY state, want no reset")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
 	// Stop the RLS server.
 	lis.Stop()
 
-	// Make another RPC similar to the first one. Since the above cache entry
-	// would have expired by now, this should trigger another RLS request. And
-	// since the RLS server is down, RLS request will fail and the cache entry
-	// will enter backoff, and we have overridden the default backoff strategy to
-	// return a value which will keep this entry in backoff for the whole duration
-	// of the test.
-	makeTestRPCAndVerifyError(ctx, t, cc, codes.Unavailable, nil)
+	// When the server is stopped, the control channel connection is closed,
+	// causing it to transition to IDLE or TRANSIENT_FAILURE.
+	// If it transitions to IDLE, we need to send a new RLS request to force
+	// it to attempt reconnection, which will fail and move it to TRANSIENT_FAILURE.
+	state := waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.Idle, connectivity.TransientFailure)
+
+	if state == connectivity.Idle {
+		// Make another RPC with different headers to bypass any cached backoff entries
+		// and force the control channel to attempt a connection, which will fail
+		// and move it to TRANSIENT_FAILURE.
+		ctxFailed := metadata.AppendToOutgoingContext(ctx, "n1", "v1")
+		makeTestRPCAndVerifyError(ctxFailed, t, cc, codes.Unavailable, nil)
+	}
+
+	// Wait for the control channel to move to TRANSIENT_FAILURE.
+	waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.TransientFailure)
 
 	// Restart the RLS server.
 	lis.Restart()
@@ -1063,9 +1135,13 @@ func (s) TestControlChannelConnectivityStateMonitoring(t *testing.T) {
 	// till the control channel comes moves back to READY. So, override the
 	// backoff strategy to perform a small backoff on this entry.
 	defaultBackoffStrategy = &fakeBackoffStrategy{backoff: defaultTestShortTimeout}
-	ctxOutgoing := metadata.AppendToOutgoingContext(ctx, "n1", "v1")
+	ctxOutgoing := metadata.AppendToOutgoingContext(ctx, "n2", "v2")
 	makeTestRPCAndExpectItToReachBackend(ctxOutgoing, t, cc, backendCh)
 
+	// Wait for the control channel to move back to READY.
+	waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.Ready)
+
+	// Verify that backoff was reset when transitioning from TRANSIENT_FAILURE to READY.
 	select {
 	case <-ctx.Done():
 		t.Fatalf("Timed out waiting for resetBackoffDone")
@@ -1081,6 +1157,118 @@ func (s) TestControlChannelConnectivityStateMonitoring(t *testing.T) {
 	verifyRLSRequest(t, rlsReqCh, true)
 }
 
+// TestControlChannelIdleTransitionNoBackoffReset tests that READY → IDLE → READY
+// transitions do not trigger backoff resets. This is a benign state change that
+// should not affect cache entry backoff state.
+func (s) TestControlChannelIdleTransitionNoBackoffReset(t *testing.T) {
+	// Create a restartable listener which can close existing connections.
+	l, err := testutils.LocalTCPListener()
+	if err != nil {
+		t.Fatalf("net.Listen() failed: %v", err)
+	}
+	lis := testutils.NewRestartableListener(l)
+
+	// Start an RLS server with the restartable listener and set the throttler to
+	// never throttle requests.
+	rlsServer, rlsReqCh := rlstest.SetupFakeRLSServer(t, lis)
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Override the reset backoff hook to get notified.
+	resetBackoffCalled := make(chan struct{}, 1)
+	origResetBackoffHook := resetBackoffHook
+	resetBackoffHook = func() { resetBackoffCalled <- struct{}{} }
+	defer func() { resetBackoffHook = origResetBackoffHook }()
+
+	// Override the backoff strategy to return a large backoff which
+	// will make sure the data cache entry remains in backoff for the
+	// duration of the test.
+	origBackoffStrategy := defaultBackoffStrategy
+	defaultBackoffStrategy = &fakeBackoffStrategy{backoff: defaultTestTimeout}
+	defer func() { defaultBackoffStrategy = origBackoffStrategy }()
+
+	// Override the connectivity state subscriber to wrap the original and
+	// make connectivity state changes visible to the test.
+	wrappedSubscriber := &wrappingConnectivityStateSubscriber{connStateCh: buffer.NewUnbounded()}
+	origConnectivityStateSubscriber := newConnectivityStateSubscriber
+	newConnectivityStateSubscriber = func(delegate grpcsync.Subscriber) grpcsync.Subscriber {
+		wrappedSubscriber.delegate = delegate
+		return wrappedSubscriber
+	}
+	defer func() { newConnectivityStateSubscriber = origConnectivityStateSubscriber }()
+
+	// Register an LB policy to act as the child policy for RLS LB policy.
+	childPolicyName := "test-child-policy" + t.Name()
+	e2e.RegisterRLSChildPolicy(childPolicyName, nil)
+	t.Logf("Registered child policy with name %q", childPolicyName)
+
+	// Build RLS service config with header matchers, and a very low value for
+	// maxAge to ensure that cache entries become invalid very soon.
+	rlsConfig := buildBasicRLSConfig(childPolicyName, rlsServer.Address)
+	rlsConfig.RouteLookupConfig.MaxAge = durationpb.New(defaultTestShortTimeout)
+
+	// Start a test backend, and set up the fake RLS server to return this as a
+	// target in the RLS response.
+	backendCh, backendAddress := startBackend(t)
+	rlsServer.SetResponseCallback(func(_ context.Context, _ *rlspb.RouteLookupRequest) *rlstest.RouteLookupResponse {
+		return &rlstest.RouteLookupResponse{Resp: &rlspb.RouteLookupResponse{Targets: []string{backendAddress}}}
+	})
+
+	// Register a manual resolver and push the RLS service config through it.
+	r := startManualResolverWithConfig(t, rlsConfig)
+
+	cc, err := grpc.NewClient(r.Scheme()+":///", grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("Failed to create gRPC client: %v", err)
+	}
+	defer cc.Close()
+
+	// Make an RPC and ensure it gets routed to the test backend.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	makeTestRPCAndExpectItToReachBackend(ctx, t, cc, backendCh)
+
+	// Make sure an RLS request is sent out.
+	verifyRLSRequest(t, rlsReqCh, true)
+
+	// Wait for the control channel to move to READY.
+	waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.Ready)
+
+	// Verify that the initial READY state did NOT trigger a backoff reset.
+	select {
+	case <-resetBackoffCalled:
+		t.Fatal("Backoff reset was triggered for initial READY state, want no reset")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	// Stop the RLS server to force the control channel to go IDLE. Stop()
+	// closes all existing connections on the listener. Since there are no active
+	// RPCs on the control channel, the subchannel transitions to IDLE instead of
+	// TRANSIENT_FAILURE.
+	lis.Stop()
+
+	// Wait for the control channel to move to IDLE.
+	waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.Idle)
+
+	// Restart the RLS server.
+	lis.Restart()
+
+	// Make another RPC to trigger reconnection. Use different headers to create
+	// a new cache entry.
+	ctxOutgoing := metadata.AppendToOutgoingContext(ctx, "n1", "v1")
+	defaultBackoffStrategy = &fakeBackoffStrategy{backoff: defaultTestShortTimeout}
+	makeTestRPCAndExpectItToReachBackend(ctxOutgoing, t, cc, backendCh)
+
+	// Wait for the control channel to move back to READY.
+	waitForConnectivityState(ctx, t, wrappedSubscriber.connStateCh, connectivity.Ready)
+
+	// Verify that the READY → IDLE → READY transition did NOT trigger a backoff reset.
+	select {
+	case <-resetBackoffCalled:
+		t.Fatal("Backoff reset was triggered for READY → IDLE → READY transition, want no reset")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
 // testCCWrapper wraps a balancer.ClientConn and overrides UpdateState and
 // stores all state updates pushed by the RLS LB policy.
 type testCCWrapper struct {