idle: decrement active call count for streaming RPCs only when the call completes

internal/idle/idle_e2e_test.go

@@ -22,6 +22,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"strings"
 	"testing"
 	"time"
@@ -193,14 +194,25 @@ func (s) TestChannelIdleness_Enabled_OngoingCall(t *testing.T) {
 	}
 	t.Cleanup(func() { cc.Close() })
 
-	// Start a test backend which keeps a unary RPC call active by blocking on a
-	// channel that is closed by the test later on. Also push an address update
-	// via the resolver.
-	blockCh := make(chan struct{})
-	backend := &stubserver.StubServer{
-		EmptyCallF: func(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
-			<-blockCh
-			return &testpb.Empty{}, nil
+	// Start a backend that implements a streaming RPC.
+	backend := stubserver.StubServer{
+		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
+			// Streaming implementation replies with a dummy response until the
+			// client closes the stream (in which case it will see an io.EOF),
+			// or an error occurs while reading/writing messages.
+			for {
+				_, err := stream.Recv()
+				if err == io.EOF {
+					return nil
+				}
+				if err != nil {
+					return err
+				}
+				payload := &testpb.Payload{Body: make([]byte, 32)}
+				if err := stream.Send(&testpb.StreamingOutputCallResponse{Payload: payload}); err != nil {
+					return err
+				}
+			}
 		},
 	}
 	if err := backend.StartServer(); err != nil {
@@ -215,15 +227,16 @@ func (s) TestChannelIdleness_Enabled_OngoingCall(t *testing.T) {
 	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
 
 	// Spawn a goroutine which checks expected state transitions and idleness
-	// channelz trace events. It eventually closes `blockCh`, thereby unblocking
-	// the server RPC handler and the unary call below.
+	// channelz trace events.
 	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()
-		testutils.AwaitNoStateChange(sCtx, t, cc, connectivity.Ready)
+		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 there are no idleness related channelz events.
 		if err := channelzTraceEventNotFound(ctx, "entering idle mode"); err != nil {
@@ -234,16 +247,14 @@ func (s) TestChannelIdleness_Enabled_OngoingCall(t *testing.T) {
 			errCh <- err
 			return
 		}
-
-		// Unblock the unary RPC on the server.
 		errCh <- nil
 	}()
 
-	// Make a unary RPC that blocks on the server, thereby ensuring that the
-	// count of active RPCs on the client is non-zero.
+	// Start a streaming RPC. Even though the stream is inactive (i,e no reads
+	// or writes), the stream should prevent the channel from moving to IDLE.
 	client := testgrpc.NewTestServiceClient(cc)
-	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
-		t.Errorf("EmptyCall RPC failed: %v", err)
+	if _, err := client.FullDuplexCall(ctx); err != nil {
+		t.Fatalf("FullDuplexCall RPC failed: %v", err)
 	}
 
 	select {

stream.go

@@ -161,12 +161,17 @@ func (cc *ClientConn) NewStream(ctx context.Context, desc *StreamDesc, method st
 	if err := cc.idlenessMgr.OnCallBegin(); err != nil {
 		return nil, err
 	}
-	defer cc.idlenessMgr.OnCallEnd()
 
 	// allow interceptor to see all applicable call options, which means those
 	// configured as defaults from dial option as well as per-call options
 	opts = combine(cc.dopts.callOptions, opts)
 
+	// Add a calloption, to decrement the count of active RPCs, that gets
+	// executed when the RPC completes.
+	opts = append(opts, OnFinish(func(error) {
+		cc.idlenessMgr.OnCallEnd()
+	}))
+
 	if cc.dopts.streamInt != nil {
 		return cc.dopts.streamInt(ctx, desc, cc, method, newClientStream, opts...)
 	}