table: add RotateBlock options and WithRotateBlockWaitGroup (#901)

Author: asubiotto

db.go

@@ -844,7 +844,7 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 				// tx == nextNonPersistedTxn, we should not persist the active
 				// block, but just create a new block.
 				table.pendingBlocks[table.active] = struct{}{}
-				go table.writeBlock(table.active, tx, db.columnStore.manualBlockRotation, false)
+				go table.writeBlock(table.active, tx, db.columnStore.manualBlockRotation)
 			}
 
 			protoEqual := false
@@ -1006,7 +1006,7 @@ func (db *DB) Close(options ...CloseOption) error {
 			// should be faster to write to local disk than upload to object
 			// storage. This would avoid a slow WAL replay on startup if we
 			// don't manage to persist in time.
-			table.writeBlock(table.ActiveBlock(), db.tx.Load(), false, false)
+			table.writeBlock(table.ActiveBlock(), db.tx.Load(), false)
 		}
 	}
 	level.Info(db.logger).Log("msg", "closed all tables")

db_test.go

@@ -1206,7 +1206,7 @@ func TestDBRecover(t *testing.T) {
 		// the second write, and the second was triggered before the third write.
 		if blockRotation {
 			// A block rotation should trigger the third snapshot.
-			require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
+			require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock()))
 			// Wait for the snapshot to complete
 			require.Eventually(t, func() bool {
 				files, err := os.ReadDir(db.snapshotsDir())
@@ -1404,7 +1404,7 @@ func TestDBRecover(t *testing.T) {
 		table, err := db.GetTable(dbAndTableName)
 		require.NoError(t, err)
 
-		require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
+		require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock()))
 
 		rec, err := dynparquet.NewTestSamples().ToRecord()
 		require.NoError(t, err)
@@ -1414,7 +1414,7 @@ func TestDBRecover(t *testing.T) {
 
 		// RotateBlock again, this should log a couple of persisted block WAL
 		// entries.
-		require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
+		require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock()))
 		require.NoError(t, c.Close())
 
 		c, err = New(
@@ -1474,7 +1474,7 @@ func TestDBRecover(t *testing.T) {
 		require.NoError(t, err)
 
 		// Rotate the block to create a new active block.
-		require.NoError(t, table.RotateBlock(ctx, block, false))
+		require.NoError(t, table.RotateBlock(ctx, block))
 
 		// Issue writes.
 		const nWrites = 5
@@ -1635,7 +1635,7 @@ func TestDBMinTXPersisted(t *testing.T) {
 	writeTx, err := table.InsertRecord(ctx, r)
 	require.NoError(t, err)
 
-	require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
+	require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock()))
 	// Writing the block is asynchronous, so wait for both the new table block
 	// txn and the block persistence txn.
 	db.Wait(writeTx + 2)
@@ -2657,7 +2657,7 @@ func Test_DB_PersistentDiskCompaction_BlockRotation(t *testing.T) {
 	validateRows(1200)
 
 	// Rotate block
-	require.NoError(t, table.RotateBlock(context.Background(), table.ActiveBlock(), false))
+	require.NoError(t, table.RotateBlock(context.Background(), table.ActiveBlock()))
 
 	validateRows(1200)
 
@@ -3151,7 +3151,7 @@ func Test_Iceberg(t *testing.T) {
 
 	validateRows(10)
 
-	require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
+	require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock()))
 	require.Eventually(t, func() bool {
 		info, err := bucket.Attributes(ctx, filepath.Join("test", "test", "metadata", "v1.metadata.json"))
 		if err != nil {
@@ -3171,7 +3171,7 @@ func Test_Iceberg(t *testing.T) {
 
 	validateRows(13)
 
-	require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
+	require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock()))
 	require.Eventually(t, func() bool {
 		info, err := bucket.Attributes(ctx, filepath.Join("test", "test", "metadata", "v2.metadata.json"))
 		if err != nil {

dst/dst_test.go

@@ -400,7 +400,7 @@ func TestDST(t *testing.T) {
 		case rotate:
 			errg.Go(func() error {
 				table := tp.GetTable()
-				if err := table.RotateBlock(ctx, table.ActiveBlock(), false); err != nil {
+				if err := table.RotateBlock(ctx, table.ActiveBlock()); err != nil {
 					return fmt.Errorf("rotate error: %s", err)
 				}
 				return nil

table.go

@@ -496,7 +496,16 @@ func (t *Table) dropPendingBlock(block *TableBlock) {
 	}
 }
 
-func (t *Table) writeBlock(block *TableBlock, nextTxn uint64, skipPersist, snapshotDB bool) {
+func (t *Table) writeBlock(
+	block *TableBlock, nextTxn uint64, snapshotDB bool, opts ...RotateBlockOption,
+) {
+	rbo := &rotateBlockOptions{}
+	for _, o := range opts {
+		o(rbo)
+	}
+	if rbo.wg != nil {
+		defer rbo.wg.Done()
+	}
 	level.Debug(t.logger).Log("msg", "syncing block", "ulid", block.ulid, "size", block.index.Size())
 	block.pendingWritersWg.Wait()
 
@@ -506,7 +515,7 @@ func (t *Table) writeBlock(block *TableBlock, nextTxn uint64, skipPersist, snaps
 
 	// Persist the block
 	var err error
-	if !skipPersist && block.index.Size() != 0 {
+	if !rbo.skipPersist && block.index.Size() != 0 {
 		err = block.Persist()
 	}
 	t.dropPendingBlock(block)
@@ -605,20 +614,51 @@ func (t *Table) writeBlock(block *TableBlock, nextTxn uint64, skipPersist, snaps
 	}
 }
 
-func (t *Table) RotateBlock(_ context.Context, block *TableBlock, skipPersist bool) error {
+type rotateBlockOptions struct {
+	skipPersist bool
+	wg          *sync.WaitGroup
+}
+
+type RotateBlockOption func(*rotateBlockOptions)
+
+// WithRotateBlockSkipPersist instructs the block rotation operation to not
+// persist the block to object storage.
+func WithRotateBlockSkipPersist() RotateBlockOption {
+	return func(o *rotateBlockOptions) {
+		o.skipPersist = true
+	}
+}
+
+// WithRotateBlockWaitGroup provides a WaitGroup. The rotate block operation
+// will call wg.Done once the block has been persisted. Otherwise, RotateBlock
+// asynchronously persists the block.
+func WithRotateBlockWaitGroup(wg *sync.WaitGroup) RotateBlockOption {
+	return func(o *rotateBlockOptions) {
+		o.wg = wg
+	}
+}
+
+func (t *Table) RotateBlock(_ context.Context, block *TableBlock, opts ...RotateBlockOption) error {
+	rbo := &rotateBlockOptions{}
+	for _, o := range opts {
+		o(rbo)
+	}
 	t.mtx.Lock()
 	defer t.mtx.Unlock()
 
 	// Need to check that we haven't already rotated this block.
 	if t.active != block {
+		if rbo.wg != nil {
+			rbo.wg.Done()
+		}
 		return nil
 	}
 
 	level.Debug(t.logger).Log(
 		"msg", "rotating block",
 		"ulid", block.ulid,
 		"size", block.Size(),
-		"skip_persist", skipPersist,
+		"skip_persist", rbo.skipPersist,
 	)
 	defer level.Debug(t.logger).Log("msg", "done rotating block", "ulid", block.ulid)
 
@@ -637,7 +677,7 @@ func (t *Table) RotateBlock(_ context.Context, block *TableBlock, skipPersist bo
 	t.metrics.blockRotated.Inc()
 	t.metrics.numParts.Set(float64(0))
 
-	if !skipPersist {
+	if !rbo.skipPersist {
 		// If skipping persist, this block rotation is simply a block discard,
 		// so no need to add this block to pending blocks. Some callers rely
 		// on the fact that blocks are not available for reads as soon as
@@ -647,7 +687,7 @@ func (t *Table) RotateBlock(_ context.Context, block *TableBlock, skipPersist bo
 	// We don't check t.db.columnStore.manualBlockRotation here because this is
 	// the entry point for users to trigger a manual block rotation and they
 	// will specify through skipPersist if they want the block to be persisted.
-	go t.writeBlock(block, tx, skipPersist, true)
+	go t.writeBlock(block, tx, true, opts...)
 
 	return nil
 }
@@ -750,7 +790,7 @@ func (t *Table) appender(ctx context.Context) (*TableBlock, func(), error) {
 		// We need to rotate the block and the writer won't actually be used.
 		finish()
 
-		err = t.RotateBlock(ctx, block, false)
+		err = t.RotateBlock(ctx, block)
 		if err != nil {
 			return nil, nil, fmt.Errorf("rotate block: %w", err)
 		}