fix: apply migration compaction and deep job pickups fixes

integration_test/partitionmigration/partitionmigration_embedded_test.go

@@ -160,9 +160,10 @@ func TestPartitionMigrationEmbeddedMode(t *testing.T) {
 		"Processor.maxLoopSleep":                                   "1s",
 		"Router.eventOrderKeyThreshold":                            "0", // we need strict event ordering guarantees for this test
 		"Router.noOfWorkers":                                       strconv.Itoa(numPartitions),
-		"Router.Network.IncludeInstanceIdInHeader":                 "true",  // for debugging in case of receiving out-of-order events
+		"Router.Network.IncludeInstanceIdInHeader":                 "true", // for debugging in case of receiving out-of-order events
+		"Router.jobIterator.maxQueries":                            "1",
 		"Gateway.allowPartialWriteWithErrors":                      "false", // not going through the lecacy gateway path
-		"PartitionMigration.Processor.SourceNode.readExcludeSleep": "5s",    // sleep a bit less than the default one to speed up the test
+		"PartitionMigration.Processor.SourceNode.readExcludeSleep": "15s",   // sleep a bit less than the default one to speed up the test
 		"PartitionMigration.SourceNode.inProgressPollSleep":        "1s",    // poll faster for test speed
 
 		// we want to create multiple datasets during the test and ensure that migration works correctly with ds limits as well

integration_test/partitionmigration/partitionmigration_gwproc_test.go

@@ -45,6 +45,23 @@ import (
 // 12. Waits for all requests to complete.
 // 13. Verifies that all requests were received successfully and in order.
 func TestPartitionMigrationGatewayProcessorMode(t *testing.T) {
+	for _, tc := range []struct {
+		name                  string
+		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
+	}{
+		{name: "normal", extraStressWorkspaces: 0, restartProcessorEvery: 30 * time.Second},
+		{name: "stress_100_workspaces", extraStressWorkspaces: 100, restartProcessorEvery: 30 * time.Second},
+		{name: "stress_1000_workspaces", extraStressWorkspaces: 1000, restartProcessorEvery: 40 * time.Second},
+		{name: "stress_5000_workspaces", extraStressWorkspaces: 5000, restartProcessorEvery: 50 * time.Second},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			testPartitionMigrationGatewayProcessorMode(t, tc.extraStressWorkspaces, tc.restartProcessorEvery)
+		})
+	}
+}
+
+func testPartitionMigrationGatewayProcessorMode(t *testing.T, extraStressWorkspaces int, restartProcessorEvery time.Duration) {
 	const (
 		namespace     = "namespace123"
 		workspaceID   = "workspace123"
@@ -54,9 +71,10 @@ func TestPartitionMigrationGatewayProcessorMode(t *testing.T) {
 
 		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
-		restartProcessorEvery     = 10 * time.Second // how often to restart processor nodes while migration is ongoing
-		readExcludeSleep          = 5 * time.Second  // sleep duration for read exclusion during migration
+		readExcludeSleep          = 15 * time.Second // sleep duration for read exclusion during migration, must not be greater than restartProcessorEvery-5s
 	)
+	require.LessOrEqual(t, readExcludeSleep, restartProcessorEvery-5*time.Second,
+		"readExcludeSleep must not be greater than restartProcessorEvery-5s")
 
 	// distribute partitions across the 2 nodes equally
 	initialMappings := map[int]int{}
@@ -206,7 +224,7 @@ func TestPartitionMigrationGatewayProcessorMode(t *testing.T) {
 		"Router.eventOrderKeyThreshold":            "0", // we need strict event ordering guarantees for this test
 		"Router.noOfWorkers":                       strconv.Itoa(numPartitions),
 		"Router.Network.IncludeInstanceIdInHeader": "true", // for debugging in case of receiving out-of-order events
-
+		"Router.jobIterator.maxQueries":            "1",
 	}
 	rsBinaryPath := filepath.Join(t.TempDir(), "rudder-server-binary")
 	rudderserver.BuildRudderServerBinary(t, "../../main.go", rsBinaryPath)

jobsdb/integration_test.go

@@ -780,11 +780,12 @@ func TestJobsDB(t *testing.T) {
 			require.NoError(t, err)
 		}
 
-		trigger() // jobs_3, jobs_4 & jobs_5 will be migrated to jobs_5_1
+		trigger() // jobs_3 & jobs_4 will be migrated to jobs_4_1; jobs_5 stays because adding it would exceed maxDSSize
 		dsList = getDSList()
-		require.Lenf(t, dsList, 2, "dsList length is not 1, got %+v", dsList)
-		require.Equal(t, prefix+"_jobs_5_1", dsList[0].JobTable) // 12 jobs
-		require.Equal(t, prefix+"_jobs_6", dsList[1].JobTable)   // 0 jobs
+		require.Lenf(t, dsList, 3, "dsList length is not 2, got %+v", dsList)
+		require.Equal(t, prefix+"_jobs_4_1", dsList[0].JobTable) // 8 jobs
+		require.Equal(t, prefix+"_jobs_5", dsList[1].JobTable)   // 4 jobs
+		require.Equal(t, prefix+"_jobs_6", dsList[2].JobTable)   // 0 jobs
 
 		jobsResult, err = jobDB.GetUnprocessed(context.Background(), GetQueryParams{
 			CustomValFilters: []string{customVal},

jobsdb/jobsdb.go

@@ -1071,7 +1071,7 @@ func (jd *Handle) loadConfig() {
 	// jobMinRowsLeftMigrateThreshold: A DS with a low number of pending rows should be eligible for migration if the number of pending rows are
 	// less than jobMinRowsLeftMigrateThreshold percent of maxDSSize (e.g. if jobMinRowsLeftMigrateThreshold is 0.5
 	// then DSs that have less than 50% of maxDSSize pending rows are eligible for migration)
-	jd.conf.migration.jobMinRowsLeftMigrateThreshold = jd.config.GetReloadableFloat64Var(0.4, jd.configKeys("jobMinRowsLeftMigrateThreshold")...)
+	jd.conf.migration.jobMinRowsLeftMigrateThreshold = jd.config.GetReloadableFloat64Var(0.6, jd.configKeys("jobMinRowsLeftMigrateThreshold")...)
 	// maxMigrateOnce: Maximum number of DSs that are migrated together into one destination
 	jd.conf.migration.maxMigrateOnce = jd.config.GetReloadableIntVar(10, 1, jd.configKeys("maxMigrateOnce")...)
 	// maxMigrateDSProbe: Maximum number of DSs that are checked from left to right if they are eligible for migration
@@ -1277,12 +1277,12 @@ func (jd *Handle) doRefreshDSList(l lock.LockToken) (dataSetTList, error) {
 	// report table count metrics before shrinking the datasetList
 	jd.statTableCount.Gauge(len(jd.datasetList))
 
-	// if the owner of this jobsdb is a writer, then shrinking datasetList to have only last two datasets
-	// this shrank datasetList is used to compute DSRangeList
-	// This is done because, writers don't care about the left datasets in the sorted datasetList
+	// If the owner of this jobsdb is a writer, then shrinking datasetList to have only last dataset
+	// which is being written to.
+	// Writers only write to the last dataset and if this dataset is full, then create a new dataset.
 	if jd.ownerType == Write {
-		if len(jd.datasetList) > 2 {
-			jd.datasetList = jd.datasetList[len(jd.datasetList)-2 : len(jd.datasetList)]
+		if len(jd.datasetList) > 1 {
+			jd.datasetList = jd.datasetList[len(jd.datasetList)-1 : len(jd.datasetList)]
 		}
 	}
 

jobsdb/migration.go

@@ -421,14 +421,8 @@ func (jd *Handle) getMigrationList(dsList []dataSetT) (migrateFrom []dsWithPendi
 	}
 
 	for idx, ds := range dsList {
-		var idxCheck bool
-		if jd.ownerType == Read {
-			// if jobsdb owner is read, exempting the last two datasets from migration.
-			// This is done to avoid dsList conflicts between reader and writer
-			idxCheck = idx == len(dsList)-1 || idx == len(dsList)-2
-		} else {
-			idxCheck = idx == len(dsList)-1
-		}
+		// exempting the last dataset from migration since it is the one being currently written to.
+		idxCheck := idx == len(dsList)-1
 
 		if liveDSCount >= jd.conf.migration.maxMigrateOnce.Load() || pendingJobsCount >= maxDSSize || idxCheck {
 			break
@@ -455,12 +449,19 @@ func (jd *Handle) getMigrationList(dsList []dataSetT) (migrateFrom []dsWithPendi
 				liveDSCount++
 			} else {
 				if waiting != nil { // have another dataset waiting for a pair
-					migrateFrom = append(migrateFrom, *waiting, dsWithPendingJobCount{ds: ds, numJobsPending: recordsLeft})
-					insertBeforeDS = dsList[idx+1]
-					pendingJobsCount += waiting.numJobsPending + recordsLeft
-					liveDSCount += 2
-					waiting = nil
+					if waiting.numJobsPending+recordsLeft > maxDSSize {
+						waiting = nil
+					} else {
+						migrateFrom = append(migrateFrom, *waiting, dsWithPendingJobCount{ds: ds, numJobsPending: recordsLeft})
+						insertBeforeDS = dsList[idx+1]
+						pendingJobsCount += waiting.numJobsPending + recordsLeft
+						liveDSCount += 2
+						waiting = nil
+					}
 				} else if pendingJobsCount > 0 { // we already know that we'll be migrating another dataset with pending jobs, so can add this one too
+					if pendingJobsCount+recordsLeft > maxDSSize {
+						break // adding this dataset would exceed maxDSSize, leave it for the next migration cycle
+					}
 					migrateFrom = append(migrateFrom, dsWithPendingJobCount{ds: ds, numJobsPending: recordsLeft})
 					insertBeforeDS = dsList[idx+1]
 					pendingJobsCount += recordsLeft

jobsdb/migration_test.go

@@ -651,3 +651,212 @@
 		require.Error(t, err)
 	})
 }
+
+func TestMigrationMaxDSSizeGuard(t *testing.T) {
+	_ = startPostgres(t)
+
+	// newJobDB creates a Handle with the given maxDSSize and jobMinRowsLeftMigrateThreshold.
+	newJobDB := func(t *testing.T, maxDSSize int, threshold float64) (*Handle, chan time.Time, *config.Config) {
+		t.Helper()
+		c := config.New()
+		c.Set("JobsDB.maxDSSize", maxDSSize)
+		c.Set("JobsDB.jobMinRowsLeftMigrateThreshold", threshold)
+		triggerAddNewDS := make(chan time.Time)
+		jd := &Handle{
+			TriggerAddNewDS:  func() <-chan time.Time { return triggerAddNewDS },
+			TriggerMigrateDS: func() <-chan time.Time { return make(chan time.Time) },
+			config:           c,
+		}
+		require.NoError(t, jd.Setup(ReadWrite, true, strings.ToLower(rand.String(5))))
+		t.Cleanup(jd.TearDown)
+		return jd, triggerAddNewDS, c
+	}
+
+	// addDS stores `jobs` with len(jobs)-pending marked as succeeded, then triggers addNewDS.
+	// Jobs must be a slice of a larger pre-created slice so that their pre-set IDs match the
+	// DB-assigned IDs (which start at 1 and increment globally per table).
+	addDS := func(t *testing.T, jd *Handle, trigger chan time.Time, jobs []*JobT, pending int) {
+		t.Helper()
+		require.NoError(t, jd.Store(context.Background(), jobs))
+		if terminal := len(jobs) - pending; terminal > 0 {
+			require.NoError(t, jd.UpdateJobStatus(context.Background(), genJobStatuses(jobs[:terminal], "executing")))
+			require.NoError(t, jd.UpdateJobStatus(context.Background(), genJobStatuses(jobs[:terminal], "succeeded")))
+		}
+		trigger <- time.Now()
+		trigger <- time.Now()
+	}
+
+	t.Run("accumulation stops before exceeding maxDSSize", func(t *testing.T) {
+		// maxDSSize=10, threshold=0.7 → pair threshold=7
+		// DS1..DS4: 3 pending each (needsPair since 3 < 7)
+		// DS5: last DS (exempt)
+		//
+		// getMigrationList walk:
+		//   DS1 → waiting
+		//   DS2 → pair: 3+3=6 ≤ 10, pendingJobsCount=6
+		//   DS3 → piggyback: 6+3=9 ≤ 10, pendingJobsCount=9
+		//   DS4 → piggyback: 9+3=12 > 10 → break
+		// expected: migrateFrom=[DS1,DS2,DS3], pendingJobsCount=9 ≤ maxDSSize
+		jd, trigger, c := newJobDB(t, 10, 0.7)
+		allJobs := genJobs(defaultWorkspaceID, "test", 50, 1) // 4 regular DSes + 1 last, 10 jobs each
+		for i := range 4 {
+			addDS(t, jd, trigger, allJobs[i*10:(i+1)*10], 3)
+		}
+		require.NoError(t, jd.Store(context.Background(), allJobs[40:]))
+
+		dsList := jd.getDSList()
+		c.Set("JobsDB."+jd.tablePrefix+"."+"maxMigrateDSProbe", len(dsList))
+
+		migrateFrom, pendingJobsCount, _, err := jd.getMigrationList(dsList)
+		require.NoError(t, err)
+		require.Len(t, migrateFrom, 3)
+		require.Equal(t, 9, pendingJobsCount)
+		require.LessOrEqual(t, pendingJobsCount, 10)
+	})
+
+	t.Run("pair exceeding maxDSSize is discarded", func(t *testing.T) {
+		// maxDSSize=10, threshold=0.7 → pair threshold=7
+		// DS1: 6 pending (needsPair since 6 < 7)
+		// DS2: 6 pending → 6+6=12 > maxDSSize → waiting cleared, nothing migrates
+		// DS3: last DS (exempt)
+		jd, trigger, c := newJobDB(t, 10, 0.7)
+		allJobs := genJobs(defaultWorkspaceID, "test", 30, 1) // 2 regular DSes + 1 last, 10 jobs each
+		addDS(t, jd, trigger, allJobs[:10], 6)
+		addDS(t, jd, trigger, allJobs[10:20], 6)
+		require.NoError(t, jd.Store(context.Background(), allJobs[20:]))
+
+		dsList := jd.getDSList()
+		c.Set("JobsDB."+jd.tablePrefix+"."+"maxMigrateDSProbe", len(dsList))
+
+		migrateFrom, _, _, err := jd.getMigrationList(dsList)
+		require.NoError(t, err)
+		require.Empty(t, migrateFrom)
+	})
+}
+
+func TestMigrationSkipsDatasets(t *testing.T) {
+	config.Reset()
+	c := config.New()
+	c.Set("JobsDB.maxDSSize", 10)
+
+	_ = startPostgres(t)
+
+	triggerAddNewDS := make(chan time.Time)
+	triggerMigrateDS := make(chan time.Time)
+
+	jobDB := Handle{
+		TriggerAddNewDS:  func() <-chan time.Time { return triggerAddNewDS },
+		TriggerMigrateDS: func() <-chan time.Time { return triggerMigrateDS },
+		config:           c,
+	}
+	tablePrefix := strings.ToLower(rand.String(5))
+	require.NoError(t, jobDB.Setup(ReadWrite, true, tablePrefix))
+	defer jobDB.TearDown()
+
+	const totalDS = 200
+	const eligibleDSPos = 150 // 0-indexed position of the dataset we'll make eligible
+
+	// Create totalDS datasets: store jobs then trigger addNewDS to create the next one
+	for i := range totalDS - 1 { // -1 because Setup already creates DS 1
+		_ = i
+		require.NoError(t, jobDB.Store(context.Background(), genJobs(defaultWorkspaceID, "test", 10, 1)))
+		triggerAddNewDS <- time.Now()
+		triggerAddNewDS <- time.Now()
+	}
+	// Store jobs in the last DS too
+	require.NoError(t, jobDB.Store(context.Background(), genJobs(defaultWorkspaceID, "test", 10, 1)))
+
+	dsList := jobDB.getDSList()
+	require.Len(t, dsList, totalDS)
+
+	// Make all jobs in the dataset at eligibleDSPos terminal (succeeded)
+	eligibleDS := dsList[eligibleDSPos]
+	rows, err := jobDB.dbHandle.Query(fmt.Sprintf(`SELECT job_id FROM %q`, eligibleDS.JobTable))
+	require.NoError(t, err)
+	var statusJobs []*JobT
+	for rows.Next() {
+		var id int64
+		require.NoError(t, rows.Scan(&id))
+		statusJobs = append(statusJobs, &JobT{JobID: id})
+	}
+	require.NoError(t, rows.Err())
+	_ = rows.Close()
+
+	require.NoError(t, jobDB.UpdateJobStatus(context.Background(), genJobStatuses(statusJobs, "succeeded")))
+
+	t.Run("intra-invocation skip", func(t *testing.T) {
+		// Allow probing enough datasets to reach the eligible one in a single call
+		c.Set("JobsDB."+tablePrefix+"."+"maxMigrateDSProbe", totalDS)
+
+		// Measure first getMigrationList call (full scan, no skip)
+		checkStart := time.Now()
+		checkResult, err := jobDB.getMigrationList(dsList, nil)
+		checkDuration := time.Since(checkStart)
+		require.NoError(t, err)
+		require.NotEmpty(t, checkResult.migrateFrom, "should find eligible datasets")
+		require.NotNil(t, checkResult.firstEligible, "should have firstEligible set")
+
+		// Measure second getMigrationList call (with skipBefore from first call)
+		lockCheckStart := time.Now()
+		lockResult, err := jobDB.getMigrationList(dsList, checkResult.firstEligible)
+		lockCheckDuration := time.Since(lockCheckStart)
+		require.NoError(t, err)
+		require.NotEmpty(t, lockResult.migrateFrom)
+
+		// Both calls should find the same eligible datasets
+		require.Equal(t, len(checkResult.migrateFrom), len(lockResult.migrateFrom))
+		for i := range checkResult.migrateFrom {
+			require.Equal(t, checkResult.migrateFrom[i].ds.Index, lockResult.migrateFrom[i].ds.Index)
+		}
+
+		t.Logf("check duration (no skip):   %v", checkDuration)
+		t.Logf("lock check duration (skip):  %v", lockCheckDuration)
+
+		// The second call with skipBefore should be significantly faster
+		require.Greater(t, checkDuration, lockCheckDuration,
+			"getMigrationList with skipBefore should be faster than without",
+		)
+	})
+
+	t.Run("cross-invocation resume", func(t *testing.T) {
+		// Set maxMigrateDSProbe to 100 so the first iteration can't reach the
+		// eligible dataset at position 150. It will need 2 iterations.
+		c.Set("JobsDB."+tablePrefix+"."+"maxMigrateDSProbe", 100)
+		jobDB.lastMigrateProbeIndex = nil
+
+		// 1st iteration: probes datasets 1..100, finds nothing, hits probe limit.
+		// Should store lastMigrateProbeIndex for resumption.
+		result1, err := jobDB.getMigrationList(dsList, jobDB.lastMigrateProbeIndex)
+		require.NoError(t, err)
+		require.Empty(t, result1.migrateFrom, "should not find eligible datasets in first 100")
+		require.True(t, result1.probeLimitReached, "should hit probe limit")
+		require.NotNil(t, result1.lastProbed, "should have lastProbed set")
+
+		// Simulate what doMigrateDS does: save the resume point
+		jobDB.lastMigrateProbeIndex = result1.lastProbed
+
+		// 2nd iteration: resumes from where the first left off, finds the eligible dataset.
+		resumeStart := time.Now()
+		result2, err := jobDB.getMigrationList(dsList, jobDB.lastMigrateProbeIndex)
+		resumeDuration := time.Since(resumeStart)
+		require.NoError(t, err)
+		require.NotEmpty(t, result2.migrateFrom, "should find eligible datasets in second iteration")
+		require.Equal(t, dsList[eligibleDSPos].Index, result2.migrateFrom[0].ds.Index)
+
+		// Compare with a full scan from scratch
+		c.Set("JobsDB."+tablePrefix+"."+"maxMigrateDSProbe", totalDS)
+		fullStart := time.Now()
+		resultFull, err := jobDB.getMigrationList(dsList, nil)
+		fullDuration := time.Since(fullStart)
+		require.NoError(t, err)
+		require.NotEmpty(t, resultFull.migrateFrom)
+
+		t.Logf("full scan duration:    %v", fullDuration)
+		t.Logf("resumed scan duration: %v", resumeDuration)
+
+		// The resumed scan should be faster than a full scan
+		require.Greater(t, fullDuration, resumeDuration,
+			"resumed scan should be faster than full scan from scratch",
+		)
+	})
+}