Create memhttp package to debug flaky testcases

connect_ext_test.go

@@ -183,14 +183,28 @@ func TestServer(t *testing.T) {
 				connect.CodeOf(stream.Err()),
 				connect.CodeInvalidArgument,
 			)
+			assert.Nil(t, stream.Close())
 		})
 		t.Run("count_up_timeout", func(t *testing.T) {
 			ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
-			defer cancel()
+			t.Cleanup(cancel)
 			_, err := client.CountUp(ctx, connect.NewRequest(&pingv1.CountUpRequest{Number: 1}))
 			assert.NotNil(t, err)
 			assert.Equal(t, connect.CodeOf(err), connect.CodeDeadlineExceeded)
 		})
+		t.Run("count_up_cancel_after_first_response", func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			request := connect.NewRequest(&pingv1.CountUpRequest{Number: 5})
+			request.Header().Set(clientHeader, headerValue)
+			stream, err := client.CountUp(ctx, request)
+			assert.Nil(t, err)
+			assert.True(t, stream.Receive())
+			cancel()
+			assert.False(t, stream.Receive())
+			assert.NotNil(t, stream.Err())
+			assert.Equal(t, connect.CodeOf(stream.Err()), connect.CodeCanceled)
+			assert.Nil(t, stream.Close())
+		})
 	}
 	testCumSum := func(t *testing.T, client pingv1connect.PingServiceClient, expectSuccess bool) { //nolint:thelper
 		t.Run("cumsum", func(t *testing.T) {
@@ -285,6 +299,7 @@ func TestServer(t *testing.T) {
 			assert.Equal(t, connect.CodeOf(err), connect.CodeCanceled)
 			assert.Equal(t, got, expect)
 			assert.False(t, connect.IsWireError(err))
+			assert.Nil(t, stream.CloseResponse())
 		})
 		t.Run("cumsum_cancel_before_send", func(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.Background())
@@ -302,6 +317,8 @@ func TestServer(t *testing.T) {
 			err := stream.Send(&pingv1.CumSumRequest{Number: 19})
 			assert.Equal(t, connect.CodeOf(err), connect.CodeCanceled, assert.Sprintf("%v", err))
 			assert.False(t, connect.IsWireError(err))
+			assert.Nil(t, stream.CloseRequest())
+			assert.Nil(t, stream.CloseResponse())
 		})
 	}
 	testErrors := func(t *testing.T, client pingv1connect.PingServiceClient) { //nolint:thelper

duplex_http_call.go

@@ -42,14 +42,15 @@ type duplexHTTPCall struct {
 	requestBodyReader *io.PipeReader
 	requestBodyWriter *io.PipeWriter
 
+	// sendRequestOnce ensures we only send the request once.
 	sendRequestOnce sync.Once
 	request         *http.Request
-	response        *http.Response
 
 	// responseReady is closed when the response is ready or when the request
 	// fails. Any error on request initialisation will be set on the
 	// responseErr. There's always a response if responseErr is nil.
 	responseReady chan struct{}
+	response      *http.Response
 	responseErr   error
 }
 
@@ -177,7 +178,9 @@ func (d *duplexHTTPCall) CloseRead() error {
 	if d.response == nil {
 		return nil
 	}
-	if _, err := discard(d.response.Body); err != nil {
+	if _, err := discard(d.response.Body); err != nil &&
+		!errors.Is(err, context.Canceled) &&
+		!errors.Is(err, context.DeadlineExceeded) {
 		_ = d.response.Body.Close()
 		return wrapIfRSTError(err)
 	}
@@ -247,6 +250,9 @@ func (d *duplexHTTPCall) makeRequest() {
 	}
 	// Once we send a message to the server, they send a message back and
 	// establish the receive side of the stream.
+	// On error, we close the request body using the Write side of the pipe.
+	// This ensures HTTP2 streams receive an io.EOF from the Read side of the
+	// pipe. Write's check for io.ErrClosedPipe and will convert this to io.EOF.
 	response, err := d.httpClient.Do(d.request) //nolint:bodyclose
 	if err != nil {
 		err = wrapIfContextError(err)
@@ -257,15 +263,15 @@ func (d *duplexHTTPCall) makeRequest() {
 			err = NewError(CodeUnavailable, err)
 		}
 		d.responseErr = err
-		d.requestBodyReader.CloseWithError(io.EOF)
+		d.requestBodyWriter.Close()
 		return
 	}
 	// We've got a response. We can now read from the response body.
-	// Closing the response body is delegated to the caller.
+	// Closing the response body is delegated to the caller even on error.
 	d.response = response
 	if err := d.validateResponse(response); err != nil {
 		d.responseErr = err
-		d.requestBodyReader.CloseWithError(io.EOF)
+		d.requestBodyWriter.Close()
 		return
 	}
 	if (d.streamType&StreamTypeBidi) == StreamTypeBidi && response.ProtoMajor < 2 {
@@ -278,7 +284,7 @@ func (d *duplexHTTPCall) makeRequest() {
 			response.ProtoMajor,
 			response.ProtoMinor,
 		)
-		d.requestBodyReader.CloseWithError(io.EOF)
+		d.requestBodyWriter.Close()
 	}
 }