Merge #156867

156867: sql/schemachanger: flush batches while making metadata changes r=fqazi a=fqazi

Previously, the declarative schema changer only flushed batches in a transaction at the end of a stage. This could lead to command sizes that are too large prevent large truncates from completing. This patch, flushes the batches once they are above a certain size.

Fixes: #156682

Release note (bug fix): Addressed a bug that prevents large TRUNCATE operations from completing due to `command is too large` errors.

Co-authored-by: Faizan Qazi <[email protected]>

Author: craig[bot]

pkg/sql/drop_test.go

@@ -1552,6 +1552,47 @@ func TestDropLargeDatabaseWithDeclarativeSchemaChanger(t *testing.T) {
 		true)
 }
 
+// TestTruncateLarge truncates a large number of tables in a single transaction.
+func TestTruncateLarge(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	skip.UnderDuress(t, "truncating a large number of tables")
+
+	testutils.RunTrueAndFalse(t, "batch limit set", func(t *testing.T, batchLimitSet bool) {
+		ctx := context.Background()
+		s, conn, _ := serverutils.StartServer(t, base.TestServerArgs{})
+		defer s.Stopper().Stop(ctx)
+		sqlDB := sqlutils.MakeSQLRunner(conn)
+		createCommand := strings.Builder{}
+		truncateCommand := strings.Builder{}
+		sqlDB.Exec(t, "SET CLUSTER SETTING kv.raft.command.max_size='5m'")
+		if batchLimitSet {
+			sqlDB.Exec(t, "SET CLUSTER SETTING sql.schema_changer.batch_flush_threshold_size='2m'")
+		}
+		// Generate the truncate and create table commands.
+		const numTables = 500
+		truncateCommand.WriteString("TRUNCATE TABLE ")
+		for i := range numTables {
+			createCommand.WriteString(fmt.Sprintf("CREATE TABLE t%d (a INT PRIMARY KEY, j INT, k INT, INDEX (j), INDEX (k), UNIQUE (j, k));\n", i))
+			truncateCommand.WriteString(fmt.Sprintf("t%d", i))
+			if i != numTables-1 {
+				truncateCommand.WriteString(", ")
+			}
+		}
+		// Execute the create commands first.
+		sqlDB.Exec(t, createCommand.String())
+
+		// The default limit is larger than our reduced raft command size, so if the
+		// limit is not modified, the truncate command will fail.
+		if batchLimitSet {
+			sqlDB.Exec(t, truncateCommand.String())
+		} else {
+			sqlDB.ExpectErr(t, "command is too large", truncateCommand.String())
+		}
+	})
+
+}
+
 func BenchmarkDropLargeDatabaseWithGenerateTestObjects(b *testing.B) {
 	defer leaktest.AfterTest(b)()
 	defer log.Scope(b).Close(b)

pkg/sql/schemachanger/scdeps/BUILD.bazel

@@ -19,6 +19,7 @@ go_library(
         "//pkg/roachpb",
         "//pkg/security/username",
         "//pkg/server/telemetry",
+        "//pkg/settings",
         "//pkg/settings/cluster",
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/catalogkeys",

pkg/sql/schemachanger/scdeps/exec_deps.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
+	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
@@ -35,6 +36,17 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
+// batchFlushThresholdSize is the size of the metadata batch that,
+// when exceeded, causes the schema changer to flush the batch to the KV store.
+var batchFlushThresholdSize = settings.RegisterByteSizeSetting(
+	settings.ApplicationLevel,
+	"sql.schema_changer.batch_flush_threshold_size",
+	"maximum size in bytes of the schema changer's metadata batch before it's flushed to the KV store. "+
+		"This setting should be smaller or equal to kv.raft.command.max_size",
+	32*1024*1024,
+	settings.IntInRange(1024*1024, 512*1024*1024),
+)
+
 // JobRegistry implements the methods the schema changer needs from the
 // job registry. Outside of tests this should always be backed by *job.Registry.
 type JobRegistry interface {
@@ -206,22 +218,38 @@ func (d *txnDeps) MustReadMutableDescriptor(
 func (d *txnDeps) CreateOrUpdateDescriptor(
 	ctx context.Context, desc catalog.MutableDescriptor,
 ) error {
-	return d.descsCollection.WriteDescToBatch(ctx, d.kvTrace, desc, d.getOrCreateBatch())
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.WriteDescToBatch(ctx, d.kvTrace, desc, b)
 }
 
 // DeleteName implements the scexec.Catalog interface.
 func (d *txnDeps) DeleteName(ctx context.Context, nameInfo descpb.NameInfo, id descpb.ID) error {
-	return d.descsCollection.DeleteNamespaceEntryToBatch(ctx, d.kvTrace, &nameInfo, d.getOrCreateBatch())
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.DeleteNamespaceEntryToBatch(ctx, d.kvTrace, &nameInfo, b)
 }
 
 // AddName implements the scexec.Catalog interface.
 func (d *txnDeps) AddName(ctx context.Context, nameInfo descpb.NameInfo, id descpb.ID) error {
-	return d.descsCollection.InsertNamespaceEntryToBatch(ctx, d.kvTrace, &nameEntry{nameInfo, id}, d.getOrCreateBatch())
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.InsertNamespaceEntryToBatch(ctx, d.kvTrace, &nameEntry{nameInfo, id}, b)
 }
 
 // DeleteDescriptor implements the scexec.Catalog interface.
 func (d *txnDeps) DeleteDescriptor(ctx context.Context, id descpb.ID) error {
-	return d.descsCollection.DeleteDescToBatch(ctx, d.kvTrace, id, d.getOrCreateBatch())
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.DeleteDescToBatch(ctx, d.kvTrace, id, b)
 }
 
 // GetZoneConfig implements the scexec.Catalog interface.
@@ -237,7 +265,11 @@ func (d *txnDeps) GetZoneConfig(ctx context.Context, id descpb.ID) (catalog.Zone
 func (d *txnDeps) WriteZoneConfigToBatch(
 	ctx context.Context, id descpb.ID, zc catalog.ZoneConfig,
 ) error {
-	err := d.descsCollection.WriteZoneConfigToBatch(ctx, d.kvTrace, d.getOrCreateBatch(), id, zc)
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	err = d.descsCollection.WriteZoneConfigToBatch(ctx, d.kvTrace, b, id, zc)
 	if err != nil {
 		return err
 	}
@@ -260,7 +292,11 @@ func (d *txnDeps) UpdateZoneConfig(ctx context.Context, id descpb.ID, zc *zonepb
 		rawBytes = oldZc.GetRawBytesInStorage()
 	}
 	newZc = zone.NewZoneConfigWithRawBytes(zc, rawBytes)
-	return d.descsCollection.WriteZoneConfigToBatch(ctx, d.kvTrace, d.getOrCreateBatch(), id, newZc)
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.WriteZoneConfigToBatch(ctx, d.kvTrace, b, id, newZc)
 }
 
 // UpdateSubzoneConfig implements the scexec.Catalog interface. Note that this
@@ -310,7 +346,11 @@ func (d *txnDeps) UpdateSubzoneConfig(
 
 // DeleteZoneConfig implements the scexec.Catalog interface.
 func (d *txnDeps) DeleteZoneConfig(ctx context.Context, id descpb.ID) error {
-	return d.descsCollection.DeleteZoneConfigInBatch(ctx, d.kvTrace, d.getOrCreateBatch(), id)
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.DeleteZoneConfigInBatch(ctx, d.kvTrace, b, id)
 }
 
 // DeleteSubzoneConfig implements the scexec.Catalog interface.
@@ -344,7 +384,11 @@ func (d *txnDeps) DeleteSubzoneConfig(
 	zc.DeleteSubzoneSpans(subzoneSpans)
 
 	newZc = zone.NewZoneConfigWithRawBytes(zc, rawBytes)
-	return d.descsCollection.WriteZoneConfigToBatch(ctx, d.kvTrace, d.getOrCreateBatch(),
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.WriteZoneConfigToBatch(ctx, d.kvTrace, b,
 		tableID, newZc)
 }
 
@@ -369,10 +413,14 @@ func (d *txnDeps) Run(ctx context.Context) error {
 }
 
 // InitializeSequence implements the scexec.Caatalog interface.
-func (d *txnDeps) InitializeSequence(id descpb.ID, startVal int64) {
-	batch := d.getOrCreateBatch()
+func (d *txnDeps) InitializeSequence(ctx context.Context, id descpb.ID, startVal int64) error {
+	batch, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
 	sequenceKey := d.codec.SequenceKey(uint32(id))
 	batch.Inc(sequenceKey, startVal)
+	return nil
 }
 
 // CheckMaxSchemaObjects implements the scexec.Catalog interface.
@@ -394,21 +442,45 @@ func (d *txnDeps) Reset(ctx context.Context) error {
 	return nil
 }
 
-func (d *txnDeps) getOrCreateBatch() *kv.Batch {
+// maybeFlushBatch flushes the current batch if it exceeds the maximum size.
+func (d *txnDeps) maybeFlushBatch(ctx context.Context) error {
+	if int64(d.batch.ApproximateMutationBytes()) > batchFlushThresholdSize.Get(&d.settings.SV) {
+		if err := d.Run(ctx); err != nil {
+			return err
+		}
+		d.batch = d.txn.KV().NewBatch()
+	}
+	return nil
+}
+
+func (d *txnDeps) getOrCreateBatch(ctx context.Context) (*kv.Batch, error) {
 	if d.batch == nil {
 		d.batch = d.txn.KV().NewBatch()
+	} else {
+		// Otherwise, flush the batch if its too big.
+		if err := d.maybeFlushBatch(ctx); err != nil {
+			return nil, err
+		}
 	}
-	return d.batch
+	return d.batch, nil
 }
 
 // UpdateComment implements the scexec.Catalog interface.
 func (d *txnDeps) UpdateComment(ctx context.Context, key catalogkeys.CommentKey, cmt string) error {
-	return d.descsCollection.WriteCommentToBatch(ctx, d.kvTrace, d.getOrCreateBatch(), key, cmt)
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.WriteCommentToBatch(ctx, d.kvTrace, b, key, cmt)
 }
 
 // DeleteComment implements the scexec.Catalog interface.
 func (d *txnDeps) DeleteComment(ctx context.Context, key catalogkeys.CommentKey) error {
-	return d.descsCollection.DeleteCommentInBatch(ctx, d.kvTrace, d.getOrCreateBatch(), key)
+	b, err := d.getOrCreateBatch(ctx)
+	if err != nil {
+		return err
+	}
+	return d.descsCollection.DeleteCommentInBatch(ctx, d.kvTrace, b, key)
 }
 
 var _ scexec.TransactionalJobRegistry = (*txnDeps)(nil)

pkg/sql/schemachanger/scdeps/sctestdeps/test_deps.go

@@ -1543,8 +1543,9 @@ func getNameEntryDescriptorType(parentID, parentSchemaID descpb.ID) string {
 }
 
 // InitializeSequence is part of the scexec.Catalog interface.
-func (s *TestState) InitializeSequence(id descpb.ID, startVal int64) {
+func (s *TestState) InitializeSequence(ctx context.Context, id descpb.ID, startVal int64) error {
 	s.LogSideEffectf("initializing sequence %d with starting value of %d", id, startVal)
+	return nil
 }
 
 // CheckMaxSchemaObjects is part of the scexec.Catalog interface.

pkg/sql/schemachanger/scexec/dependencies.go

@@ -118,7 +118,7 @@ type Catalog interface {
 	Reset(ctx context.Context) error
 
 	// InitializeSequence initializes the initial value for a sequence.
-	InitializeSequence(id descpb.ID, startVal int64)
+	InitializeSequence(ctx context.Context, id descpb.ID, startVal int64) error
 
 	// CheckMaxSchemaObjects checks if the number of schema objects in the
 	// cluster plus the new objects being created would exceed the configured

pkg/sql/schemachanger/scexec/exec_immediate_mutation.go

@@ -272,7 +272,9 @@ func (s *immediateState) exec(ctx context.Context, c Catalog) error {
 		}
 	}
 	for _, s := range s.sequencesToInit {
-		c.InitializeSequence(s.id, s.startVal)
+		if err := c.InitializeSequence(ctx, s.id, s.startVal); err != nil {
+			return err
+		}
 	}
 	for tempIdxId, tempIdxToRegister := range s.temporarySchemasToRegister {
 		c.InsertTemporarySchema(tempIdxToRegister.schemaName, tempIdxToRegister.parentID, tempIdxId)