Make IsSticky() accurate for mixed-mode workflow pollers

Move the sticky/normal poll decision before slot reservation so that
SlotSupplier.ReserveSlot receives an accurate IsSticky() value even
when the worker uses a mixed-mode poller (the default for non-autoscaling
workers).

Key changes:
- Add pollHint struct to carry pre-decided sticky/normal choice
- Add decideNextPollKind() / undoPollDecision() on workflowTaskPoller
- Modify getNextPollRequest() to accept a *pollHint and skip inline
  re-decision when a hint is provided
- Thread the hint through taskPoller.PollTask -> poll -> getNextPollRequest
- Handle cancellation cleanup (undo counters if reservation fails)
- Expand test coverage for counter lifecycle and hint propagation

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: billrich2001

internal/internal_nexus_task_poller.go

@@ -98,7 +98,7 @@ func (ntp *nexusTaskPoller) Cleanup() error {
 }
 
 // PollTask polls a new task
-func (ntp *nexusTaskPoller) PollTask() (taskForWorker, error) {
+func (ntp *nexusTaskPoller) PollTask(_ *pollHint) (taskForWorker, error) {
 	return ntp.doPoll(ntp.poll)
 }
 

internal/internal_task_handlers_test.go

@@ -1832,7 +1832,7 @@ func (t *TaskHandlersTestSuite) TestLocalActivityRetry_Workflow() {
 	laTaskPoller := newLocalActivityPoller(params, laTunnel, nil, nil, stopCh)
 	go func() {
 		for {
-			task, _ := laTaskPoller.PollTask()
+			task, _ := laTaskPoller.PollTask(nil)
 			_ = laTaskPoller.ProcessTask(task)
 			// Quit after we've polled enough times
 			if laFailures.Load() == 4 {
@@ -1915,7 +1915,7 @@ func (t *TaskHandlersTestSuite) TestLocalActivityRetry_WorkflowTaskHeartbeatFail
 	doneCh := make(chan struct{})
 	go func() {
 		// laTaskPoller needs to poll the local activity and process it
-		task, err := laTaskPoller.PollTask()
+		task, err := laTaskPoller.PollTask(nil)
 		t.NoError(err)
 		err = laTaskPoller.ProcessTask(task)
 		t.NoError(err)

internal/internal_task_pollers.go

@@ -48,11 +48,19 @@ const (
 	Sticky
 )
 
+// pollHint carries a pre-decided sticky/normal decision from runPoller to PollTask.
+// When non-nil, the workflow task poller uses the hint instead of re-deciding.
+type pollHint struct {
+	isSticky bool
+}
+
 type (
 	// taskPoller interface to poll for tasks
 	taskPoller interface {
-		// PollTask polls for one new task
-		PollTask() (taskForWorker, error)
+		// PollTask polls for one new task. The hint, when non-nil, carries a
+		// pre-decided sticky/normal choice for workflow task pollers. Non-workflow
+		// pollers should ignore it.
+		PollTask(hint *pollHint) (taskForWorker, error)
 		// Called when the poller will no longer be polled. Presently only useful for
 		// workflow workers.
 		Cleanup() error
@@ -373,10 +381,13 @@ func (wtp *workflowTaskPoller) Cleanup() error {
 	return err
 }
 
-// PollTask polls a new task
-func (wtp *workflowTaskPoller) PollTask() (taskForWorker, error) {
+// PollTask polls a new task. If hint is non-nil, the pre-decided sticky/normal
+// choice is used instead of re-deciding in getNextPollRequest.
+func (wtp *workflowTaskPoller) PollTask(hint *pollHint) (taskForWorker, error) {
 	// Get the task.
-	workflowTask, err := wtp.doPoll(wtp.poll)
+	workflowTask, err := wtp.doPoll(func(ctx context.Context) (taskForWorker, error) {
+		return wtp.poll(ctx, hint)
+	})
 	if err != nil {
 		return nil, err
 	}
@@ -727,7 +738,7 @@ func (latp *localActivityTaskPoller) Cleanup() error {
 	return nil
 }
 
-func (latp *localActivityTaskPoller) PollTask() (taskForWorker, error) {
+func (latp *localActivityTaskPoller) PollTask(_ *pollHint) (taskForWorker, error) {
 	return latp.laTunnel.getTask(), nil
 }
 
@@ -901,16 +912,50 @@ func (wtp *workflowTaskPoller) updateBacklog(taskQueueKind enumspb.TaskQueueKind
 	wtp.requestLock.Unlock()
 }
 
+// decideNextPollKind commits to a sticky/normal decision for the next poll.
+// It increments the appropriate pending counter and returns whether the poll
+// should target the sticky queue. Only meaningful in Mixed mode; for Sticky and
+// NonSticky modes the answer is deterministic and no counter is changed.
+func (wtp *workflowTaskPoller) decideNextPollKind() (isSticky bool) {
+	if wtp.mode != Mixed || wtp.stickyCacheSize <= 0 {
+		return wtp.mode == Sticky
+	}
+	wtp.requestLock.Lock()
+	defer wtp.requestLock.Unlock()
+	if wtp.stickyBacklog > 0 || wtp.pendingStickyPollCount <= wtp.pendingRegularPollCount {
+		wtp.pendingStickyPollCount++
+		return true
+	}
+	wtp.pendingRegularPollCount++
+	return false
+}
+
+// undoPollDecision reverses the counter increment made by decideNextPollKind.
+// Call this when a pre-decided poll will not happen (e.g. slot reservation
+// was cancelled or failed).
+func (wtp *workflowTaskPoller) undoPollDecision(isSticky bool) {
+	if wtp.mode != Mixed || wtp.stickyCacheSize <= 0 {
+		return
+	}
+	if isSticky {
+		wtp.release(enumspb.TASK_QUEUE_KIND_STICKY)
+	} else {
+		wtp.release(enumspb.TASK_QUEUE_KIND_NORMAL)
+	}
+}
+
 // getNextPollRequest returns appropriate next poll request based on poller configuration and mode.
 // Simple rules:
 //  1. if mode is NonSticky, always poll from regular task queue
 //  2. if mode is Sticky, always poll from sticky task queue
 //  3. if mode is Mixed
 //     3.1. if sticky execution is disabled, always poll for regular task queue
-//     3.2. otherwise:
-//     3.2.1) if sticky task queue has backlog, always prefer to process sticky task first
-//     3.2.2) poll from the task queue that has less pending requests (prefer sticky when they are the same).
-func (wtp *workflowTaskPoller) getNextPollRequest() (request *workflowservice.PollWorkflowTaskQueueRequest) {
+//     3.2. otherwise, if a hint is provided, use the pre-decided kind (counter
+//     was already incremented by decideNextPollKind)
+//     3.3. otherwise (no hint / nil hint):
+//     3.3.1) if sticky task queue has backlog, always prefer to process sticky task first
+//     3.3.2) poll from the task queue that has less pending requests (prefer sticky when they are the same).
+func (wtp *workflowTaskPoller) getNextPollRequest(hint *pollHint) (request *workflowservice.PollWorkflowTaskQueueRequest) {
 	taskQueue := &taskqueuepb.TaskQueue{
 		Name: wtp.taskQueueName,
 		Kind: enumspb.TASK_QUEUE_KIND_NORMAL,
@@ -923,16 +968,27 @@ func (wtp *workflowTaskPoller) getNextPollRequest() (request *workflowservice.Po
 		taskQueue.Kind = enumspb.TASK_QUEUE_KIND_STICKY
 		taskQueue.NormalName = wtp.taskQueueName
 	} else if wtp.mode == Mixed {
-		wtp.requestLock.Lock()
-		if wtp.stickyBacklog > 0 || wtp.pendingStickyPollCount <= wtp.pendingRegularPollCount {
-			wtp.pendingStickyPollCount++
-			taskQueue.Name = getWorkerTaskQueue(wtp.stickyUUID)
-			taskQueue.Kind = enumspb.TASK_QUEUE_KIND_STICKY
-			taskQueue.NormalName = wtp.taskQueueName
+		if hint != nil {
+			// Use the pre-decided kind; counter was already incremented
+			// by decideNextPollKind.
+			if hint.isSticky {
+				taskQueue.Name = getWorkerTaskQueue(wtp.stickyUUID)
+				taskQueue.Kind = enumspb.TASK_QUEUE_KIND_STICKY
+				taskQueue.NormalName = wtp.taskQueueName
+			}
 		} else {
-			wtp.pendingRegularPollCount++
+			// Fallback: decide inline (original behavior).
+			wtp.requestLock.Lock()
+			if wtp.stickyBacklog > 0 || wtp.pendingStickyPollCount <= wtp.pendingRegularPollCount {
+				wtp.pendingStickyPollCount++
+				taskQueue.Name = getWorkerTaskQueue(wtp.stickyUUID)
+				taskQueue.Kind = enumspb.TASK_QUEUE_KIND_STICKY
+				taskQueue.NormalName = wtp.taskQueueName
+			} else {
+				wtp.pendingRegularPollCount++
+			}
+			wtp.requestLock.Unlock()
 		}
-		wtp.requestLock.Unlock()
 	} else {
 		panic("unknown workflow task poller mode")
 	}
@@ -973,12 +1029,12 @@ func (wtp *workflowTaskPoller) pollWorkflowTaskQueue(ctx context.Context, reques
 }
 
 // Poll for a single workflow task from the service
-func (wtp *workflowTaskPoller) poll(ctx context.Context) (taskForWorker, error) {
+func (wtp *workflowTaskPoller) poll(ctx context.Context, hint *pollHint) (taskForWorker, error) {
 	traceLog(func() {
 		wtp.logger.Debug("workflowTaskPoller::Poll")
 	})
 
-	request := wtp.getNextPollRequest()
+	request := wtp.getNextPollRequest(hint)
 	defer wtp.release(request.TaskQueue.GetKind())
 
 	response, err := wtp.pollWorkflowTaskQueue(ctx, request)
@@ -1207,7 +1263,7 @@ func (atp *activityTaskPoller) Cleanup() error {
 }
 
 // PollTask polls a new task
-func (atp *activityTaskPoller) PollTask() (taskForWorker, error) {
+func (atp *activityTaskPoller) PollTask(_ *pollHint) (taskForWorker, error) {
 	// Get the task.
 	activityTask, err := atp.doPoll(atp.poll)
 	if err != nil {

internal/internal_worker_base.go

@@ -444,16 +444,26 @@ func (bw *baseWorker) runPoller(taskWorker scalableTaskPoller) {
 				}
 			}
 
+			// Pre-decide sticky/normal for workflow task pollers so that
+			// IsSticky() is accurate at slot reservation time.
 			data := bw.options.slotReservationData
-			if wtp, ok := taskWorker.taskPoller.(*workflowTaskPoller); ok && wtp.mode == Sticky {
-				data.isSticky = true
+			var hint *pollHint
+			wtp, isWorkflowPoller := taskWorker.taskPoller.(*workflowTaskPoller)
+			if isWorkflowPoller {
+				isSticky := wtp.decideNextPollKind()
+				data.isSticky = isSticky
+				hint = &pollHint{isSticky: isSticky}
 			}
 
 			bw.stopWG.Add(1)
 			go func() {
 				defer bw.stopWG.Done()
 				s, err := bw.slotSupplier.ReserveSlot(ctx, &data)
 				if err != nil {
+					// Undo pre-decided counter since the poll will not happen.
+					if isWorkflowPoller {
+						wtp.undoPollDecision(hint.isSticky)
+					}
 					if !errors.Is(err, context.Canceled) {
 						bw.logger.Error("Error while trying to reserve slot", "error", err)
 						select {
@@ -467,6 +477,10 @@ func (bw *baseWorker) runPoller(taskWorker scalableTaskPoller) {
 				select {
 				case reserveChan <- s:
 				case <-ctx.Done():
+					// Worker stopped: undo pre-decided counter and release the slot.
+					if isWorkflowPoller {
+						wtp.undoPollDecision(hint.isSticky)
+					}
 					bw.releaseSlot(s, SlotReleaseReasonUnused)
 				}
 			}()
@@ -488,7 +502,7 @@ func (bw *baseWorker) runPoller(taskWorker scalableTaskPoller) {
 				if bw.pollerBalancer != nil {
 					bw.pollerBalancer.incrementPoller(taskWorker.taskPollerType)
 				}
-				bw.pollTask(taskWorker, permit)
+				bw.pollTask(taskWorker, permit, hint)
 				if bw.pollerBalancer != nil {
 					bw.pollerBalancer.decrementPoller(taskWorker.taskPollerType)
 				}
@@ -595,7 +609,7 @@ func (bw *baseWorker) runEagerTaskDispatcher() {
 	}
 }
 
-func (bw *baseWorker) pollTask(taskWorker scalableTaskPoller, slotPermit *SlotPermit) {
+func (bw *baseWorker) pollTask(taskWorker scalableTaskPoller, slotPermit *SlotPermit, hint *pollHint) {
 	var err error
 	var task taskForWorker
 	didSendTask := false
@@ -607,7 +621,7 @@ func (bw *baseWorker) pollTask(taskWorker scalableTaskPoller, slotPermit *SlotPe
 
 	bw.retrier.Throttle(bw.stopCh)
 	if bw.pollLimiter == nil || bw.pollLimiter.Wait(bw.limiterContext) == nil {
-		task, err = taskWorker.taskPoller.PollTask()
+		task, err = taskWorker.taskPoller.PollTask(hint)
 		bw.logPollTaskError(err)
 		if err != nil {
 			// We retry "non retriable" errors while long polling for a while, because some proxies return

internal/tuning.go

@@ -60,9 +60,7 @@ type SlotReservationInfo interface {
 	MetricsHandler() metrics.Handler
 	// IsSticky returns true if the slot being reserved will be used to poll a sticky task queue.
 	// This is only meaningful for workflow task slots. For activity and local activity slots,
-	// this will always return false. When the worker is configured with a mixed-mode poller
-	// (the default when not using autoscaling), this will also return false because the
-	// sticky-vs-normal decision is made after the slot is reserved.
+	// this will always return false.
 	IsSticky() bool
 }