Make test match prod behavior

Author: yuandrew

internal/internal_worker_base.go

@@ -579,10 +579,9 @@ func (bw *baseWorker) runTaskDispatcher() {
 		// is cancelled, so Wait returns immediately — we still process the
 		// task rather than dropping it.
 		if _, isPolledTask := task.(*polledTask); isPolledTask {
-			if err := bw.taskLimiter.Wait(bw.limiterContext); err != nil {
-				// Context cancelled during shutdown — skip rate limiting
-				// but still process remaining tasks.
-			}
+			// Ignore error: during shutdown the limiter context is
+			// cancelled, but we still process remaining tasks.
+			_ = bw.taskLimiter.Wait(bw.limiterContext)
 		}
 		bw.processTaskAsync(task)
 	}

internal/internal_worker_base_test.go

@@ -322,7 +322,8 @@ func TestTaskNotDroppedDuringShutdown(t *testing.T) {
 	taskProcessed := make(chan struct{})
 	pollStarted := make(chan struct{})
 
-	// A poller that blocks until returnTask is closed, then returns a task.
+	// A poller that blocks until returnTask is closed, then returns a task
+	// exactly once. Subsequent polls return nil so the poller can exit.
 	tp := &shutdownTaskPoller{
 		pollStarted: pollStarted,
 		returnTask:  make(chan struct{}),
@@ -348,26 +349,35 @@ func TestTaskNotDroppedDuringShutdown(t *testing.T) {
 
 	bw.Start()
 
-	// Wait for the poller to start polling
+	// Wait for the poller to be actively polling.
 	<-pollStarted
 
-	// Signal the poller to return a task, then stop the worker.
-	// The task should be processed, not dropped.
-	bw.noRepoll.Store(true)
-	close(bw.stopCh)
+	// Release the poller so it returns a task, then stop the worker.
+	// The poller returns a task and then nil on subsequent polls,
+	// allowing it to exit via noRepoll/stopCh during Stop().
 	close(tp.returnTask)
-	bw.limiterContextCancel()
+
+	// Stop exercises the real shutdown path: noRepoll, close(stopCh),
+	// limiterContextCancel, and awaitWaitGroup.
+	stopDone := make(chan struct{})
+	go func() {
+		bw.Stop()
+		close(stopDone)
+	}()
 
 	select {
 	case <-taskProcessed:
-		// Success: the task was dispatched and processed
+		// Success: the task was dispatched and processed during shutdown
 	case <-time.After(5 * time.Second):
 		t.Fatal("task polled during shutdown was not processed (dropped)")
 	}
 
-	// Wait for full cleanup. We already closed stopCh manually, so
-	// replicate the remaining Stop() logic.
-	awaitWaitGroup(&bw.stopWG, bw.options.stopTimeout)
+	select {
+	case <-stopDone:
+		// Stop completed cleanly
+	case <-time.After(5 * time.Second):
+		t.Fatal("Stop() did not return in time")
+	}
 }
 
 // shutdownTaskPoller blocks until returnTask is closed, then returns a task