feat(processor): fork event processing for configured isolated destinations

Author: atzoum

app/apphandlers/setup_partitionmigration.go

@@ -208,6 +208,8 @@ func setupProcessorPartitionMigrator(ctx context.Context,
 	if procRWDB != nil {
 		procBuffRWHandle := jobsdb.NewForReadWrite(
 			"proc_buf",
+			// proc_buf doesn't need to be multi-consumer because it is used only as a temporary storage.
+			// During flush, the jobs will be moved to the proc jobsdb which is multi-consumer.
 			jobsdb.WithClearDB(false),
 			jobsdb.WithDSLimit(config.GetReloadableIntVar(0, 1, "JobsDB.proc_buf.dsLimit", "JobsDB.dsLimit")),
 			jobsdb.WithSkipMaintenanceErr(config.GetBoolVar(true, "JobsDB.proc_buf.skipMaintenanceError", "JobsDB.buff.skipMaintenanceError", "JobsDB.skipMaintenanceError")),

integration_test/partitionmigration/partitionmigration_embedded_test.go

@@ -51,23 +51,25 @@ func TestPartitionMigrationEmbeddedMode(t *testing.T) {
 	for _, tc := range []struct {
 		name                string
 		jobsDBFanoutEnabled bool // whether source nodes declare per-jobsdb fan-out on migration acknowledgement
+		forkDestinations    bool // fork the destination to the (multi-consumer) proc jobsdb, exercising proc-jobsdb migration
 	}{
-		{name: "normal", jobsDBFanoutEnabled: true},
-		{name: "legacy_no_jobsdb_fanout", jobsDBFanoutEnabled: false},
+		{name: "normal", jobsDBFanoutEnabled: true, forkDestinations: true},
+		{name: "legacy_no_jobsdb_fanout", jobsDBFanoutEnabled: false, forkDestinations: false},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			testPartitionMigrationEmbeddedMode(t, tc.jobsDBFanoutEnabled)
+			testPartitionMigrationEmbeddedMode(t, tc.jobsDBFanoutEnabled, tc.forkDestinations)
 		})
 	}
 }
 
-func testPartitionMigrationEmbeddedMode(t *testing.T, jobsDBFanoutEnabled bool) {
+func testPartitionMigrationEmbeddedMode(t *testing.T, jobsDBFanoutEnabled, forkDestinations bool) {
 	const (
-		namespace     = "namespace123"
-		workspaceID   = "workspace123"
-		sourceID      = "source123"
-		destinationID = "destination123"
-		writeKey      = "writekey123"
+		namespace          = "namespace123"
+		workspaceID        = "workspace123"
+		sourceID           = "source123"
+		destinationID      = "destination123"
+		abortDestinationID = "destination-abort-123" // second destination, aborted at the router
+		writeKey           = "writekey123"
 
 		numPartitions             = 4 // needs to be a power of 2 (e.g., 2, 4, 8, 16, ...)
 		jobsPerPartitionPerSecond = 50
@@ -116,22 +118,31 @@ func testPartitionMigrationEmbeddedMode(t *testing.T, jobsDBFanoutEnabled bool)
 	wh := newTestWebhook(t)
 	t.Cleanup(wh.Close)
 
-	// start a test backendconfig with 1 source connected to the webhook
+	// start a test backendconfig with 1 source connected to the webhook. When forking, add a
+	// second destination on the same source that the router aborts (never delivered), so each
+	// forked proc job carries two consumers and migration must move all pending consumers.
+	sourceBuilder := backendconfigtest.NewSourceBuilder().
+		WithWorkspaceID(workspaceID).
+		WithID(sourceID).
+		WithWriteKey(writeKey).
+		WithConnection(
+			backendconfigtest.NewDestinationBuilder("WEBHOOK").
+				WithID(destinationID).
+				WithConfigOption("webhookMethod", "POST").
+				WithConfigOption("webhookUrl", wh.URL).
+				Build())
+	if forkDestinations {
+		sourceBuilder = sourceBuilder.WithConnection(
+			backendconfigtest.NewDestinationBuilder("WEBHOOK").
+				WithID(abortDestinationID).
+				WithConfigOption("webhookMethod", "POST").
+				WithConfigOption("webhookUrl", "http://localhost:1234"). // aborted at the router, never delivered
+				Build())
+	}
 	bc := backendconfigtest.NewBuilder().
 		WithNamespace(namespace, backendconfigtest.NewConfigBuilder().
 			WithWorkspaceID(workspaceID).
-			WithSource(
-				backendconfigtest.NewSourceBuilder().
-					WithWorkspaceID(workspaceID).
-					WithID(sourceID).
-					WithWriteKey(writeKey).
-					WithConnection(
-						backendconfigtest.NewDestinationBuilder("WEBHOOK").
-							WithID(destinationID).
-							WithConfigOption("webhookMethod", "POST").
-							WithConfigOption("webhookUrl", wh.URL).
-							Build()).
-					Build()).
+			WithSource(sourceBuilder.Build()).
 			Build()).
 		Build()
 
@@ -218,6 +229,12 @@ func testPartitionMigrationEmbeddedMode(t *testing.T, jobsDBFanoutEnabled bool)
 		"JobsDB.dsLimit":                        "2",
 		"JobsDB.refreshDSListLoopSleepDuration": "5s",
 	}
+	if forkDestinations {
+		// fork both destinations to the proc jobsdb (isolation already enabled above) so each
+		// proc job carries two consumers; abort the second at the router so it is never delivered
+		commonEnv["Processor.DestinationIsolation.enabledDestinations.all"] = "true"
+		commonEnv["Router.toAbortDestinationIDs"] = abortDestinationID
+	}
 	rsBinaryPath := filepath.Join(t.TempDir(), "rudder-server-binary")
 	rudderserver.BuildRudderServerBinary(t, "../../main.go", rsBinaryPath)
 	node0Name := "proc-node-0"
@@ -282,6 +299,14 @@ func testPartitionMigrationEmbeddedMode(t *testing.T, jobsDBFanoutEnabled bool)
 	// wait for some time to let events flow
 	time.Sleep(10 * time.Second)
 
+	if forkDestinations {
+		// the fan-out siphon must be active: forked jobs land in the proc jobsdb of some node,
+		// giving partition migration multi-consumer proc jobs to move
+		require.Eventually(t, func() bool {
+			return procJobsCount(t, pg0.DB)+procJobsCount(t, pg1.DB) > 0
+		}, 30*time.Second, 500*time.Millisecond, "forked jobs should appear in the proc jobsdb")
+	}
+
 	var srcRouterAcks []string
 	err = rudoacker.NewSrcrouterAcker(ctx, g, etcdResource.Client, namespace, []string{"srcrouter"}).
 		WithEventListener(func(key string, value etcdtypes.ReloadSrcRouterCommand) {

integration_test/partitionmigration/partitionmigration_gwproc_test.go

@@ -2,6 +2,7 @@ package partitionmigration_test
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"path/filepath"
 	"strconv"
@@ -55,26 +56,26 @@ func TestPartitionMigrationGatewayProcessorMode(t *testing.T) {
 		extraStressWorkspaces int           // number of extra workspace migrations to include (0 = normal mode)
 		restartProcessorEvery time.Duration // how often to restart processor nodes while migration is ongoing
 		jobsDBFanoutEnabled   bool          // whether source nodes declare per-jobsdb fan-out on migration acknowledgement
+		forkDestinations      bool          // fork the destination to the (multi-consumer) proc jobsdb, exercising proc-jobsdb migration
 	}{
-		{name: "normal", extraStressWorkspaces: 0, restartProcessorEvery: 25 * time.Second, jobsDBFanoutEnabled: true},
-		{name: "stress_100_workspaces", extraStressWorkspaces: 100, restartProcessorEvery: 30 * time.Second, jobsDBFanoutEnabled: true},
-		{name: "stress_1000_workspaces", extraStressWorkspaces: 1000, restartProcessorEvery: 35 * time.Second, jobsDBFanoutEnabled: true},
-		{name: "stress_5000_workspaces", extraStressWorkspaces: 5000, restartProcessorEvery: 50 * time.Second, jobsDBFanoutEnabled: true},
-		{name: "legacy_no_jobsdb_fanout", extraStressWorkspaces: 0, restartProcessorEvery: 25 * time.Second, jobsDBFanoutEnabled: false},
+		{name: "normal", extraStressWorkspaces: 0, restartProcessorEvery: 25 * time.Second, jobsDBFanoutEnabled: true, forkDestinations: true},
+		{name: "legacy_no_jobsdb_fanout", extraStressWorkspaces: 0, restartProcessorEvery: 25 * time.Second, jobsDBFanoutEnabled: false, forkDestinations: false},
+		{name: "stress_5000_workspaces", extraStressWorkspaces: 5000, restartProcessorEvery: 50 * time.Second, jobsDBFanoutEnabled: true, forkDestinations: false},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			testPartitionMigrationGatewayProcessorMode(t, tc.extraStressWorkspaces, tc.restartProcessorEvery, tc.jobsDBFanoutEnabled)
+			testPartitionMigrationGatewayProcessorMode(t, tc.extraStressWorkspaces, tc.restartProcessorEvery, tc.jobsDBFanoutEnabled, tc.forkDestinations)
 		})
 	}
 }
 
-func testPartitionMigrationGatewayProcessorMode(t *testing.T, extraStressWorkspaces int, restartProcessorEvery time.Duration, jobsDBFanoutEnabled bool) {
+func testPartitionMigrationGatewayProcessorMode(t *testing.T, extraStressWorkspaces int, restartProcessorEvery time.Duration, jobsDBFanoutEnabled, forkDestinations bool) {
 	const (
-		namespace     = "namespace123"
-		workspaceID   = "workspace123"
-		sourceID      = "source123"
-		destinationID = "destination123"
-		writeKey      = "writekey123"
+		namespace          = "namespace123"
+		workspaceID        = "workspace123"
+		sourceID           = "source123"
+		destinationID      = "destination123"
+		abortDestinationID = "destination-abort-123" // second destination, aborted at the router
+		writeKey           = "writekey123"
 
 		numPartitions             = 4                // needs to be a power of 2 (e.g., 2, 4, 8, 16, ...)
 		jobsPerPartitionPerSecond = 50               // number of jobs to send per partition per second from the gateway client
@@ -137,22 +138,31 @@ func testPartitionMigrationGatewayProcessorMode(t *testing.T, extraStressWorkspa
 	wh := newTestWebhook(t)
 	t.Cleanup(wh.Close)
 
-	// start a test backendconfig with 1 source connected to the webhook
+	// start a test backendconfig with 1 source connected to the webhook. When forking, add a
+	// second destination on the same source that the router aborts (never delivered), so each
+	// forked proc job carries two consumers and migration must move all pending consumers.
+	sourceBuilder := backendconfigtest.NewSourceBuilder().
+		WithWorkspaceID(workspaceID).
+		WithID(sourceID).
+		WithWriteKey(writeKey).
+		WithConnection(
+			backendconfigtest.NewDestinationBuilder("WEBHOOK").
+				WithID(destinationID).
+				WithConfigOption("webhookMethod", "POST").
+				WithConfigOption("webhookUrl", wh.URL).
+				Build())
+	if forkDestinations {
+		sourceBuilder = sourceBuilder.WithConnection(
+			backendconfigtest.NewDestinationBuilder("WEBHOOK").
+				WithID(abortDestinationID).
+				WithConfigOption("webhookMethod", "POST").
+				WithConfigOption("webhookUrl", "http://localhost:1234"). // aborted at the router, never delivered
+				Build())
+	}
 	bc := backendconfigtest.NewBuilder().
 		WithNamespace(namespace, backendconfigtest.NewConfigBuilder().
 			WithWorkspaceID(workspaceID).
-			WithSource(
-				backendconfigtest.NewSourceBuilder().
-					WithWorkspaceID(workspaceID).
-					WithID(sourceID).
-					WithWriteKey(writeKey).
-					WithConnection(
-						backendconfigtest.NewDestinationBuilder("WEBHOOK").
-							WithID(destinationID).
-							WithConfigOption("webhookMethod", "POST").
-							WithConfigOption("webhookUrl", wh.URL).
-							Build()).
-					Build()).
+			WithSource(sourceBuilder.Build()).
 			Build()).
 		Build()
 	bc.URL = strings.Replace(bc.URL, "127.0.0.1", localIp, 1) // replace localhost with local IP for docker containers to access
@@ -271,6 +281,12 @@ func testPartitionMigrationGatewayProcessorMode(t *testing.T, extraStressWorkspa
 		"Router.Network.IncludeInstanceIdInHeader": "true", // for debugging in case of receiving out-of-order events
 		"Router.jobIterator.maxQueries":            "1",
 	}
+	if forkDestinations {
+		// fork both destinations to the proc jobsdb (isolation already enabled above) so each
+		// proc job carries two consumers; abort the second at the router so it is never delivered
+		procCommonEnv["Processor.DestinationIsolation.enabledDestinations.all"] = "true"
+		procCommonEnv["Router.toAbortDestinationIDs"] = abortDestinationID
+	}
 	rsBinaryPath := filepath.Join(t.TempDir(), "rudder-server-binary")
 	rudderserver.BuildRudderServerBinary(t, "../../main.go", rsBinaryPath)
 	gwNode0Name := "gw-node-0"
@@ -363,6 +379,14 @@ func testPartitionMigrationGatewayProcessorMode(t *testing.T, extraStressWorkspa
 	// wait for some time to let events flow
 	time.Sleep(10 * time.Second)
 
+	if forkDestinations {
+		// the fan-out siphon must be active: forked jobs land in the proc jobsdb of some node,
+		// giving partition migration multi-consumer proc jobs to move
+		require.Eventually(t, func() bool {
+			return procJobsCount(t, pg0.DB)+procJobsCount(t, pg1.DB) > 0
+		}, 30*time.Second, 500*time.Millisecond, "forked jobs should appear in the proc jobsdb")
+	}
+
 	var srcRouterAcks []string
 	err = rudoacker.NewSrcrouterAcker(ctx, g, etcdResource.Client, namespace, []string{"srcrouter"}).
 		WithEventListener(func(key string, value etcdtypes.ReloadSrcRouterCommand) {
@@ -499,3 +523,13 @@ func restartingProcessorServer(t *testing.T, ctx context.Context, g *errgroup.Gr
 		}
 	})
 }
+
+// procJobsCount returns the number of jobs currently present in the intermediate (proc)
+// jobsdb of a node. Used to assert the destination-isolation fan-out siphon is active so
+// partition migration has proc jobs to move.
+func procJobsCount(t *testing.T, db *sql.DB) int {
+	t.Helper()
+	var count int
+	require.NoError(t, db.QueryRow("SELECT count(DISTINCT job_id) FROM unionjobsdbmetadata('proc', 10)").Scan(&count))
+	return count
+}

jobsdb/jobsdb_multiconsumer_test.go

@@ -417,3 +417,42 @@ func TestMultiConsumerJobsDB_Cache(t *testing.T) {
 	require.NoError(t, jd.UpdateJobStatus(ctx, []*JobStatusT{succeedFor(res.Jobs[0], "B")}))
 	require.True(t, cacheHit("A"), "B's status update must not invalidate A's cache")
 }
+
+// TestStoreConsumersOnSingleConsumerHandle verifies that a single-consumer jobsdb preserves a
+// job's explicitly-set Consumers on store and round-trips them through GetUnprocessed. This is
+// relied upon by the partition-migration buffer (proc_buf): it is a single-consumer handle that
+// relays multi-consumer proc jobs between nodes, so it must carry their consumers through.
+// Without it, migrated jobs would resurface under the legacy ” consumer and be dropped at the
+// target (unknown destination ""). Jobs without explicit consumers keep the legacy ” consumer.
+func TestStoreConsumersOnSingleConsumerHandle(t *testing.T) {
+	postgres := startPostgres(t)
+
+	prefix := strings.ToLower(rand.String(5))
+	jd := NewForReadWrite(prefix, // NOT multi-consumer
+		WithDBHandle(postgres.DB),
+		WithConfig(config.New()),
+	)
+	require.NoError(t, jd.Start())
+	defer jd.TearDown()
+
+	ctx := context.Background()
+	const customVal = "SC"
+	newJob := func(userID string, consumers []string) *JobT {
+		return &JobT{
+			UUID: uuid.New(), UserID: userID, CustomVal: customVal,
+			Parameters: []byte(`{}`), EventPayload: []byte(`{}`), EventCount: 1,
+			WorkspaceId: "w", Consumers: consumers,
+		}
+	}
+	require.NoError(t, jd.Store(ctx, []*JobT{
+		newJob("with", []string{"A", "B"}),
+		newJob("without", nil),
+	}))
+
+	res, err := jd.GetUnprocessed(ctx, GetQueryParams{CustomValFilters: []string{customVal}, JobsLimit: 10})
+	require.NoError(t, err)
+	require.Len(t, res.Jobs, 2)
+	// GetUnprocessed returns jobs ordered by job_id, i.e. insertion order.
+	require.Equal(t, []string{"A", "B"}, res.Jobs[0].Consumers, "explicit consumers must be preserved on a single-consumer handle")
+	require.Equal(t, []string{""}, res.Jobs[1].Consumers, "a job without consumers keeps the legacy '' consumer")
+}

jobsdb/jobsdb_pending_events.go

@@ -84,6 +84,13 @@ func (pejdb *pendingEventsJobsDB) Store(ctx context.Context, jobList []*JobT) er
 }
 
 func (pejdb *pendingEventsJobsDB) StoreInTx(ctx context.Context, tx StoreSafeTx, jobList []*JobT) error {
+	// Register the pending-events increase only after the underlying store succeeds. The
+	// store may be retried internally on a stale dataset list; when it runs inside a
+	// caller-provided transaction (WithStoreSafeTxFromTx) that transaction is reused across
+	// retries, so registering before the store would double-count on every retry.
+	if err := pejdb.JobsDB.StoreInTx(ctx, tx, jobList); err != nil {
+		return err
+	}
 	tx.Tx().AddSuccessListener(func() {
 		counters := make(map[string]map[string]map[string]float64) // workspaceID -> destType -> destinationID -> count
 		for _, job := range jobList {
@@ -107,7 +114,7 @@ func (pejdb *pendingEventsJobsDB) StoreInTx(ctx context.Context, tx StoreSafeTx,
 			}
 		}
 	})
-	return pejdb.JobsDB.StoreInTx(ctx, tx, jobList)
+	return nil
 }
 
 func (pejdb *pendingEventsJobsDB) UpdateJobStatus(ctx context.Context, statusList []*JobStatusT) error {
@@ -117,6 +124,12 @@ func (pejdb *pendingEventsJobsDB) UpdateJobStatus(ctx context.Context, statusLis
 }
 
 func (pejdb *pendingEventsJobsDB) UpdateJobStatusInTx(ctx context.Context, tx UpdateSafeTx, statusList []*JobStatusT) error {
+	// Register the pending-events decrease only after the underlying update succeeds: the
+	// update retries internally on a stale dataset list against the (possibly reused)
+	// transaction, so registering before it would double-count on every retry.
+	if err := pejdb.JobsDB.UpdateJobStatusInTx(ctx, tx, statusList); err != nil {
+		return err
+	}
 	tx.Tx().AddSuccessListener(func() {
 		counters := make(map[string]map[string]map[string]float64) // workspaceID -> destType -> destinationID -> count
 		for _, status := range statusList {
@@ -142,5 +155,5 @@ func (pejdb *pendingEventsJobsDB) UpdateJobStatusInTx(ctx context.Context, tx Up
 			}
 		}
 	})
-	return pejdb.JobsDB.UpdateJobStatusInTx(ctx, tx, statusList)
+	return nil
 }

jobsdb/jobsdb_pending_events_test.go

@@ -11,6 +11,8 @@ import (
 
 	"github.com/rudderlabs/rudder-go-kit/config"
 	"github.com/rudderlabs/rudder-go-kit/testhelper/rand"
+
+	txpkg "github.com/rudderlabs/rudder-server/utils/tx"
 )
 
 func TestPendingEventsJobsDB(t *testing.T) {
@@ -331,6 +333,48 @@ func TestPendingEventsJobsDB(t *testing.T) {
 		require.Equal(t, float64(1), registry.decreasedByDest["A"])
 		require.Equal(t, float64(0), registry.decreasedByDest["B"])
 	})
+
+	t.Run("retry on a shared transaction increases pending events exactly once", func(t *testing.T) {
+		// Regression: with WithStoreSafeTxFromTx the caller's transaction is reused across
+		// the store's internal stale-dataset-list retry. The pending-events increase must be
+		// registered only after the store succeeds, so a failed-then-successful attempt on the
+		// same committed transaction counts once, not once per attempt.
+		_ = startPostgres(t)
+		jobDB := &Handle{config: config.New()}
+		require.NoError(t, jobDB.Setup(ReadWrite, false, strings.ToLower(rand.String(5))))
+		defer jobDB.TearDown()
+
+		registry := &mockPendingEventsRegistry{}
+		// underlying store fails the first attempt (as a stale-DS retry would) and succeeds the
+		// second, without touching the real DB — the real tx below is what actually commits.
+		underlying := &failThenSucceedStoreJobsDB{JobsDB: jobDB, errs: []error{ErrStaleDsList, nil}}
+		decoratedDB := NewPendingEventsJobsDB(underlying, registry).(*pendingEventsJobsDB)
+
+		jobs := genJobsWithDestination(3)
+		require.NoError(t, jobDB.WithTx(context.Background(), func(tx *txpkg.Tx) error {
+			stx := &storeSafeTx{tx: tx, identity: underlying.Identifier()}
+			require.ErrorIs(t, decoratedDB.StoreInTx(context.Background(), stx, jobs), ErrStaleDsList)
+			require.NoError(t, decoratedDB.StoreInTx(context.Background(), stx, jobs))
+			return nil // commit → fire success listeners
+		}))
+
+		require.Equal(t, 1, registry.increaseCallCount, "increase must fire once despite the retry")
+		require.Equal(t, float64(3), registry.totalIncreased)
+	})
+}
+
+// failThenSucceedStoreJobsDB is a JobsDB whose StoreInTx returns the queued errors in order
+// (modelling a stale-dataset-list retry) without performing any real store. All other
+// behaviour is inherited from the embedded handle.
+type failThenSucceedStoreJobsDB struct {
+	JobsDB
+	errs []error
+}
+
+func (f *failThenSucceedStoreJobsDB) StoreInTx(_ context.Context, _ StoreSafeTx, _ []*JobT) error {
+	err := f.errs[0]
+	f.errs = f.errs[1:]
+	return err
 }
 
 // mockPendingEventsRegistry tracks calls to the pending events registry