otelgrpc: fix goroutine leak after context cancellation

instrumentation/google.golang.org/grpc/otelgrpc/interceptor.go

@@ -200,7 +200,7 @@ const (
 	receiveEndedState
 )
 
-func wrapClientStream(s grpc.ClientStream, desc *grpc.StreamDesc) *clientStream {
+func wrapClientStream(ctx context.Context, s grpc.ClientStream, desc *grpc.StreamDesc) *clientStream {
 	events := make(chan streamEvent)
 	eventsDone := make(chan struct{})
 	finished := make(chan error)
@@ -211,19 +211,25 @@ func wrapClientStream(s grpc.ClientStream, desc *grpc.StreamDesc) *clientStream
 		// Both streams have to be closed
 		state := byte(0)
 
-		for event := range events {
-			switch event.Type {
-			case closeEvent:
-				state |= clientClosedState
-			case receiveEndEvent:
-				state |= receiveEndedState
-			case errorEvent:
-				finished <- event.Err
-				return
-			}
-
-			if state == clientClosedState|receiveEndedState {
-				finished <- nil
+		for {
+			select {
+			case event := <-events:
+				switch event.Type {
+				case closeEvent:
+					state |= clientClosedState
+				case receiveEndEvent:
+					state |= receiveEndedState
+				case errorEvent:
+					finished <- event.Err
+					return
+				}
+				if state == clientClosedState|receiveEndedState {
+					finished <- nil
+					return
+				}
+
+			case <-ctx.Done():
+				finished <- ctx.Err()
 				return
 			}
 		}
@@ -284,7 +290,7 @@ func StreamClientInterceptor(opts ...Option) grpc.StreamClientInterceptor {
 			span.End()
 			return s, err
 		}
-		stream := wrapClientStream(s, desc)
+		stream := wrapClientStream(ctx, s, desc)
 
 		go func() {
 			err := <-stream.finished

instrumentation/google.golang.org/grpc/otelgrpc/interceptor_test.go

@@ -372,6 +372,59 @@ func TestStreamClientInterceptor(t *testing.T) {
 	_ = streamClient.CloseSend()
 }
 
+// TestStreamClientInterceptorCancelContext tests a cancel context situation.
+// There should be no goleaks
+func TestStreamClientInterceptorCancelContext(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	clientConn, err := grpc.Dial("fake:connection", grpc.WithInsecure())
+	if err != nil {
+		t.Fatalf("failed to create client connection: %v", err)
+	}
+	defer clientConn.Close()
+
+	// tracer
+	sr := NewSpanRecorder()
+	tp := oteltest.NewTracerProvider(oteltest.WithSpanRecorder(sr))
+	streamCI := StreamClientInterceptor(WithTracerProvider(tp))
+
+	var mockClStr mockClientStream
+	method := "/github.com.serviceName/bar"
+	name := "github.com.serviceName/bar"
+
+	// create a context with cancel
+	cancelCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	streamClient, err := streamCI(
+		cancelCtx,
+		&grpc.StreamDesc{ServerStreams: true},
+		clientConn,
+		method,
+		func(ctx context.Context,
+			desc *grpc.StreamDesc,
+			cc *grpc.ClientConn,
+			method string,
+			opts ...grpc.CallOption) (grpc.ClientStream, error) {
+			mockClStr = mockClientStream{Desc: desc, Ctx: ctx}
+			return mockClStr, nil
+		},
+	)
+	require.NoError(t, err, "initialize grpc stream client")
+	_, ok := getSpanFromRecorder(sr, name)
+	require.False(t, ok, "span should not ended while stream is open")
+
+	req := &mockProtoMessage{}
+	reply := &mockProtoMessage{}
+
+	// send and receive fake data
+	for i := 0; i < 10; i++ {
+		_ = streamClient.SendMsg(req)
+		_ = streamClient.RecvMsg(reply)
+	}
+
+	// close client stream
+	_ = streamClient.CloseSend()
+}
+
 // TestStreamClientInterceptorWithError tests a situation that streamer returns an error.
 func TestStreamClientInterceptorWithError(t *testing.T) {
 	defer goleak.VerifyNone(t)