context cancellation callbacks for unary and streaming

Author: Pranjali-2501

stream.go

@@ -514,24 +514,25 @@ func newClientStreamWithParams(ctx context.Context, desc *StreamDesc, cc *Client
 		}
 	}
 
-	if desc != unaryStreamDesc {
-		// Listen on cc and stream contexts to cleanup when the user closes the
-		// ClientConn or cancels the stream context.  In all other cases, an error
-		// should already be injected into the recv buffer by the transport, which
-		// the client will eventually receive, and then we will cancel the stream's
-		// context in clientStream.finish.
-		go func() {
-			select {
-			case <-cc.ctx.Done():
-				cs.finish(ErrClientConnClosing)
-			case <-ctx.Done():
-				cs.finish(toRPCErr(ctx.Err()))
-			}
-		}()
-	}
+	cs.registerContextCallbacks()
 	return cs, nil
 }
 
+// registerContextCallbacks registers callbacks to clean up the stream when the
+// RPC context expires or the ClientConn is closed.
+func (cs *clientStream) registerContextCallbacks() {
+	cs.mu.Lock()
+	defer cs.mu.Unlock()
+	if !cs.finished {
+		cs.stopCtx = context.AfterFunc(cs.ctx, func() {
+			cs.finish(toRPCErr(cs.ctx.Err()))
+		})
+		cs.stopCC = context.AfterFunc(cs.cc.ctx, func() {
+			cs.finish(ErrClientConnClosing)
+		})
+	}
+}
+
 // newAttemptLocked creates a new csAttempt without a transport or stream.
 func (cs *clientStream) newAttemptLocked(isTransparent bool) (*csAttempt, error) {
 	if err := cs.ctx.Err(); err != nil {
@@ -713,6 +714,9 @@ type clientStream struct {
 	replayBuffer     []replayOp // operations to replay on retry
 	replayBufferSize int        // current size of replayBuffer
 
+	stopCtx func() bool // stops the callback registered to run on RPC context expiration
+	stopCC  func() bool // stops the callback registered to run on ClientConn context expiration
+
 	// Bool fields are grouped at the tail to eliminate the alignment padding
 	// that would otherwise follow each bool when the next field is pointer- or
 	// int-sized. See https://github.com/grpc/grpc-go/issues/9280 for benchmarks.
@@ -1195,6 +1199,10 @@ func (cs *clientStream) finish(err error) {
 		return
 	}
 	cs.finished = true
+	stopCtx := cs.stopCtx
+	stopCC := cs.stopCC
+	cs.stopCtx = nil
+	cs.stopCC = nil
 	cs.commitAttemptLocked()
 	attemptCreated := cs.attempt != nil
 	if attemptCreated {
@@ -1208,6 +1216,12 @@ func (cs *clientStream) finish(err error) {
 	}
 
 	cs.mu.Unlock()
+	if stopCtx != nil {
+		stopCtx()
+	}
+	if stopCC != nil {
+		stopCC()
+	}
 	// Only one of cancel or trailer needs to be logged.
 	if len(cs.binlogs) != 0 {
 		switch err {
@@ -1472,29 +1486,28 @@ func newNonRetryClientStream(ctx context.Context, desc *StreamDesc, method strin
 	as.transportStream = s
 	as.parser = parser{r: s, bufferPool: ac.dopts.copts.BufferPool}
 	ac.incrCallsStarted()
-	if desc != unaryStreamDesc {
-		// Listen on stream context to cleanup when the stream context is
-		// canceled.  Also listen for the addrConn's context in case the
-		// addrConn is closed or reconnects to a different address.  In all
-		// other cases, an error should already be injected into the recv
-		// buffer by the transport, which the client will eventually receive,
-		// and then we will cancel the stream's context in
-		// addrConnStream.finish.
-		go func() {
-			ac.mu.Lock()
-			acCtx := ac.ctx
-			ac.mu.Unlock()
-			select {
-			case <-acCtx.Done():
-				as.finish(status.Error(codes.Canceled, "grpc: the SubConn is closing"))
-			case <-ctx.Done():
-				as.finish(toRPCErr(ctx.Err()))
-			}
-		}()
-	}
+	as.registerContextCallbacks()
 	return &clientStreamWrapper{ClientStream: as, desc: desc}, nil
 }
 
+// registerContextCallbacks registers callbacks to clean up the stream when the
+// RPC context expires or the SubConn is closed.
+func (as *addrConnStream) registerContextCallbacks() {
+	as.mu.Lock()
+	defer as.mu.Unlock()
+	if !as.finished {
+		as.stopCtx = context.AfterFunc(as.ctx, func() {
+			as.finish(toRPCErr(as.ctx.Err()))
+		})
+		as.ac.mu.Lock()
+		acCtx := as.ac.ctx
+		as.ac.mu.Unlock()
+		as.stopCC = context.AfterFunc(acCtx, func() {
+			as.finish(status.Error(codes.Canceled, "grpc: the SubConn is closing"))
+		})
+	}
+}
+
 type addrConnStream struct {
 	transportStream  *transport.ClientStream
 	ac               *addrConn
@@ -1515,6 +1528,9 @@ type addrConnStream struct {
 	mu     sync.Mutex
 	parser parser
 
+	stopCtx func() bool // stops the callback registered to run on RPC context expiration
+	stopCC  func() bool // stops the callback registered to run on SubConn context expiration
+
 	// Bool fields are grouped at the tail to eliminate the alignment padding
 	// that would otherwise follow each bool when the next field is pointer- or
 	// int-sized. See https://github.com/grpc/grpc-go/issues/9348 for benchmarks.
@@ -1656,6 +1672,11 @@ func (as *addrConnStream) finish(err error) {
 		return
 	}
 	as.finished = true
+	stopCtx := as.stopCtx
+	stopCC := as.stopCC
+	as.stopCtx = nil
+	as.stopCC = nil
+
 	if err == io.EOF {
 		// Ending a stream with EOF indicates a success.
 		err = nil
@@ -1671,6 +1692,13 @@ func (as *addrConnStream) finish(err error) {
 	}
 	as.cancel()
 	as.mu.Unlock()
+
+	if stopCtx != nil {
+		stopCtx()
+	}
+	if stopCC != nil {
+		stopCC()
+	}
 }
 
 // ServerStream defines the server-side behavior of a streaming RPC.

test/context_canceled_test.go

@@ -162,3 +162,99 @@ func (s) TestCancelWhileRecvingWithCompression(t *testing.T) {
 		t.Fatalf("Close failed with %v, want nil", err)
 	}
 }
+
+// Test verifies that an in-flight Unary RPC fails promptly when the ClientConn
+// is closed (canceling ClientConn context).
+func (s) TestUnary_ClientConnContextExpires(t *testing.T) {
+	rpcStarted := make(chan struct{})
+	ss := &stubserver.StubServer{
+		EmptyCallF: func(ctx context.Context, _ *testpb.Empty) (*testpb.Empty, error) {
+			close(rpcStarted)
+			<-ctx.Done()
+			return nil, ctx.Err()
+		},
+	}
+	if err := ss.Start(nil); err != nil {
+		t.Fatalf("Error starting endpoint server: %v", err)
+	}
+	defer ss.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	errChan := make(chan error, 1)
+	go func() {
+		_, err := ss.Client.EmptyCall(ctx, &testpb.Empty{})
+		errChan <- err
+	}()
+
+	select {
+	case <-rpcStarted:
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("timed out waiting for RPC to start on server")
+	}
+
+	// Close ClientConn while Unary RPC is in-flight.
+	ss.CC.Close()
+
+	select {
+	case err := <-errChan:
+		if status.Code(err) != codes.Canceled {
+			t.Fatalf("EmptyCall returned error: %v (code: %v), want Canceled", err, status.Code(err))
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Unary RPC did not return promptly after ClientConn was closed")
+	}
+}
+
+// Test verifies that an in-flight Streaming RPC fails promptly when the
+// ClientConn is closed (canceling ClientConn context).
+func (s) TestStreaming_ClientConnContextExpires(t *testing.T) {
+	rpcStarted := make(chan struct{})
+	ss := &stubserver.StubServer{
+		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
+			close(rpcStarted)
+			<-stream.Context().Done()
+			return stream.Context().Err()
+		},
+	}
+	if err := ss.Start(nil); err != nil {
+		t.Fatalf("Error starting endpoint server: %v", err)
+	}
+	defer ss.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	stream, err := ss.Client.FullDuplexCall(ctx)
+	if err != nil {
+		t.Fatalf("FullDuplexCall failed: %v", err)
+	}
+
+	// Send an initial message so the stream is active on the server.
+	if err := stream.Send(&testpb.StreamingOutputCallRequest{}); err != nil {
+		t.Fatalf("Send failed: %v", err)
+	}
+
+	select {
+	case <-rpcStarted:
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("timed out waiting for stream to start on server")
+	}
+
+	errChan := make(chan error, 1)
+	go func() {
+		_, err := stream.Recv()
+		errChan <- err
+	}()
+
+	// Close ClientConn while Streaming RPC is waiting in Recv.
+	ss.CC.Close()
+
+	select {
+	case err := <-errChan:
+		if status.Code(err) != codes.Canceled {
+			t.Fatalf("Recv returned error: %v (code: %v), want Canceled", err, status.Code(err))
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Streaming RPC did not return promptly after ClientConn was closed")
+	}
+}