Merge #160627 #160659

160627: blobs: fix resource leak in GetStream reader close r=jeffswenson a=jeffswenson

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

160659: workload/schemachanger: fix dependency detection for DROP COLUMN r=fqazi a=fqazi

Previously, the DROP COLUMN operation in the schema changer could fail, since it detected self referential foreign key dependencies as external ones. To address this, this patch excludes them from the detection.

Fixes: #160318
Fixes: #159952

Release note: None

Co-authored-by: Jeff Swenson <jeffswenson@betterthannull.com>
Co-authored-by: Faizan Qazi <faizan@cockroachlabs.com>

Author: 3 people

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},
 	}
 }
 

pkg/workload/schemachange/error_screening.go

@@ -150,7 +150,7 @@ func (og *operationGenerator) tableHasDependencies(
 }
 
 func (og *operationGenerator) columnIsDependedOn(
-	ctx context.Context, tx pgx.Tx, tableName *tree.TableName, columnName tree.Name,
+	ctx context.Context, tx pgx.Tx, tableName *tree.TableName, columnName tree.Name, includeFKs bool,
 ) (bool, error) {
 	// To see if a column is depended on, the ordinal_position of the column is looked up in
 	// information_schema.columns. Then, this position is used to see if that column has view dependencies
@@ -187,6 +187,7 @@ func (og *operationGenerator) columnIsDependedOn(
 			            WHERE fd.descriptor_id
 			                  = $1::REGCLASS
                     AND fd.dependedonby_type != 'sequence'
+										AND ($5::BOOL || fd.dependedonby_type != 'fk')
 			          )
 			   UNION  (
 			           SELECT unnest(confkey) AS column_id
@@ -208,7 +209,7 @@ func (og *operationGenerator) columnIsDependedOn(
 			      AND column_name = $4
 			  ) AS source ON source.column_id = cons.column_id
 )
-`, tableName.String(), tableName.Schema(), tableName.Object(), columnName)
+`, tableName.String(), tableName.Schema(), tableName.Object(), columnName, includeFKs)
 }
 
 // colIsRefByComputed determines if a column is referenced by a computed column.

pkg/workload/schemachange/operation_generator.go

@@ -1713,7 +1713,7 @@ func (og *operationGenerator) dropColumn(ctx context.Context, tx pgx.Tx) (*opStm
 	if err != nil {
 		return nil, err
 	}
-	columnIsDependedOn, err := og.columnIsDependedOn(ctx, tx, tableName, columnName)
+	columnIsDependedOn, err := og.columnIsDependedOn(ctx, tx, tableName, columnName, false /* includeFKs */)
 	if err != nil {
 		return nil, err
 	}
@@ -2675,7 +2675,7 @@ func (og *operationGenerator) setColumnType(ctx context.Context, tx pgx.Tx) (*op
 		return nil, err
 	}
 
-	columnHasDependencies, err := og.columnIsDependedOn(ctx, tx, tableName, columnForTypeChange.name)
+	columnHasDependencies, err := og.columnIsDependedOn(ctx, tx, tableName, columnForTypeChange.name, true /* includeFKs */)
 	if err != nil {
 		return nil, err
 	}