Drain polled tasks on shutdown

[yuandrew]

What was changed

Refactored worker shutdown to use a two-stage approach: pollers shut down first, then the task dispatcher drains any remaining tasks before exiting. This ensures tasks polled during shutdown are processed rather than silently dropped.

Key changes:

• Added pollerWG to baseWorker to track poller goroutines separately from the global stopWG
• A closer goroutine waits for all pollers to finish, then closes taskQueueCh
• runTaskDispatcher now ranges over taskQueueCh
• pollTask always sends to taskQueueCh
• Removed the 5s timeout hack in doPoll from Fix test flakes #2253 (no longer needed)

Why?

PR #2199 changed shutdown to let the server complete in-flight polls, instead of cancelling them. This exposed a pre-existing race when a poller receives a task during shutdown, Go would silently dropping the task. The dispatcher had the same issue — it could exit on stopCh before reading pending tasks from the channel.

This aligns the Go SDK's shutdown with how Core SDK handles it:

1. Set a flag so pollers stop polling after their current attempt
2. Close channels from pollers → task processing
3. Wait for all in-flight tasks to complete & channels to be empty

Checklist

1. Closes Drain polled tasks on shutdown #1197

2. How was this tested:

   • New unit test TestTaskNotDroppedDuringShutdown — verifies a task polled during shutdown is processed, not dropped
   • Existing TestDoPollGracefulShutdown — validates both graceful and legacy poll completion

3. Any docs updates needed?
   No

---

Note

Medium Risk
Touches worker shutdown concurrency (poller lifecycle, channel closing, dispatcher draining), which can affect graceful termination and task processing ordering. Adds blocking sends and new wait conditions that could deadlock if assumptions break, so needs careful review/testing.

Overview
Implements a two-stage worker shutdown so tasks polled during shutdown are still dispatched and processed instead of being dropped: pollers exit first, then the task dispatcher drains remaining items from taskQueueCh until it is closed.

This adds a dedicated pollerWG to track poller goroutines, closes taskQueueCh only after all pollers finish, updates runTaskDispatcher to range over the channel (no early exit on stopCh), and makes pollTask always send polled tasks to the dispatcher. It also removes the prior 5s shutdown timeout hack in doPoll (now waiting for poll completion) and adds a unit test TestTaskNotDroppedDuringShutdown to cover the shutdown race.

^{Written by Cursor Bugbot for commit be5c0e4. This will update automatically on new commits. Configure here.}

[Sushisource]

Just checking the stop channel is still used elsewhere since we removed it in two spots

[yuandrew]

yep! still used in plenty of other places

[yuandrew]

recheck

[cursor]

Unbounded pollerWG.Wait() bypasses user-configured stopTimeout

Medium Severity

bw.pollerWG.Wait() in Stop() blocks without any timeout, and runs before awaitWaitGroup(&bw.stopWG, bw.options.stopTimeout). Combined with doPoll now waiting unconditionally on <-doneC (bounded only by pollTaskServiceTimeOut = 70s), Stop() can block for up to 70 seconds before the user's stopTimeout even begins counting. Previously, a 5-second fallback cancellation bounded this. In failure scenarios (broken gRPC connection, unresponsive server), total Stop() duration becomes ~70s + stopTimeout instead of just stopTimeout.

Additional Locations (1)

• internal/internal_task_pollers.go#L316-L317

 

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

[cursor]

Debug print statements left in production code

High Severity

Two fmt.Println("DEBUG: ...") statements were left in the graceful shutdown path of doPoll. These will print to stdout in production whenever a worker shuts down with workerPollCompleteOnShutdown enabled. No other production code in this file uses fmt.Println — the existing usages are in test files and CLI tools.