snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/idle/idle.go

@@ -0,0 +1,289 @@
+/*
+ *
+ * Copyright 2023 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package idle contains a component for managing idleness (entering and exiting)
+// based on RPC activity.
+package idle
+
+import (
+	"math"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// For overriding in unit tests.
+var timeAfterFunc = func(d time.Duration, f func()) *time.Timer {
+	return time.AfterFunc(d, f)
+}
+
+// ClientConn is the functionality provided by grpc.ClientConn to enter and exit
+// from idle mode.
+type ClientConn interface {
+	ExitIdleMode()
+	EnterIdleMode()
+}
+
+// Manager implements idleness detection and calls the ClientConn to enter/exit
+// idle mode when appropriate. Must be created by NewManager.
+type Manager struct {
+	// State accessed atomically.
+	lastCallEndTime           int64 // Unix timestamp in nanos; time when the most recent RPC completed.
+	activeCallsCount          int32 // Count of active RPCs; -math.MaxInt32 means channel is idle or is trying to get there.
+	activeSinceLastTimerCheck int32 // Boolean; True if there was an RPC since the last timer callback.
+	closed                    int32 // Boolean; True when the manager is closed.
+
+	// Can be accessed without atomics or mutex since these are set at creation
+	// time and read-only after that.
+	cc      ClientConn // Functionality provided by grpc.ClientConn.
+	timeout time.Duration
+
+	// idleMu is used to guarantee mutual exclusion in two scenarios:
+	// - Opposing intentions:
+	//   - a: Idle timeout has fired and handleIdleTimeout() is trying to put
+	//     the channel in idle mode because the channel has been inactive.
+	//   - b: At the same time an RPC is made on the channel, and OnCallBegin()
+	//     is trying to prevent the channel from going idle.
+	// - Competing intentions:
+	//   - The channel is in idle mode and there are multiple RPCs starting at
+	//     the same time, all trying to move the channel out of idle. Only one
+	//     of them should succeed in doing so, while the other RPCs should
+	//     piggyback on the first one and be successfully handled.
+	idleMu       sync.RWMutex
+	actuallyIdle bool
+	timer        *time.Timer
+}
+
+// NewManager creates a new idleness manager implementation for the
+// given idle timeout.  It begins in idle mode.
+func NewManager(cc ClientConn, timeout time.Duration) *Manager {
+	return &Manager{
+		cc:               cc,
+		timeout:          timeout,
+		actuallyIdle:     true,
+		activeCallsCount: -math.MaxInt32,
+	}
+}
+
+// resetIdleTimerLocked resets the idle timer to the given duration.  Called
+// when exiting idle mode or when the timer fires and we need to reset it.
+func (m *Manager) resetIdleTimerLocked(d time.Duration) {
+	if m.isClosed() || m.timeout == 0 || m.actuallyIdle {
+		return
+	}
+
+	// It is safe to ignore the return value from Reset() because this method is
+	// only ever called from the timer callback or when exiting idle mode.
+	if m.timer != nil {
+		m.timer.Stop()
+	}
+	m.timer = timeAfterFunc(d, m.handleIdleTimeout)
+}
+
+func (m *Manager) resetIdleTimer(d time.Duration) {
+	m.idleMu.Lock()
+	defer m.idleMu.Unlock()
+	m.resetIdleTimerLocked(d)
+}
+
+// handleIdleTimeout is the timer callback that is invoked upon expiry of the
+// configured idle timeout. The channel is considered inactive if there are no
+// ongoing calls and no RPC activity since the last time the timer fired.
+func (m *Manager) handleIdleTimeout() {
+	if m.isClosed() {
+		return
+	}
+
+	if atomic.LoadInt32(&m.activeCallsCount) > 0 {
+		m.resetIdleTimer(m.timeout)
+		return
+	}
+
+	// There has been activity on the channel since we last got here. Reset the
+	// timer and return.
+	if atomic.LoadInt32(&m.activeSinceLastTimerCheck) == 1 {
+		// Set the timer to fire after a duration of idle timeout, calculated
+		// from the time the most recent RPC completed.
+		atomic.StoreInt32(&m.activeSinceLastTimerCheck, 0)
+		m.resetIdleTimer(time.Duration(atomic.LoadInt64(&m.lastCallEndTime)-time.Now().UnixNano()) + m.timeout)
+		return
+	}
+
+	// Now that we've checked that there has been no activity, attempt to enter
+	// idle mode, which is very likely to succeed.
+	if m.tryEnterIdleMode(true) {
+		// Successfully entered idle mode. No timer needed until we exit idle.
+		return
+	}
+
+	// Failed to enter idle mode due to a concurrent RPC that kept the channel
+	// active, or because of an error from the channel. Undo the attempt to
+	// enter idle, and reset the timer to try again later.
+	m.resetIdleTimer(m.timeout)
+}
+
+// tryEnterIdleMode instructs the channel to enter idle mode. But before
+// that, it performs a last minute check to ensure that no new RPC has come in,
+// making the channel active.
+//
+// checkActivity controls if a check for RPC activity, since the last time the
+// idle_timeout fired, is made.
+
+// Return value indicates whether or not the channel moved to idle mode.
+//
+// Holds idleMu which ensures mutual exclusion with exitIdleMode.
+func (m *Manager) tryEnterIdleMode(checkActivity bool) bool {
+	// Setting the activeCallsCount to -math.MaxInt32 indicates to OnCallBegin()
+	// that the channel is either in idle mode or is trying to get there.
+	if !atomic.CompareAndSwapInt32(&m.activeCallsCount, 0, -math.MaxInt32) {
+		// This CAS operation can fail if an RPC started after we checked for
+		// activity in the timer handler, or one was ongoing from before the
+		// last time the timer fired, or if a test is attempting to enter idle
+		// mode without checking.  In all cases, abort going into idle mode.
+		return false
+	}
+	// N.B. if we fail to enter idle mode after this, we must re-add
+	// math.MaxInt32 to m.activeCallsCount.
+
+	m.idleMu.Lock()
+	defer m.idleMu.Unlock()
+
+	if atomic.LoadInt32(&m.activeCallsCount) != -math.MaxInt32 {
+		// We raced and lost to a new RPC. Very rare, but stop entering idle.
+		atomic.AddInt32(&m.activeCallsCount, math.MaxInt32)
+		return false
+	}
+	if checkActivity && atomic.LoadInt32(&m.activeSinceLastTimerCheck) == 1 {
+		// A very short RPC could have come in (and also finished) after we
+		// checked for calls count and activity in handleIdleTimeout(), but
+		// before the CAS operation. So, we need to check for activity again.
+		atomic.AddInt32(&m.activeCallsCount, math.MaxInt32)
+		return false
+	}
+
+	// No new RPCs have come in since we set the active calls count value to
+	// -math.MaxInt32. And since we have the lock, it is safe to enter idle mode
+	// unconditionally now.
+	m.cc.EnterIdleMode()
+	m.actuallyIdle = true
+	return true
+}
+
+// EnterIdleModeForTesting instructs the channel to enter idle mode.
+func (m *Manager) EnterIdleModeForTesting() {
+	m.tryEnterIdleMode(false)
+}
+
+// OnCallBegin is invoked at the start of every RPC.
+func (m *Manager) OnCallBegin() {
+	if m.isClosed() {
+		return
+	}
+
+	if atomic.AddInt32(&m.activeCallsCount, 1) > 0 {
+		// Channel is not idle now. Set the activity bit and allow the call.
+		atomic.StoreInt32(&m.activeSinceLastTimerCheck, 1)
+		return
+	}
+
+	// Channel is either in idle mode or is in the process of moving to idle
+	// mode. Attempt to exit idle mode to allow this RPC.
+	m.ExitIdleMode()
+	atomic.StoreInt32(&m.activeSinceLastTimerCheck, 1)
+}
+
+// ExitIdleMode instructs m to call the ClientConn's ExitIdleMode and update its
+// internal state.
+func (m *Manager) ExitIdleMode() {
+	// Holds idleMu which ensures mutual exclusion with tryEnterIdleMode.
+	m.idleMu.Lock()
+	defer m.idleMu.Unlock()
+
+	if m.isClosed() || !m.actuallyIdle {
+		// This can happen in three scenarios:
+		// - handleIdleTimeout() set the calls count to -math.MaxInt32 and called
+		//   tryEnterIdleMode(). But before the latter could grab the lock, an RPC
+		//   came in and OnCallBegin() noticed that the calls count is negative.
+		// - Channel is in idle mode, and multiple new RPCs come in at the same
+		//   time, all of them notice a negative calls count in OnCallBegin and get
+		//   here. The first one to get the lock would get the channel to exit idle.
+		// - Channel is not in idle mode, and the user calls Connect which calls
+		//   m.ExitIdleMode.
+		//
+		// In any case, there is nothing to do here.
+		return
+	}
+
+	m.cc.ExitIdleMode()
+
+	// Undo the idle entry process. This also respects any new RPC attempts.
+	atomic.AddInt32(&m.activeCallsCount, math.MaxInt32)
+	m.actuallyIdle = false
+
+	// Start a new timer to fire after the configured idle timeout.
+	m.resetIdleTimerLocked(m.timeout)
+}
+
+// UnsafeSetNotIdle instructs the Manager to update its internal state to
+// reflect the reality that the channel is no longer in IDLE mode.
+//
+// N.B. This method is intended only for internal use by the gRPC client
+// when it exits IDLE mode **manually** from `Dial`. The callsite must ensure:
+//   - The channel was **actually in IDLE mode** immediately prior to the call.
+//   - There is **no concurrent activity** that could cause the channel to exit
+//     IDLE mode *naturally* at the same time.
+func (m *Manager) UnsafeSetNotIdle() {
+	m.idleMu.Lock()
+	defer m.idleMu.Unlock()
+
+	atomic.AddInt32(&m.activeCallsCount, math.MaxInt32)
+	m.actuallyIdle = false
+	m.resetIdleTimerLocked(m.timeout)
+}
+
+// OnCallEnd is invoked at the end of every RPC.
+func (m *Manager) OnCallEnd() {
+	if m.isClosed() {
+		return
+	}
+
+	// Record the time at which the most recent call finished.
+	atomic.StoreInt64(&m.lastCallEndTime, time.Now().UnixNano())
+
+	// Decrement the active calls count. This count can temporarily go negative
+	// when the timer callback is in the process of moving the channel to idle
+	// mode, but one or more RPCs come in and complete before the timer callback
+	// can get done with the process of moving to idle mode.
+	atomic.AddInt32(&m.activeCallsCount, -1)
+}
+
+func (m *Manager) isClosed() bool {
+	return atomic.LoadInt32(&m.closed) == 1
+}
+
+// Close stops the timer associated with the Manager, if it exists.
+func (m *Manager) Close() {
+	atomic.StoreInt32(&m.closed, 1)
+
+	m.idleMu.Lock()
+	if m.timer != nil {
+		m.timer.Stop()
+		m.timer = nil
+	}
+	m.idleMu.Unlock()
+}

internal/idle/idle_e2e_test.go

@@ -0,0 +1,589 @@
+/*
+ *
+ * Copyright 2023 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package idle_test
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/roundrobin"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/stubserver"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/grpc/status"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+)
+
+func init() {
+	channelz.TurnOn()
+}
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+const (
+	defaultTestTimeout          = 10 * time.Second
+	defaultTestShortTimeout     = 100 * time.Millisecond
+	defaultTestShortIdleTimeout = 500 * time.Millisecond
+)
+
+// channelzTraceEventFound looks up the top-channels in channelz (expects a
+// single one), and checks if there is a trace event on the channel matching the
+// provided description string.
+func channelzTraceEventFound(ctx context.Context, wantDesc string) error {
+	for ctx.Err() == nil {
+		tcs, _ := channelz.GetTopChannels(0, 0)
+		if l := len(tcs); l != 1 {
+			return fmt.Errorf("when looking for channelz trace event with description %q, found %d top-level channels, want 1", wantDesc, l)
+		}
+		trace := tcs[0].Trace()
+		if trace == nil {
+			return fmt.Errorf("when looking for channelz trace event with description %q, no trace events found for top-level channel", wantDesc)
+		}
+
+		for _, e := range trace.Events {
+			if strings.Contains(e.Desc, wantDesc) {
+				return nil
+			}
+		}
+	}
+	return fmt.Errorf("when looking for channelz trace event with description %q, %w", wantDesc, ctx.Err())
+}
+
+// Registers a wrapped round_robin LB policy for the duration of this test that
+// retains all the functionality of the round_robin LB policy and makes the
+// balancer close event available for inspection by the test.
+//
+// Returns a channel that gets pinged when the balancer is closed.
+func registerWrappedRoundRobinPolicy(t *testing.T) chan struct{} {
+	rrBuilder := balancer.Get(roundrobin.Name)
+	closeCh := make(chan struct{}, 1)
+	stub.Register(roundrobin.Name, stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = rrBuilder.Build(bd.ClientConn, bd.BuildOptions)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			select {
+			case closeCh <- struct{}{}:
+			default:
+			}
+			bd.ChildBalancer.Close()
+		},
+	})
+	t.Cleanup(func() { balancer.Register(rrBuilder) })
+
+	return closeCh
+}
+
+// Tests the case where channel idleness is disabled by passing an idle_timeout
+// of 0. Verifies that a READY channel with no RPCs does not move to IDLE.
+func (s) TestChannelIdleness_Disabled_NoActivity(t *testing.T) {
+	closeCh := registerWrappedRoundRobinPolicy(t)
+
+	// Create a ClientConn with idle_timeout set to 0.
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(0), // Disable idleness.
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	cc.Connect()
+
+	// Start a test backend and push an address update via the resolver.
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// Verify that the ClientConn stays in READY.
+	sCtx, sCancel := context.WithTimeout(ctx, 3*defaultTestShortIdleTimeout)
+	defer sCancel()
+	testutils.AwaitNoStateChange(sCtx, t, cc, connectivity.Ready)
+
+	// Verify that the LB policy is not closed which is expected to happen when
+	// the channel enters IDLE.
+	sCtx, sCancel = context.WithTimeout(ctx, defaultTestShortIdleTimeout)
+	defer sCancel()
+	select {
+	case <-sCtx.Done():
+	case <-closeCh:
+		t.Fatal("LB policy closed when expected not to")
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Verifies that a READY channel with no RPCs moves to IDLE, and
+// the connection to the backend is closed.
+func (s) TestChannelIdleness_Enabled_NoActivity(t *testing.T) {
+	closeCh := registerWrappedRoundRobinPolicy(t)
+
+	// Create a ClientConn with a short idle_timeout.
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	cc.Connect()
+	// Start a test backend and push an address update via the resolver.
+	lis := testutils.NewListenerWrapper(t, nil)
+	backend := stubserver.StartTestService(t, &stubserver.StubServer{Listener: lis})
+	defer backend.Stop()
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// Retrieve the wrapped conn from the listener.
+	v, err := lis.NewConnCh.Receive(ctx)
+	if err != nil {
+		t.Fatalf("Failed to retrieve conn from test listener: %v", err)
+	}
+	conn := v.(*testutils.ConnWrapper)
+
+	// Verify that the ClientConn moves to IDLE as there is no activity.
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+
+	// Verify idleness related channelz events.
+	if err := channelzTraceEventFound(ctx, "entering idle mode"); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify that the previously open connection is closed.
+	if _, err := conn.CloseCh.Receive(ctx); err != nil {
+		t.Fatalf("Failed when waiting for connection to be closed after channel entered IDLE: %v", err)
+	}
+
+	// Verify that the LB policy is closed.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for LB policy to be closed after the channel enters IDLE")
+	case <-closeCh:
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Verifies that a READY channel with an ongoing RPC stays READY.
+func (s) TestChannelIdleness_Enabled_OngoingCall(t *testing.T) {
+	tests := []struct {
+		name    string
+		makeRPC func(ctx context.Context, client testgrpc.TestServiceClient) error
+	}{
+		{
+			name: "unary",
+			makeRPC: func(ctx context.Context, client testgrpc.TestServiceClient) error {
+				if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+					return fmt.Errorf("EmptyCall RPC failed: %v", err)
+				}
+				return nil
+			},
+		},
+		{
+			name: "streaming",
+			makeRPC: func(ctx context.Context, client testgrpc.TestServiceClient) error {
+				stream, err := client.FullDuplexCall(ctx)
+				if err != nil {
+					t.Fatalf("FullDuplexCall RPC failed: %v", err)
+				}
+				if _, err := stream.Recv(); err != nil && err != io.EOF {
+					t.Fatalf("stream.Recv() failed: %v", err)
+				}
+				return nil
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			closeCh := registerWrappedRoundRobinPolicy(t)
+
+			// Create a ClientConn with a short idle_timeout.
+			r := manual.NewBuilderWithScheme("whatever")
+			dopts := []grpc.DialOption{
+				grpc.WithTransportCredentials(insecure.NewCredentials()),
+				grpc.WithResolvers(r),
+				grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+				grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+			}
+			cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+			if err != nil {
+				t.Fatalf("grpc.NewClient() failed: %v", err)
+			}
+			defer cc.Close()
+			cc.Connect()
+			// Start a test backend that keeps the RPC call active by blocking
+			// on a channel that is closed by the test later on.
+			blockCh := make(chan struct{})
+			backend := &stubserver.StubServer{
+				EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+					<-blockCh
+					return &testpb.Empty{}, nil
+				},
+				FullDuplexCallF: func(testgrpc.TestService_FullDuplexCallServer) error {
+					<-blockCh
+					return nil
+				},
+			}
+			if err := backend.StartServer(); err != nil {
+				t.Fatalf("Failed to start backend: %v", err)
+			}
+			defer backend.Stop()
+
+			// Push an address update containing the address of the above
+			// backend via the manual resolver.
+			r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+			// Verify that the ClientConn moves to READY.
+			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+			defer cancel()
+			testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+			// Spawn a goroutine to check for expected behavior while a blocking
+			// RPC all is made from the main test goroutine.
+			errCh := make(chan error, 1)
+			go func() {
+				defer close(blockCh)
+
+				// Verify that the ClientConn stays in READY.
+				sCtx, sCancel := context.WithTimeout(ctx, 3*defaultTestShortIdleTimeout)
+				defer sCancel()
+				if cc.WaitForStateChange(sCtx, connectivity.Ready) {
+					errCh <- fmt.Errorf("state changed from %q to %q when no state change was expected", connectivity.Ready, cc.GetState())
+					return
+				}
+
+				// Verify that the LB policy is not closed which is expected to happen when
+				// the channel enters IDLE.
+				sCtx, sCancel = context.WithTimeout(ctx, defaultTestShortIdleTimeout)
+				defer sCancel()
+				select {
+				case <-sCtx.Done():
+				case <-closeCh:
+					errCh <- fmt.Errorf("LB policy closed when expected not to")
+				}
+				errCh <- nil
+			}()
+
+			if err := test.makeRPC(ctx, testgrpc.NewTestServiceClient(cc)); err != nil {
+				t.Fatalf("%s rpc failed: %v", test.name, err)
+			}
+
+			select {
+			case err := <-errCh:
+				if err != nil {
+					t.Fatal(err)
+				}
+			case <-ctx.Done():
+				t.Fatalf("Timeout when trying to verify that an active RPC keeps channel from moving to IDLE")
+			}
+		})
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Verifies that activity on a READY channel (frequent and short
+// RPCs) keeps it from moving to IDLE.
+func (s) TestChannelIdleness_Enabled_ActiveSinceLastCheck(t *testing.T) {
+	closeCh := registerWrappedRoundRobinPolicy(t)
+
+	// Create a ClientConn with a short idle_timeout.
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	cc.Connect()
+	// Start a test backend and push an address update via the resolver.
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// For a duration of three times the configured idle timeout, making RPCs
+	// every now and then and ensure that the channel does not move out of
+	// READY.
+	sCtx, sCancel := context.WithTimeout(ctx, 3*defaultTestShortIdleTimeout)
+	defer sCancel()
+	go func() {
+		for ; sCtx.Err() == nil; <-time.After(defaultTestShortIdleTimeout / 4) {
+			client := testgrpc.NewTestServiceClient(cc)
+			if _, err := client.EmptyCall(sCtx, &testpb.Empty{}); err != nil {
+				// While iterating through this for loop, at some point in time,
+				// the context deadline will expire. It is safe to ignore that
+				// error code.
+				if status.Code(err) != codes.DeadlineExceeded {
+					t.Errorf("EmptyCall RPC failed: %v", err)
+					return
+				}
+			}
+		}
+	}()
+
+	// Verify that the ClientConn stays in READY.
+	testutils.AwaitNoStateChange(sCtx, t, cc, connectivity.Ready)
+
+	// Verify that the LB policy is not closed which is expected to happen when
+	// the channel enters IDLE.
+	select {
+	case <-sCtx.Done():
+	case <-closeCh:
+		t.Fatal("LB policy closed when expected not to")
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Verifies that a READY channel with no RPCs moves to IDLE. Also
+// verifies that a subsequent RPC on the IDLE channel kicks it out of IDLE.
+func (s) TestChannelIdleness_Enabled_ExitIdleOnRPC(t *testing.T) {
+	closeCh := registerWrappedRoundRobinPolicy(t)
+
+	// Start a test backend and set the bootstrap state of the resolver to
+	// include this address. This will ensure that when the resolver is
+	// restarted when exiting idle, it will push the same address to grpc again.
+	r := manual.NewBuilderWithScheme("whatever")
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Create a ClientConn with a short idle_timeout.
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	cc.Connect()
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// Verify that the ClientConn moves to IDLE as there is no activity.
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+
+	// Verify idleness related channelz events.
+	if err := channelzTraceEventFound(ctx, "entering idle mode"); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify that the LB policy is closed.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for LB policy to be closed after the channel enters IDLE")
+	case <-closeCh:
+	}
+
+	// Make an RPC and ensure that it succeeds and moves the channel back to
+	// READY.
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("EmptyCall RPC failed: %v", err)
+	}
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+	if err := channelzTraceEventFound(ctx, "exiting idle mode"); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Simulates a race between the idle timer firing and RPCs being
+// initiated, after a period of inactivity on the channel.
+//
+// After a period of inactivity (for the configured idle timeout duration), when
+// RPCs are started, there are two possibilities:
+//   - the idle timer wins the race and puts the channel in idle. The RPCs then
+//     kick it out of idle.
+//   - the RPCs win the race, and therefore the channel never moves to idle.
+//
+// In either of these cases, all RPCs must succeed.
+func (s) TestChannelIdleness_Enabled_IdleTimeoutRacesWithRPCs(t *testing.T) {
+	// Start a test backend and set the bootstrap state of the resolver to
+	// include this address. This will ensure that when the resolver is
+	// restarted when exiting idle, it will push the same address to grpc again.
+	r := manual.NewBuilderWithScheme("whatever")
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Create a ClientConn with a short idle_timeout.
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Errorf("EmptyCall RPC failed: %v", err)
+	}
+
+	// Make an RPC every defaultTestShortTimeout duration so as to race with the
+	// idle timeout. Whether the idle timeout wins the race or the RPC wins the
+	// race, RPCs must succeed.
+	for i := 0; i < 20; i++ {
+		<-time.After(defaultTestShortTimeout)
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+			t.Fatalf("EmptyCall RPC failed: %v", err)
+		}
+		t.Logf("Iteration %d succeeded", i)
+	}
+}
+
+// Tests the case where the channel is IDLE and we call cc.Connect.
+func (s) TestChannelIdleness_Connect(t *testing.T) {
+	// Start a test backend and set the bootstrap state of the resolver to
+	// include this address. This will ensure that when the resolver is
+	// restarted when exiting idle, it will push the same address to grpc again.
+	r := manual.NewBuilderWithScheme("whatever")
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Create a ClientConn with a short idle_timeout.
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	// Verify that the ClientConn moves to IDLE.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+
+	// Connect should exit channel idleness.
+	cc.Connect()
+
+	// Verify that the ClientConn moves back to READY.
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+}
+
+// TestChannelIdleness_RPCFailingBeforeAttempt_FinishesOnce verifies that an
+// RPC which fails during stream creation, before any attempt is started, is
+// reported as finished exactly once.
+func (s) TestChannelIdleness_RPCFailingBeforeAttempt_FinishesOnce(t *testing.T) {
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+
+	cc, err := grpc.NewClient(fmt.Sprintf("dns:///%s", backend.Address), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	client := testgrpc.NewTestServiceClient(cc)
+
+	// Warm-up RPC: moves the channel to READY and fires the resolver's
+	// first-resolution event, so that the next RPC does not block waiting for
+	// addresses and instead proceeds to fail during stream creation.
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("EmptyCall() failed: %v", err)
+	}
+
+	// Issue an RPC with an already-canceled context so that it fails during
+	// stream creation, before any attempt is started, and count how many times
+	// gRPC reports it as finished. The contract is exactly once.
+	var finishCount atomic.Int32
+	cancel()
+	_, err = client.EmptyCall(ctx, &testpb.Empty{}, grpc.OnFinish(func(error) { finishCount.Add(1) }))
+	if got, want := status.Code(err), codes.Canceled; got != want {
+		t.Fatalf("EmptyCall() with canceled context returned code %v, want %v (err: %v)", got, want, err)
+	}
+	if got := finishCount.Load(); got != 1 {
+		t.Fatalf("OnFinish callback invoked %d times, want exactly 1", got)
+	}
+}

internal/idle/idle_test.go

@@ -0,0 +1,378 @@
+/*
+ *
+ * Copyright 2023 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package idle
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+const (
+	defaultTestTimeout      = 10 * time.Second
+	defaultTestIdleTimeout  = 500 * time.Millisecond // A short idle_timeout for tests.
+	defaultTestShortTimeout = 10 * time.Millisecond  // A small deadline to wait for events expected to not happen.
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+type testEnforcer struct {
+	exitIdleCh  chan struct{}
+	enterIdleCh chan struct{}
+}
+
+func (ti *testEnforcer) ExitIdleMode() {
+	ti.exitIdleCh <- struct{}{}
+}
+
+func (ti *testEnforcer) EnterIdleMode() {
+	ti.enterIdleCh <- struct{}{}
+}
+
+func newTestEnforcer() *testEnforcer {
+	return &testEnforcer{
+		exitIdleCh:  make(chan struct{}, 1),
+		enterIdleCh: make(chan struct{}, 1),
+	}
+}
+
+// overrideNewTimer overrides the new timer creation function by ensuring that a
+// message is pushed on the returned channel everytime the timer fires.
+func overrideNewTimer(t *testing.T) <-chan struct{} {
+	t.Helper()
+
+	ch := make(chan struct{}, 1)
+	origTimeAfterFunc := timeAfterFunc
+	timeAfterFunc = func(d time.Duration, callback func()) *time.Timer {
+		return time.AfterFunc(d, func() {
+			select {
+			case ch <- struct{}{}:
+			default:
+			}
+			callback()
+		})
+	}
+	t.Cleanup(func() { timeAfterFunc = origTimeAfterFunc })
+	return ch
+}
+
+// TestManager_Disabled tests the case where the idleness manager is
+// disabled by passing an idle_timeout of 0. Verifies the following things:
+//   - timer callback does not fire
+//   - an RPC triggers a call to ExitIdleMode on the ClientConn
+//   - more calls to RPC termination (as compared to RPC initiation) does not
+//     result in an error log
+func (s) TestManager_Disabled(t *testing.T) {
+	callbackCh := overrideNewTimer(t)
+
+	// Create an idleness manager that is disabled because of idleTimeout being
+	// set to `0`.
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(0))
+
+	// Ensure that the timer callback does not fire within a short deadline.
+	select {
+	case <-callbackCh:
+		t.Fatal("Idle timer callback fired when manager is disabled")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	// The first invocation of OnCallBegin() should lead to a call to
+	// ExitIdleMode() on the enforcer.
+	go mgr.OnCallBegin()
+	select {
+	case <-enforcer.exitIdleCh:
+	case <-time.After(defaultTestShortTimeout):
+		t.Fatal("Timeout waiting for channel to move out of idle mode")
+	}
+
+	// If the number of calls to OnCallEnd() exceeds the number of calls to
+	// OnCallBegin(), the idleness manager is expected to throw an error log
+	// (which will cause our TestLogger to fail the test). But since the manager
+	// is disabled, this should not happen.
+	mgr.OnCallEnd()
+	mgr.OnCallEnd()
+
+	// The idleness manager is explicitly not closed here. But since the manager
+	// is disabled, it will not start the run goroutine, and hence we expect the
+	// leak checker to not find any leaked goroutines.
+}
+
+// TestManager_Enabled_TimerFires tests the case where the idle manager
+// is enabled. Ensures that when there are no RPCs, the timer callback is
+// invoked and the EnterIdleMode() method is invoked on the enforcer.
+func (s) TestManager_Enabled_TimerFires(t *testing.T) {
+	callbackCh := overrideNewTimer(t)
+
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(defaultTestIdleTimeout))
+	defer mgr.Close()
+	mgr.ExitIdleMode()
+
+	// Ensure that the timer callback fires within an appropriate amount of time.
+	select {
+	case <-callbackCh:
+	case <-time.After(2 * defaultTestIdleTimeout):
+		t.Fatal("Timeout waiting for idle timer callback to fire")
+	}
+
+	// Ensure that the channel moves to idle mode eventually.
+	select {
+	case <-enforcer.enterIdleCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Timeout waiting for channel to move to idle")
+	}
+}
+
+// TestManager_Enabled_OngoingCall tests the case where the idle manager
+// is enabled. Ensures that when there is an ongoing RPC, the channel does not
+// enter idle mode.
+func (s) TestManager_Enabled_OngoingCall(t *testing.T) {
+	callbackCh := overrideNewTimer(t)
+
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(defaultTestIdleTimeout))
+	defer mgr.Close()
+	mgr.ExitIdleMode()
+
+	// Fire up a goroutine that simulates an ongoing RPC that is terminated
+	// after the timer callback fires for the first time.
+	timerFired := make(chan struct{})
+	go func() {
+		mgr.OnCallBegin()
+		<-timerFired
+		mgr.OnCallEnd()
+	}()
+
+	// Ensure that the timer callback fires and unblock the above goroutine.
+	select {
+	case <-callbackCh:
+		close(timerFired)
+	case <-time.After(2 * defaultTestIdleTimeout):
+		t.Fatal("Timeout waiting for idle timer callback to fire")
+	}
+
+	// The invocation of the timer callback should not put the channel in idle
+	// mode since we had an ongoing RPC.
+	select {
+	case <-enforcer.enterIdleCh:
+		t.Fatalf("EnterIdleMode() called on enforcer when active RPC exists")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	// Since we terminated the ongoing RPC and we have no other active RPCs, the
+	// channel must move to idle eventually.
+	select {
+	case <-enforcer.enterIdleCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Timeout waiting for channel to move to idle")
+	}
+}
+
+// TestManager_Enabled_ActiveSinceLastCheck tests the case where the
+// idle manager is enabled. Ensures that when there are active RPCs in the last
+// period (even though there is no active call when the timer fires), the
+// channel does not enter idle mode.
+func (s) TestManager_Enabled_ActiveSinceLastCheck(t *testing.T) {
+	callbackCh := overrideNewTimer(t)
+
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(defaultTestIdleTimeout))
+	defer mgr.Close()
+	mgr.ExitIdleMode()
+
+	// Fire up a goroutine that simulates unary RPCs until the timer callback
+	// fires.
+	timerFired := make(chan struct{})
+	go func() {
+		for ; ; <-time.After(defaultTestShortTimeout) {
+			mgr.OnCallBegin()
+			mgr.OnCallEnd()
+
+			select {
+			case <-timerFired:
+				return
+			default:
+			}
+		}
+	}()
+
+	// Ensure that the timer callback fires, and that we don't enter idle as
+	// part of this invocation of the timer callback, since we had some RPCs in
+	// this period.
+	select {
+	case <-callbackCh:
+		close(timerFired)
+	case <-time.After(2 * defaultTestIdleTimeout):
+		close(timerFired)
+		t.Fatal("Timeout waiting for idle timer callback to fire")
+	}
+	select {
+	case <-enforcer.enterIdleCh:
+		t.Fatalf("EnterIdleMode() called on enforcer when one RPC completed in the last period")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	// Since the unary RPC terminated and we have no other active RPCs, the
+	// channel must move to idle eventually.
+	select {
+	case <-enforcer.enterIdleCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Timeout waiting for channel to move to idle")
+	}
+}
+
+// TestManager_Enabled_ExitIdleOnRPC tests the case where the idle
+// manager is enabled. Ensures that the channel moves out of idle when an RPC is
+// initiated.
+func (s) TestManager_Enabled_ExitIdleOnRPC(t *testing.T) {
+	overrideNewTimer(t)
+
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(defaultTestIdleTimeout))
+	defer mgr.Close()
+
+	mgr.ExitIdleMode()
+	<-enforcer.exitIdleCh
+	// Ensure that the channel moves to idle since there are no RPCs.
+	select {
+	case <-enforcer.enterIdleCh:
+	case <-time.After(2 * defaultTestIdleTimeout):
+		t.Fatal("Timeout waiting for channel to move to idle mode")
+	}
+
+	for i := 0; i < 100; i++ {
+		// A call to OnCallBegin and OnCallEnd simulates an RPC.
+		go func() {
+			mgr.OnCallBegin()
+			mgr.OnCallEnd()
+		}()
+	}
+
+	// Ensure that the channel moves out of idle as a result of the above RPC.
+	select {
+	case <-enforcer.exitIdleCh:
+	case <-time.After(2 * defaultTestIdleTimeout):
+		t.Fatal("Timeout waiting for channel to move out of idle mode")
+	}
+
+	// Ensure that only one call to exit idle mode is made to the CC.
+	sCtx, sCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer sCancel()
+	select {
+	case <-enforcer.exitIdleCh:
+		t.Fatal("More than one call to exit idle mode on the ClientConn; only one expected")
+	case <-sCtx.Done():
+	}
+}
+
+type racyState int32
+
+const (
+	stateInitial racyState = iota
+	stateEnteredIdle
+	stateExitedIdle
+	stateActiveRPCs
+)
+
+// racyEnforcer is a test idleness enforcer used specifically to test the
+// race between idle timeout and incoming RPCs.
+type racyEnforcer struct {
+	t       *testing.T
+	state   *racyState // Accessed atomically.
+	started bool
+}
+
+// ExitIdleMode sets the internal state to stateExitedIdle. We should only ever
+// exit idle when we are currently in idle.
+func (ri *racyEnforcer) ExitIdleMode() {
+	// Set only on the initial ExitIdleMode
+	if ri.started == false {
+		if !atomic.CompareAndSwapInt32((*int32)(ri.state), int32(stateInitial), int32(stateInitial)) {
+			ri.t.Errorf("idleness enforcer's first ExitIdleMode after EnterIdleMode")
+			return
+		}
+		ri.started = true
+		return
+	}
+	if !atomic.CompareAndSwapInt32((*int32)(ri.state), int32(stateEnteredIdle), int32(stateExitedIdle)) {
+		ri.t.Errorf("idleness enforcer asked to exit idle when it did not enter idle earlier")
+		return
+	}
+}
+
+// EnterIdleMode attempts to set the internal state to stateEnteredIdle. We should only ever enter idle before RPCs start.
+func (ri *racyEnforcer) EnterIdleMode() {
+	if !atomic.CompareAndSwapInt32((*int32)(ri.state), int32(stateInitial), int32(stateEnteredIdle)) {
+		ri.t.Errorf("idleness enforcer asked to enter idle after rpcs started")
+	}
+}
+
+// TestManager_IdleTimeoutRacesWithOnCallBegin tests the case where firing of
+// the idle timeout races with an incoming RPC. The test verifies that if the
+// timer callback wins the race and puts the channel in idle, the RPCs can kick
+// it out of idle. And if the RPCs win the race and keep the channel active,
+// then the timer callback should not attempt to put the channel in idle mode.
+func (s) TestManager_IdleTimeoutRacesWithOnCallBegin(t *testing.T) {
+	// Run multiple iterations to simulate different possibilities.
+	for i := 0; i < 20; i++ {
+		t.Run(fmt.Sprintf("iteration=%d", i), func(t *testing.T) {
+			var idlenessState racyState
+			enforcer := &racyEnforcer{t: t, state: &idlenessState}
+
+			// Configure a large idle timeout so that we can control the
+			// race between the timer callback and RPCs.
+			mgr := NewManager(enforcer, time.Duration(10*time.Minute))
+			defer mgr.Close()
+			mgr.ExitIdleMode()
+
+			var wg sync.WaitGroup
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				<-time.After(defaultTestIdleTimeout / 50)
+				mgr.handleIdleTimeout()
+			}()
+			for j := 0; j < 5; j++ {
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					// Wait for the configured idle timeout and simulate an RPC to
+					// race with the idle timeout timer callback.
+					<-time.After(defaultTestIdleTimeout / 50)
+					mgr.OnCallBegin()
+					atomic.StoreInt32((*int32)(&idlenessState), int32(stateActiveRPCs))
+					mgr.OnCallEnd()
+				}()
+			}
+			wg.Wait()
+		})
+	}
+}