Merge #158655

158655: bulkmerge: implement task distribution and loopback coordination r=mw5h a=mw5h

Previously, the bulkmerge processors had only a skeleton implementation. This change implements the core coordination logic:

1. Task Distribution: The merge coordinator now maintains a task set and distributes tasks across worker SQL instances using a loopback channel. As workers complete tasks, they are dynamically assigned new tasks until all tasks are processed.

2. Loopback Communication: Implements a loopback processor that receives task assignments from the coordinator and routes them back through the merge processors. Uses a flow-scoped channel map to enable communication between the coordinator and loopback processors.

3. Result Aggregation: The coordinator collects SST outputs from all processors and aggregates them into a final result proto.

4. Testing: Adds comprehensive tests for both single-node and multi-node scenarios, verifying that all tasks are processed and distributed across available instances.

The implementation uses protobuf marshaling for SST output and includes proper error handling throughout the pipeline.

Epic: None

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Matt White <matt.white@cockroachlabs.com>

Author: craig[bot]

pkg/sql/bulkmerge/BUILD.bazel

@@ -25,6 +25,9 @@ go_library(
         "//pkg/sql/rowexec",
         "//pkg/sql/sem/tree",
         "//pkg/sql/types",
+        "//pkg/util/protoutil",
+        "//pkg/util/syncutil",
+        "//pkg/util/taskset",
         "@com_github_cockroachdb_errors//:errors",
     ],
 )
@@ -55,6 +58,7 @@ go_test(
         "//pkg/testutils/testcluster",
         "//pkg/util/leaktest",
         "//pkg/util/log",
+        "//pkg/util/protoutil",
         "//pkg/util/randutil",
         "@com_github_stretchr_testify//require",
     ],

pkg/sql/bulkmerge/merge_coordinator.go

@@ -14,6 +14,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/rowexec"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+	"github.com/cockroachdb/cockroach/pkg/util/taskset"
 	"github.com/cockroachdb/errors"
 )
 
