chore: improve batch router

.github/workflows/verify.yml

@@ -77,5 +77,5 @@ jobs:
       - name: golangci-lint
         uses: golangci/golangci-lint-action@v7
         with:
-          version: v2.0.2
+          version: v2.1.1
           args: -v

Makefile

@@ -6,7 +6,7 @@ TESTFILE=_testok
 MOUNT_PATH=/local
 
 # go tools versions
-GOLANGCI=github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.0.2
+GOLANGCI=github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.1.1
 gofumpt=mvdan.cc/gofumpt@latest
 govulncheck=golang.org/x/vuln/cmd/govulncheck@latest
 goimports=golang.org/x/tools/cmd/goimports@latest

integration_test/docker_test/docker_test.go

@@ -175,12 +175,12 @@ func testCases(t *testing.T) {
 			eventSql := "select anonymous_id, user_id from dev_integration_test_1.identifies limit 1;"
 			_ = db.QueryRow(eventSql).Scan(&myEvent.anonymousID, &myEvent.userID)
 			return myEvent.anonymousID == "anonymousId_1"
-		}, time.Minute, 10*time.Millisecond)
+		}, 2*time.Minute, 10*time.Millisecond)
 		require.Eventually(t, func() bool {
 			eventSql := "select count(*) from dev_integration_test_1.identifies;"
 			_ = db.QueryRow(eventSql).Scan(&myEvent.count)
 			return myEvent.count == "2"
-		}, time.Minute, 10*time.Millisecond)
+		}, 2*time.Minute, 10*time.Millisecond)
 
 		// Verify User Transformation
 		eventSql := "select context_myuniqueid,context_id,context_ip from dev_integration_test_1.identifies;"
@@ -194,12 +194,12 @@ func testCases(t *testing.T) {
 			eventSql := "select anonymous_id, user_id from dev_integration_test_1.users limit 1;"
 			_ = db.QueryRow(eventSql).Scan(&myEvent.anonymousID, &myEvent.userID)
 			return myEvent.anonymousID == "anonymousId_1"
-		}, time.Minute, 10*time.Millisecond)
+		}, 2*time.Minute, 10*time.Millisecond)
 		require.Eventually(t, func() bool {
 			eventSql := "select count(*) from dev_integration_test_1.users;"
 			_ = db.QueryRow(eventSql).Scan(&myEvent.count)
 			return myEvent.count == "1"
-		}, time.Minute, 10*time.Millisecond)
+		}, 2*time.Minute, 10*time.Millisecond)
 
 		// Verify User Transformation
 		eventSql = "select context_myuniqueid,context_id,context_ip from dev_integration_test_1.users;"
@@ -213,12 +213,12 @@ func testCases(t *testing.T) {
 			eventSql := "select anonymous_id, user_id from dev_integration_test_1.screens limit 1;"
 			_ = db.QueryRow(eventSql).Scan(&myEvent.anonymousID, &myEvent.userID)
 			return myEvent.anonymousID == "anonymousId_1"
-		}, time.Minute, 10*time.Millisecond)
+		}, 2*time.Minute, 10*time.Millisecond)
 		require.Eventually(t, func() bool {
 			eventSql := "select count(*) from dev_integration_test_1.screens;"
 			_ = db.QueryRow(eventSql).Scan(&myEvent.count)
 			return myEvent.count == "1"
-		}, time.Minute, 10*time.Millisecond)
+		}, 2*time.Minute, 10*time.Millisecond)
 
 		// Verify User Transformation
 		eventSql = "select prop_key,myuniqueid,ip from dev_integration_test_1.screens;"

router/batchrouter/batchrouter_test.go

