blobs: fix resource leak in GetStream reader close

Previously, closing a GetStream reader before reading all data was a
no-op, leaving gRPC streams and goroutines active. This caused resource
leaks and OOM errors in distributed merge operations.

Now, ReadFile creates a cancelable context and passes the cancel
function to the stream reader, which calls it on Close() to properly
terminate the stream.

Release note: none
Epic: CRDB-48845

Author: jeffswenson

pkg/blobs/client.go

@@ -63,11 +63,16 @@ func (c *remoteClient) ReadFile(
 	if err != nil {
 		return nil, 0, err
 	}
+	ctx, cancel := context.WithCancel(ctx)
 	stream, err := c.blobClient.GetStream(ctx, &blobspb.GetRequest{
 		Filename: file,
 		Offset:   offset,
 	})
-	return newGetStreamReader(stream), st.Filesize, errors.Wrap(err, "fetching file")
+	if err != nil {
+		cancel()
+		return nil, 0, errors.Wrap(err, "fetching file")
+	}
+	return newGetStreamReader(stream, cancel), st.Filesize, nil
 }
 
 type streamWriter struct {

pkg/blobs/client_test.go

@@ -543,3 +543,32 @@ func TestBlobClientStat(t *testing.T) {
 		})
 	}
 }
+
+type mockClient struct {
+	blobspb.RPCBlobClient
+	ctx context.Context
+}
+
+func (m *mockClient) GetStream(
+	ctx context.Context, in *blobspb.GetRequest,
+) (blobspb.RPCBlob_GetStreamClient, error) {
+	m.ctx = ctx
+	return nil, nil
+}
+
+func (m *mockClient) Stat(ctx context.Context, in *blobspb.StatRequest) (*blobspb.BlobStat, error) {
+	return &blobspb.BlobStat{}, nil
+}
+
+// TestBlobClientReadFileClose tests that closing a reader cancels the context
+// passed to the rpc stream.
+func TestBlobClientReadFileClose(t *testing.T) {
+	mock := &mockClient{}
+	r := remoteClient{blobClient: mock}
+
+	reader, _, err := r.ReadFile(context.Background(), "doesn't matter", 0)
+	require.NoError(t, err)
+	require.NoError(t, reader.Close(context.Background()))
+
+	<-mock.ctx.Done()
+}

pkg/blobs/stream.go

@@ -39,22 +39,24 @@ type streamReceiver interface {
 	Recv() (*blobspb.StreamChunk, error)
 }
 
-// nopSendAndClose creates a GetStreamClient that has a nop SendAndClose function.
-// This is needed as Blob_GetStreamClient does not have a Close() function, whereas
-// the other sender, Blob_PutStreamServer, does.
-type nopSendAndClose struct {
+// closeFuncSendAndClose creates a GetStreamClient that calls a close function
+// when SendAndClose is called. The close function is expected to cancel the context
+// and cause the stream to clean itself up.
+type closeFuncSendAndClose struct {
 	blobspb.RPCBlob_GetStreamClient
+	close func()
 }
 
-func (*nopSendAndClose) SendAndClose(*blobspb.StreamResponse) error {
+func (c *closeFuncSendAndClose) SendAndClose(*blobspb.StreamResponse) error {
+	c.close()
 	return nil
 }
 
 // newGetStreamReader creates an io.ReadCloser that uses gRPC's streaming API
 // to read chunks of data.
-func newGetStreamReader(client blobspb.RPCBlob_GetStreamClient) ioctx.ReadCloserCtx {
+func newGetStreamReader(client blobspb.RPCBlob_GetStreamClient, close func()) ioctx.ReadCloserCtx {
 	return &blobStreamReader{
-		stream: &nopSendAndClose{client},
+		stream: &closeFuncSendAndClose{client, close},
 	}
 }