@@ -31,34 +33,152 @@ var mergeCoordinatorOutputTypes = []*types.T{
 
 type mergeCoordinator struct {
 	execinfra.ProcessorBase
+
 	input execinfra.RowSource
+	spec  execinfrapb.MergeCoordinatorSpec
+	tasks taskset.TaskSet
+
+	loopback chan rowenc.EncDatumRow
+	cleanup  func()
+
+	done    bool
+	results execinfrapb.BulkMergeSpec_Output
+}
+
+type mergeCoordinatorInput struct {
+	sqlInstanceID string
+	taskID        taskset.TaskID
+	outputSSTs    []execinfrapb.BulkMergeSpec_SST
+}
+
+// parseCoordinatorInput ensures each column has the correct type and unmarshals
+// the output SSTs.
+func parseCoordinatorInput(row rowenc.EncDatumRow) (mergeCoordinatorInput, error) {
+	if len(row) != 3 {
+		return mergeCoordinatorInput{}, errors.Newf("expected 3 columns, got %d", len(row))
+	}
+	if err := row[0].EnsureDecoded(types.Bytes, nil); err != nil {
+		return mergeCoordinatorInput{}, err
+	}
+	sqlInstanceID, ok := row[0].Datum.(*tree.DBytes)
+	if !ok {
+		return mergeCoordinatorInput{},
+			errors.Newf("expected bytes column for sqlInstanceID, got %s", row[0].Datum.String())
+	}
+	if err := row[1].EnsureDecoded(types.Int4, nil); err != nil {
+		return mergeCoordinatorInput{}, err
+	}
+	taskID, ok := row[1].Datum.(*tree.DInt)
+	if !ok {
+		return mergeCoordinatorInput{},
+			errors.Newf("expected int4 column for taskID, got %s", row[1].Datum.String())
+	}
+	if err := row[2].EnsureDecoded(types.Bytes, nil); err != nil {
+		return mergeCoordinatorInput{}, err
+	}
+	outputBytes, ok := row[2].Datum.(*tree.DBytes)
+	if !ok {
+		return mergeCoordinatorInput{},
+			errors.Newf("expected bytes column for outputSSTs, got %s", row[2].Datum.String())
+	}
+	results := execinfrapb.BulkMergeSpec_Output{}
+	if err := protoutil.Unmarshal([]byte(*outputBytes), &results); err != nil {
+		return mergeCoordinatorInput{}, err
+	}
+	return mergeCoordinatorInput{
+		sqlInstanceID: string(*sqlInstanceID),
+		taskID:        taskset.TaskID(*taskID),
+		outputSSTs:    results.SSTs,
+	}, nil
 }
 
 // Next implements execinfra.RowSource.
 func (m *mergeCoordinator) Next() (rowenc.EncDatumRow, *execinfrapb.ProducerMetadata) {
 	for m.State == execinfra.StateRunning {
 		row, meta := m.input.Next()
 		switch {
-		case row == nil && meta == nil:
-			m.MoveToDraining(nil /* err */)
-		case meta != nil && meta.Err != nil:
+		case row != nil:
+			err := m.handleRow(row)
+			if err != nil {
+				m.MoveToDraining(err)
+			}
+		case meta == nil:
+			if m.done {
+				m.MoveToDraining(nil /* err */)
+				break
+			}
+			m.done = true
+			return m.emitResults()
+		case meta.Err != nil:
 			m.MoveToDraining(meta.Err)
-		case meta != nil:
+		default:
 			m.MoveToDraining(errors.Newf("unexpected meta: %v", meta))
-		case row != nil:
-			base := *row[2].Datum.(*tree.DBytes)
-			return rowenc.EncDatumRow{
-				rowenc.EncDatum{Datum: tree.NewDBytes(base + "->coordinator")},
-			}, nil
 		}
 	}
 	return nil, m.DrainHelper()
 }
 
+func (m *mergeCoordinator) emitResults() (rowenc.EncDatumRow, *execinfrapb.ProducerMetadata) {
+	marshaled, err := protoutil.Marshal(&m.results)
+	if err != nil {
+		m.MoveToDraining(errors.Wrap(err, "failed to marshal results"))
+		return nil, m.DrainHelper()
+	}
+	return rowenc.EncDatumRow{
+		rowenc.EncDatum{Datum: tree.NewDBytes(tree.DBytes(marshaled))},
+	}, nil
+}
+
+func (m *mergeCoordinator) publishInitialTasks() {
+	for _, sqlInstanceID := range m.spec.WorkerSqlInstanceIds {
+		taskID := m.tasks.ClaimFirst()
+		if taskID.IsDone() {
+			m.closeLoopback()
+			return
+		}
+		m.loopback <- rowenc.EncDatumRow{
+			rowenc.EncDatum{Datum: tree.NewDBytes(tree.DBytes(sqlInstanceID))},
+			rowenc.EncDatum{Datum: tree.NewDInt(tree.DInt(taskID))},
+		}
+	}
+}
+
+func (m *mergeCoordinator) closeLoopback() {
+	if m.cleanup != nil {
+		m.cleanup()
+		m.cleanup = nil
+	}
+}
+
+// handleRow accepts a row output by the merge processor, marks its task as
+// complete
+func (m *mergeCoordinator) handleRow(row rowenc.EncDatumRow) error {
+	input, err := parseCoordinatorInput(row)
+	if err != nil {
+		return err
+	}
+
+	m.results.SSTs = append(m.results.SSTs, input.outputSSTs...)
+
+	next := m.tasks.ClaimNext(input.taskID)
+	if next.IsDone() {
+		m.closeLoopback()
+		return nil
+	}
+
+	m.loopback <- rowenc.EncDatumRow{
+		rowenc.EncDatum{Datum: tree.NewDBytes(tree.DBytes(input.sqlInstanceID))},
+		rowenc.EncDatum{Datum: tree.NewDInt(tree.DInt(next))},
+	}
+
+	return nil
+}
+
 // Start implements execinfra.RowSource.
 func (m *mergeCoordinator) Start(ctx context.Context) {
 	m.StartInternal(ctx, "mergeCoordinator")
 	m.input.Start(ctx)
+	m.publishInitialTasks()
 }
 
 func init() {
@@ -70,8 +190,13 @@ func init() {
 		postSpec *execinfrapb.PostProcessSpec,
 		input execinfra.RowSource,
 	) (execinfra.Processor, error) {
+		channel, cleanup := loopback.create(flow)
 		mc := &mergeCoordinator{
-			input: input,
+			input:    input,
+			tasks:    taskset.MakeTaskSet(spec.TaskCount, int64(len(spec.WorkerSqlInstanceIds))),
+			loopback: channel,
+			cleanup:  cleanup,
+			spec:     spec,
 		}
 		err := mc.Init(
 			ctx, mc, postSpec, mergeCoordinatorOutputTypes, flow, flowID, nil,

pkg/sql/bulkmerge/merge_loopback.go

@@ -10,13 +10,48 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
-	"github.com/cockroachdb/cockroach/pkg/sql/physicalplan"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowexec"
-	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/errors"
 )
 
+// loopbackMap allows the mergeLoopback processor to communicate with the merge
+// coordinator by mapping flow IDs to channels.
+type loopbackMap struct {
+	syncutil.Mutex
+	loopback map[execinfrapb.FlowID]chan rowenc.EncDatumRow
+}
+
+var loopback = &loopbackMap{
+	loopback: make(map[execinfrapb.FlowID]chan rowenc.EncDatumRow),
+}
+
+// get returns the channel for the given id if it exists.
+func (l *loopbackMap) get(flowCtx *execinfra.FlowCtx) (chan rowenc.EncDatumRow, bool) {
+	l.Lock()
+	defer l.Unlock()
+	id := flowCtx.ID
+	channel, ok := l.loopback[id]
+	return channel, ok
+}
+
+// create returns a channel for the given id and a function to close it.
+func (l *loopbackMap) create(flowCtx *execinfra.FlowCtx) (chan rowenc.EncDatumRow, func()) {
+	l.Lock()
+	defer l.Unlock()
+	id := flowCtx.ID
+	ch := make(chan rowenc.EncDatumRow)
+	l.loopback[id] = ch
+	return ch, func() {
+		l.Lock()
+		defer l.Unlock()
+		delete(l.loopback, id)
+		close(ch)
+	}
+}
+
 var (
 	_ execinfra.Processor = &mergeLoopback{}
 	_ execinfra.RowSource = &mergeLoopback{}
@@ -32,29 +67,47 @@ var mergeLoopbackOutputTypes = []*types.T{
 
 type mergeLoopback struct {
 	execinfra.ProcessorBase
-	done bool
+	loopback chan rowenc.EncDatumRow
 }
 
 // Next implements execinfra.RowSource.
 func (m *mergeLoopback) Next() (rowenc.EncDatumRow, *execinfrapb.ProducerMetadata) {
-	if m.done {
-		m.MoveToDraining(nil)
-		return nil, m.DrainHelper()
+	// Read from the loopback channel until it's closed
+	if m.State == execinfra.StateRunning {
+		row, ok := <-m.loopback
+		if !ok {
+			m.MoveToDraining(nil)
+			return nil, m.DrainHelper()
+		}
+		return row, nil
 	}
-	m.done = true
-	// Generate a routing key for the current SQL instance (where this processor is running).
-	// This ensures the routing key matches one of the spans in the range router.
-	routingDatum, _ := physicalplan.RoutingDatumsForSQLInstance(m.FlowCtx.NodeID.SQLInstanceID())
-	return rowenc.EncDatumRow{
-		routingDatum,
-		rowenc.EncDatum{Datum: tree.NewDInt(1)},
-	}, nil
+	return nil, m.DrainHelper()
 }
 
 // Start implements execinfra.RowSource.
 func (m *mergeLoopback) Start(ctx context.Context) {
 	m.StartInternal(ctx, "mergeLoopback")
-	// TODO(jeffswenson): create the initial set of tasks
+	var ok bool
+	m.loopback, ok = loopback.get(m.FlowCtx)
+	if !ok {
+		m.MoveToDraining(errors.New("loopback channel not found"))
+		return
+	}
+}
+
+func (m *mergeLoopback) DrainHelper() *execinfrapb.ProducerMetadata {
+	// First drain any inputs coming back from the coordinator.
+	for again := true; again; {
+		select {
+		case _, ok := <-m.loopback:
+			if ok {
+				continue
+			}
+		default:
+		}
+		again = false
+	}
+	return m.ProcessorBase.DrainHelper()
 }
 
 func init() {

pkg/sql/bulkmerge/merge_planning.go

@@ -16,7 +16,7 @@ import (
 )
 
 func newBulkMergePlan(
-	ctx context.Context, execCtx sql.JobExecContext,
+	ctx context.Context, execCtx sql.JobExecContext, taskCount int,
 ) (*sql.PhysicalPlan, *sql.PlanningCtx, error) {
 	// NOTE: This implementation is inspired by the physical plan created by
 	// restore in `pkg/backup/restore_processor_planning.go`
@@ -32,6 +32,11 @@ func newBulkMergePlan(
 	// Use the gateway node as the coordinator, which is where the job was initiated.
 	coordinatorID := plan.GatewaySQLInstanceID
 
+	keys := make([][]byte, 0, len(sqlInstanceIDs))
+	for _, id := range sqlInstanceIDs {
+		keys = append(keys, physicalplan.RoutingKeyForSQLInstance(id))
+	}
+
 	router, err := physicalplan.MakeInstanceRouter(sqlInstanceIDs)
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "unable to make instance router")
@@ -85,7 +90,8 @@ func newBulkMergePlan(
 
 	plan.AddSingleGroupStage(ctx, coordinatorID, execinfrapb.ProcessorCoreUnion{
 		MergeCoordinator: &execinfrapb.MergeCoordinatorSpec{
-			// TODO fill in the rest of the spec
+			TaskCount:            int64(taskCount),
+			WorkerSqlInstanceIds: keys,
 		},
 	}, execinfrapb.PostProcessSpec{}, mergeCoordinatorOutputTypes, nil /* finalizeLastStageCb */)