@@ -302,7 +302,7 @@ var _ = Describe("BatchRouter", func() {
 
 			<-batchrouter.backendConfigInitialized
 			batchrouter.minIdleSleep = config.SingleValueLoader(time.Microsecond)
-			batchrouter.uploadFreq = config.SingleValueLoader(time.Microsecond)
+			batchrouter.uploadFreq = config.SingleValueLoader(5 * time.Millisecond)
 			batchrouter.mainLoopFreq = config.SingleValueLoader(time.Microsecond)
 			ctx, cancel := context.WithCancel(context.Background())
 			var wg sync.WaitGroup
@@ -409,19 +409,13 @@ var _ = Describe("BatchRouter", func() {
 				},
 			).AnyTimes()
 
-			c.mockBatchRouterJobsDB.EXPECT().UpdateJobStatus(gomock.Any(), gomock.Any(), []string{CustomVal["S3"]}, gomock.Any()).Times(1).
-				Do(func(ctx context.Context, statuses []*jobsdb.JobStatusT, _, _ interface{}) {
-					assertJobStatus(toRetryJobsList[0], statuses[0], jobsdb.Executing.State, `{}`, 130)
-					assertJobStatus(toRetryJobsList[1], statuses[1], jobsdb.Executing.State, `{}`, 4)
-				}).Return(nil)
-
 			c.mockBatchRouterJobsDB.EXPECT().WithUpdateSafeTx(gomock.Any(), gomock.Any()).Times(1).Do(func(ctx context.Context, f func(tx jobsdb.UpdateSafeTx) error) {
 				_ = f(jobsdb.EmptyUpdateSafeTx())
 			}).Return(nil)
 			c.mockBatchRouterJobsDB.EXPECT().UpdateJobStatusInTx(gomock.Any(), gomock.Any(), gomock.Any(), []string{CustomVal["S3"]}, gomock.Any()).Times(1).
 				Do(func(ctx context.Context, _ interface{}, statuses []*jobsdb.JobStatusT, _, _ interface{}) {
-					assertJobStatus(toRetryJobsList[0], statuses[0], jobsdb.Aborted.State, fmt.Sprintf(`{"firstAttemptedAt": "%s", "Error": "BRT: Batch destination source not found in config for sourceID: %s"}`, attempt1.Format(misc.RFC3339Milli), SourceIDEnabled+"random"), 130)
-					assertJobStatus(toRetryJobsList[1], statuses[1], jobsdb.Aborted.State, fmt.Sprintf(`{"firstAttemptedAt": "%s", "Error": "BRT: Batch destination source not found in config for sourceID: %s"}`, attempt2.Format(misc.RFC3339Milli), SourceIDEnabled+"random"), 4)
+					assertJobStatus(toRetryJobsList[0], statuses[0], jobsdb.Aborted.State, "{\"reason\":\"source_not_found\"}", 130)
+					assertJobStatus(toRetryJobsList[1], statuses[1], jobsdb.Aborted.State, "{\"reason\":\"source_not_found\"}", 4)
 				}).Return(nil)
 			c.mockProcErrorsDB.EXPECT().Store(gomock.Any(), gomock.Any()).Times(1).DoAndReturn(
 				func(ctx context.Context, _ []*jobsdb.JobT) error {
@@ -432,7 +426,7 @@ var _ = Describe("BatchRouter", func() {
 
 			<-batchrouter.backendConfigInitialized
 			batchrouter.minIdleSleep = config.SingleValueLoader(time.Microsecond)
-			batchrouter.uploadFreq = config.SingleValueLoader(time.Microsecond)
+			batchrouter.uploadFreq = config.SingleValueLoader(5 * time.Millisecond)
 			batchrouter.mainLoopFreq = config.SingleValueLoader(time.Microsecond)
 			done := make(chan struct{})
 			go func() {
@@ -620,11 +614,11 @@ func TestBatchRouter(t *testing.T) {
 					"table": "tracks"
 				}
 			}`),
-			Parameters: jsonb.RawMessage([]byte(fmt.Sprintf(`{
+			Parameters: []byte(fmt.Sprintf(`{
 				"source_id": %[1]q,
 				"destination_id": %[2]q,
 				"receivedAt": %[3]q
-			}`, bc.Sources[0].ID, s3Dest.ID, time.Now().Format(time.RFC3339)))),
+			}`, bc.Sources[0].ID, s3Dest.ID, time.Now().Format(time.RFC3339))),
 			CustomVal: s3Dest.DestinationDefinition.Name,
 			CreatedAt: time.Now(),
 		})
@@ -649,8 +643,8 @@ func TestBatchRouter(t *testing.T) {
 	)
 
 	batchrouter.minIdleSleep = config.SingleValueLoader(time.Microsecond)
-	batchrouter.uploadFreq = config.SingleValueLoader(time.Microsecond)
-	batchrouter.mainLoopFreq = config.SingleValueLoader(time.Microsecond)
+	batchrouter.uploadFreq = config.SingleValueLoader(5 * time.Millisecond)
+	batchrouter.mainLoopFreq = config.SingleValueLoader(time.Second)
 
 	err = routerDB.Store(context.Background(), jobs)
 	require.NoError(t, err)
@@ -665,7 +659,7 @@ func TestBatchRouter(t *testing.T) {
 			return false
 		}
 		return len(minioContents) == len(bcs)
-	}, 5*time.Second, 200*time.Millisecond)
+	}, 2*time.Minute, 200*time.Millisecond)
 
 	minioContents, err := minioResource.Contents(context.Background(), "")
 	require.NoError(t, err)

router/batchrouter/consumer_worker.go

@@ -0,0 +1,125 @@
+package batchrouter
+
+import (
+	"context"
+	"time"
+
+	"github.com/rudderlabs/rudder-go-kit/logger"
+	"github.com/rudderlabs/rudder-server/jobsdb"
+)
+
+// ConsumerWorker handles job consumption for a specific source-destination pair.
+// It acts as a buffer between job ingestion and batch processing, storing jobs
+// in partition queues for the batch worker to process.
+type ConsumerWorker struct {
+	sourceID  string
+	destID    string
+	jobChan   chan *jobsdb.JobT
+	ctx       context.Context
+	logger    logger.Logger
+	callbacks *ConsumerCallbacks
+}
+
+// ConsumerCallbacks contains callback functions needed by the consumer worker
+type ConsumerCallbacks struct {
+	GetUploadFrequency func() time.Duration
+	GetMaxBatchSize    func() int
+	ProcessJobs        func(ctx context.Context, jobs []*jobsdb.JobT) error
+	ReleaseWorker      func(sourceID, destID string)
+	PartitionJobs      func(jobs []*jobsdb.JobT) [][]*jobsdb.JobT
+}
+
+// NewConsumerWorker creates a new consumer worker for a source-destination pair
+func NewConsumerWorker(
+	sourceID string,
+	destID string,
+	jobChan chan *jobsdb.JobT,
+	ctx context.Context,
+	logger logger.Logger,
+	callbacks *ConsumerCallbacks,
+) *ConsumerWorker {
+	return &ConsumerWorker{
+		sourceID:  sourceID,
+		destID:    destID,
+		jobChan:   jobChan,
+		ctx:       ctx,
+		logger:    logger.Child("consumer-worker").With("sourceID", sourceID, "destID", destID),
+		callbacks: callbacks,
+	}
+}
+
+// Work implements the workerpool.Worker interface.
+// It consumes jobs from the channel and processes them in batches.
+// When either the batch size threshold or the upload frequency timer is reached,
+// it sends the collected jobs for processing using the callbacks.ProcessJobs method.
+// This directly processes the jobs without sending them back to the buffer.
+// Returns true if any jobs were processed.
+func (c *ConsumerWorker) Work() bool {
+	defer c.callbacks.ReleaseWorker(c.sourceID, c.destID)
+
+	uploadFrequency := c.callbacks.GetUploadFrequency()
+	maxBatchSize := c.callbacks.GetMaxBatchSize()
+
+	var jobs []*jobsdb.JobT
+	timer := time.NewTimer(uploadFrequency)
+	defer timer.Stop()
+
+	jobsProcessed := false
+
+	for {
+		select {
+		case <-c.ctx.Done():
+			if len(jobs) > 0 {
+				partitionedJobs := c.callbacks.PartitionJobs(jobs)
+				for _, jobBatch := range partitionedJobs {
+					if err := c.callbacks.ProcessJobs(c.ctx, jobBatch); err != nil {
+						c.logger.Errorf("Error processing jobs: %v", err)
+					}
+				}
+				jobsProcessed = true
+			}
+			return jobsProcessed
+
+		case job := <-c.jobChan:
+			jobs = append(jobs, job)
+			if len(jobs) >= maxBatchSize {
+				partitionedJobs := c.callbacks.PartitionJobs(jobs)
+				for _, jobBatch := range partitionedJobs {
+					if err := c.callbacks.ProcessJobs(c.ctx, jobBatch); err != nil {
+						c.logger.Errorf("Error processing jobs: %v", err)
+					}
+				}
+				jobs = nil
+				timer.Reset(uploadFrequency)
+				jobsProcessed = true
+			}
+
+		case <-timer.C:
+			if len(jobs) > 0 {
+				partitionedJobs := c.callbacks.PartitionJobs(jobs)
+				for _, jobBatch := range partitionedJobs {
+					if err := c.callbacks.ProcessJobs(c.ctx, jobBatch); err != nil {
+						c.logger.Errorf("Error processing jobs: %v", err)
+					}
+				}
+				jobsProcessed = true
+			}
+			timer.Reset(uploadFrequency)
+			return jobsProcessed
+		}
+	}
+}
+
+// SleepDurations returns the min and max sleep durations for the worker when idle
+func (c *ConsumerWorker) SleepDurations() (min, max time.Duration) {
+	return time.Millisecond * 100, time.Second * 5
+}
+
+// Stop implements the workerpool.Worker interface
+func (c *ConsumerWorker) Stop() {
+	// No cleanup needed
+}
+
+func (c *ConsumerWorker) IsIdle() bool {
+	return len(c.jobChan) == 0
+}

router/batchrouter/handle.go

@@ -109,6 +109,8 @@ type Handle struct {
 	backgroundCancel context.CancelFunc
 	backgroundWait   func() error
 
+	jobBuffer *JobBuffer // Added for channel-based job buffering
+
 	backendConfigInitializedOnce sync.Once
 	backendConfigInitialized     chan bool
 
@@ -121,9 +123,8 @@ type Handle struct {
 	encounteredMergeRuleMap   map[string]map[string]bool
 
 	limiter struct {
-		read    kitsync.Limiter
-		process kitsync.Limiter
-		upload  kitsync.Limiter
+		read   kitsync.Limiter
+		upload kitsync.Limiter
 	}
 
 	lastExecTimesMu sync.RWMutex
@@ -190,8 +191,7 @@ func (brt *Handle) getWorkerJobs(partition string) (workerJobs []*DestinationJob
 	if brt.skipFetchingJobs(partition) {
 		return
 	}
-
-	defer brt.limiter.read.Begin("")()
+	defer brt.limiter.read.Begin(partition)()
 
 	brt.configSubscriberMu.RLock()
 	destinationsMap := brt.destinationsMap
@@ -621,6 +621,7 @@ func (brt *Handle) updateJobStatus(batchJobs *BatchedJobs, isWarehouse bool, err
 	brt.failingDestinationsMu.Lock()
 	brt.failingDestinations[batchJobs.Connection.Destination.ID] = batchReqMetric.batchRequestFailed > 0
 	brt.failingDestinationsMu.Unlock()
+
 	var statusList []*jobsdb.JobStatusT
 
 	if isWarehouse && notifyWarehouseErr {

router/batchrouter/handle_lifecycle.go

@@ -141,15 +141,6 @@ func (brt *Handle) Setup(
 			return time.After(limiterStatsPeriod)
 		}),
 	)
-	brt.limiter.process = kitsync.NewLimiter(ctx, &limiterGroup, "brt_process",
-		getBatchRouterConfigInt("Limiter.process.limit", brt.destType, 20),
-		stats.Default,
-		kitsync.WithLimiterDynamicPeriod(config.GetDuration("BatchRouter.Limiter.process.dynamicPeriod", 1, time.Second)),
-		kitsync.WithLimiterTags(map[string]string{"destType": brt.destType}),
-		kitsync.WithLimiterStatsTriggerFunc(func() <-chan time.Time {
-			return time.After(limiterStatsPeriod)
-		}),
-	)
 	brt.limiter.upload = kitsync.NewLimiter(ctx, &limiterGroup, "brt_upload",
 		getBatchRouterConfigInt("Limiter.upload.limit", brt.destType, 50),
 		stats.Default,
@@ -161,9 +152,17 @@ func (brt *Handle) Setup(
 	)
 
 	brt.logger.Infof("BRT: Batch Router started: %s", destType)
-
 	brt.crashRecover()
 
+	if asynccommon.IsAsyncDestination(brt.destType) {
+		brt.startAsyncDestinationManager()
+	}
+
+	brt.backgroundGroup.Go(crash.Wrapper(func() error {
+		brt.backendConfigSubscriber()
+		return nil
+	}))
+
 	// periodically publish a zero counter for ensuring that stuck processing pipeline alert
 	// can always detect a stuck batch router
 	brt.backgroundGroup.Go(func() error {
@@ -192,14 +191,8 @@ func (brt *Handle) Setup(
 		return nil
 	}))
 
-	if asynccommon.IsAsyncDestination(destType) {
-		brt.startAsyncDestinationManager()
-	}
-
-	brt.backgroundGroup.Go(crash.Wrapper(func() error {
-		brt.backendConfigSubscriber()
-		return nil
-	}))
+	// Initialize the job buffer
+	brt.jobBuffer = NewJobBuffer(brt)
 }
 
 func (brt *Handle) setupReloadableVars() {
@@ -262,8 +255,21 @@ func (brt *Handle) Start() {
 
 // Shutdown stops the batch router
 func (brt *Handle) Shutdown() {
+	// Signal all goroutines to stop via context cancellation
+	brt.logger.Info("Initiating batch router shutdown")
 	brt.backgroundCancel()
+
+	// Stop all job buffer timers and workerPool
+	if brt.jobBuffer != nil {
+		brt.logger.Debug("Stopping JobBuffer and its resources")
+
+		// Stop the job buffer (which will also stop the worker pool)
+		brt.jobBuffer.Shutdown()
+	}
+
+	// Wait for all background goroutines to complete
 	_ = brt.backgroundWait()
+	brt.logger.Info("Batch router shutdown complete")
 }
 
 func (brt *Handle) initAsyncDestinationStruct(destination *backendconfig.DestinationT) {
@@ -399,7 +405,9 @@ func (brt *Handle) backendConfigSubscriber() {
 						if destination.DestinationDefinition.Name == brt.destType && destination.Enabled {
 							if _, ok := destinationsMap[destination.ID]; !ok {
 								destinationsMap[destination.ID] = &routerutils.DestinationWithSources{Destination: destination, Sources: []backendconfig.SourceT{}}
-								uploadIntervalMap[destination.ID] = brt.uploadInterval(destination.Config)
+								if asynccommon.IsAsyncDestination(brt.destType) {
+									uploadIntervalMap[destination.ID] = brt.uploadInterval(destination.Config)
+								}
 							}
 							destinationsMap[destination.ID].Sources = append(destinationsMap[destination.ID].Sources, source)
 							brt.refreshDestination(destination)