Skip to content

Report poller target on each poll and in the worker heartbeat #1688

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
veeral-patel wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Report poller target on each poll and in the worker heartbeat #1688

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
492de21
Report poller target on each poll
veeral-patel Oct 15, 2026
464e07f
Report poller target in the worker heartbeat
veeral-patel Oct 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions crates/protos/protos/api_upstream/temporal/api/taskqueue/v1/message.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,14 @@ message TaskQueueMetadata {
google.protobuf.DoubleValue max_tasks_per_second = 1;
}

// Worker-reported state about the poller pool issuing a poll, used by the server to spread
// poller scaling suggestions across a fleet.
message PollerScalingInfo {
// Number of concurrent poll requests this worker is currently targeting for this task
// queue and task type. 0 if unknown, such as a poller pool of fixed size.
int32 poller_target = 1;
}

message TaskQueueVersioningInfo {
// Specifies which Deployment Version should receive new workflow executions and tasks of
// existing unversioned or AutoUpgrade workflows.
Expand Down
4 changes: 4 additions & 0 deletions crates/protos/protos/api_upstream/temporal/api/worker/v1/message.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ message WorkerPollerInfo {

// Set true if the number of concurrent pollers is auto-scaled
bool is_autoscaling = 3;

// Number of concurrent poll requests this poller pool is currently targeting, which may
// differ from current_pollers. 0 if unknown, such as a poller pool of fixed size.
int32 poller_target = 4;
}

message WorkerSlotsInfo {
Expand Down
8 changes: 8 additions & 0 deletions crates/protos/protos/api_upstream/temporal/api/workflowservice/v1/request_response.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,6 +314,10 @@ message PollWorkflowTaskQueueRequest {
// Worker deployment options that user has set in the worker.
temporal.api.deployment.v1.WorkerDeploymentOptions deployment_options = 6;

// Worker-reported poller pool state, used to spread poller scaling suggestions across
// a fleet. Optional.
temporal.api.taskqueue.v1.PollerScalingInfo poller_scaling_info = 11;

// Removed in 1.55.0; was temporal.api.worker.v1.WorkerHeartbeat worker_heartbeat
reserved 7;
reserved "worker_heartbeat";
Expand Down Expand Up @@ -544,6 +548,10 @@ message PollActivityTaskQueueRequest {
// Worker deployment options that user has set in the worker.
temporal.api.deployment.v1.WorkerDeploymentOptions deployment_options = 6;

// Worker-reported poller pool state, used to spread poller scaling suggestions across
// a fleet. Optional.
temporal.api.taskqueue.v1.PollerScalingInfo poller_scaling_info = 11;

// Removed in 1.55.0; was temporal.api.worker.v1.WorkerHeartbeat worker_heartbeat
reserved 7;
reserved "worker_heartbeat";
Expand Down
56 changes: 50 additions & 6 deletions crates/sdk-core/src/pollers/poll_buffer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,7 @@ impl LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind> {
num_pollers_handler: Option<impl Fn(usize) + Send + Sync + 'static>,
options: WorkflowTaskOptions,
last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
reported_target: Arc<AtomicCell<Option<i32>>>,
capabilities: Arc<NamespaceCapabilities>,
) -> Self {
let is_sticky = sticky_queue.is_some();
Expand All @@ -133,6 +134,7 @@ impl LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind> {
num_pollers_handler,
shutdown.clone(),
last_successful_poll_time,
reported_target,
capabilities.clone(),
);
if let Some(wftps) = options.wft_poller_shared.as_ref() {
Expand Down Expand Up @@ -168,7 +170,7 @@ impl LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind> {
} else {
None
};
let poll_fn = move |timeout_override: Option<Duration>| {
let poll_fn = move |timeout_override: Option<Duration>, poller_target: Option<i32>| {
let client = client.clone();
let task_queue = task_queue.clone();
let sticky_queue_name = sticky_queue.clone();
Expand All @@ -179,6 +181,7 @@ impl LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind> {
task_queue,
no_retry,
timeout_override,
poller_target,
},
PollWorkflowOptions { sticky_queue_name },
)
Expand Down Expand Up @@ -208,6 +211,7 @@ impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
num_pollers_handler: Option<impl Fn(usize) + Send + Sync + 'static>,
options: ActivityTaskOptions,
last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
reported_target: Arc<AtomicCell<Option<i32>>>,
capabilities: Arc<NamespaceCapabilities>,
) -> Self {
let pre_permit_delay = options
Expand All @@ -226,7 +230,7 @@ impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
} else {
None
};
let poll_fn = move |timeout_override: Option<Duration>| {
let poll_fn = move |timeout_override: Option<Duration>, poller_target: Option<i32>| {
let client = client.clone();
let task_queue = task_queue.clone();
async move {
Expand All @@ -236,6 +240,7 @@ impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
task_queue,
no_retry,
timeout_override,
poller_target,
},
PollActivityOptions {
max_tasks_per_sec: options.max_tps,
Expand All @@ -250,6 +255,7 @@ impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
num_pollers_handler,
shutdown.clone(),
last_successful_poll_time,
reported_target,
capabilities.clone(),
);
Self::new(
Expand All @@ -274,6 +280,7 @@ impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
shutdown: CancellationToken,
num_pollers_handler: Option<impl Fn(usize) + Send + Sync + 'static>,
last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
reported_target: Arc<AtomicCell<Option<i32>>>,
capabilities: Arc<NamespaceCapabilities>,
worker_commands_queue: bool,
) -> Self {
Expand All @@ -284,7 +291,7 @@ impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
} else {
None
};
let poll_fn = move |timeout_override: Option<Duration>| {
let poll_fn = move |timeout_override: Option<Duration>, poller_target: Option<i32>| {
let client = client.clone();
let task_queue = task_queue.clone();
async move {
Expand All @@ -294,6 +301,7 @@ impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
task_queue,
no_retry,
timeout_override,
poller_target,
},
PollNexusOptions {
worker_commands_queue,
Expand All @@ -311,6 +319,7 @@ impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
num_pollers_handler,
shutdown,
last_successful_poll_time,
reported_target,
capabilities.clone(),
),
None::<fn() -> BoxFuture<'static, ()>>,
Expand All @@ -335,7 +344,7 @@ where
SK: SlotKind + 'static,
{
fn new<FT, DelayFut, F>(
poll_fn: impl Fn(Option<Duration>) -> FT + Send + Sync + 'static,
poll_fn: impl Fn(Option<Duration>, Option<i32>) -> FT + Send + Sync + 'static,
permit_dealer: MeteredPermitDealer<SK>,
shutdown: CancellationToken,
mut poll_scaler: PollScaler<F>,
Expand Down Expand Up @@ -412,18 +421,19 @@ where
} else {
None
};
let poller_target = report_handle.reportable_target();
let capabilities = capabilities.clone();
let poll_task = tokio::spawn(async move {
let r = if capabilities.graceful_poll_shutdown() {
pf(timeout_override).await
pf(timeout_override, poller_target).await
} else {
let poll_interruptor = shutdown.cancelled().then(|_| async move {
if let Some(w) = poll_shutdown_interrupt_wait {
tokio::time::sleep(w).await;
}
});
tokio::select! {
r = pf(timeout_override) => r,
r = pf(timeout_override, poller_target) => r,
_ = poll_interruptor => return,
}
};
Expand Down Expand Up @@ -535,6 +545,7 @@ where
num_pollers_handler: Option<F>,
shutdown: CancellationToken,
last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
reported_target: Arc<AtomicCell<Option<i32>>>,
capabilities: Arc<NamespaceCapabilities>,
) -> Self {
let (active_tx, active_rx) = watch::channel(0);
Expand All @@ -559,10 +570,13 @@ where
ingested_last_period: Default::default(),
scale_up_allowed: AtomicBool::new(true),
last_successful_poll_time,
reported_target,
persistent_error_warning_state: Default::default(),
exponential_backoff: parking_lot::Mutex::new(TASK_POLL_BACKOFF.build()),
resource_exhausted_backoff: parking_lot::Mutex::new(THROTTLE_POLL_BACKOFF.build()),
});
// Seed the cell so a pool that never changes target still reports one.
report_handle.publish_target();
let rhc = report_handle.clone();
let ingestor_task = if behavior.is_autoscaling() {
// Here we spawn a task to periodically check if we should permit increasing the
Expand Down Expand Up @@ -623,6 +637,10 @@ struct PollScalerReportHandle {
ingested_last_period: AtomicUsize,
scale_up_allowed: AtomicBool,
last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
/// Mirrors reportable_target() so the heartbeat -- constructed before any
/// poller exists -- can read the live target. Same shape as
/// last_successful_poll_time above, which solves the same ordering problem.
reported_target: Arc<AtomicCell<Option<i32>>>,
persistent_error_warning_state: parking_lot::Mutex<PersistentPollErrorWarningState>,

// Exponential backoff for normal errors and resource exhausted errors
Expand All @@ -631,6 +649,19 @@ struct PollScalerReportHandle {
}

impl PollScalerReportHandle {
/// The pool's current scaling target, reported to the server on each poll so it can
/// spread poller-scaling suggestions across a fleet by share rather than letting them
/// land on whichever worker happens to receive the most tasks.
///
/// `None` for a fixed-size pool: it has no target to balance, and the server leaves
/// such a worker out of fairness entirely rather than assuming anything about it.
Comment on lines +652 to +657

@Sushisource Sushisource Oct 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/// The pool's current scaling target, reported to the server on each poll so it can
/// spread poller-scaling suggestions across a fleet by share rather than letting them
/// land on whichever worker happens to receive the most tasks.
///
/// `None` for a fixed-size pool: it has no target to balance, and the server leaves
/// such a worker out of fairness entirely rather than assuming anything about it.
/// The pool's current scaling target. `None` for a fixed-size pool.

fn reportable_target(&self) -> Option<i32> {
if !self.behavior.is_autoscaling() {
return None;
}
Some((*self.target.borrow()).try_into().unwrap_or(i32::MAX))
}

/// Returns (should_forward, backoff_duration)
/// - should_forward: true if the response should be passed on, false if it should be swallowed
/// - backoff_duration: Some(duration) if we should sleep before the next poll
Expand Down Expand Up @@ -738,6 +769,12 @@ impl PollScalerReportHandle {
*target = new_target;
true
});
self.publish_target();
}

/// Mirrors the current target into the cell the heartbeat reads.
fn publish_target(&self) {
self.reported_target.store(self.reportable_target());
}

/// We want to avoid scaling down on empty polls if the server has never made any scaling
Expand Down Expand Up @@ -950,6 +987,7 @@ mod tests {
wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))),
},
Arc::new(AtomicCell::new(None)),
Arc::new(AtomicCell::new(None)),
Arc::new(NamespaceCapabilities::default()),
);

Expand Down Expand Up @@ -1007,6 +1045,7 @@ mod tests {
wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(1)))),
},
Arc::new(AtomicCell::new(None)),
Arc::new(AtomicCell::new(None)),
Arc::new(NamespaceCapabilities::default()),
);

Expand Down Expand Up @@ -1084,6 +1123,7 @@ mod tests {
wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))),
},
Arc::new(AtomicCell::new(None)),
Arc::new(AtomicCell::new(None)),
Arc::new(NamespaceCapabilities::default()),
);

Expand Down Expand Up @@ -1193,6 +1233,7 @@ mod tests {
wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(INITIAL_POLLERS)))),
},
Arc::new(AtomicCell::new(None)),
Arc::new(AtomicCell::new(None)),
Arc::new(NamespaceCapabilities::default()),
));

Expand Down Expand Up @@ -1274,6 +1315,7 @@ mod tests {
wft_poller_shared: None,
},
Arc::new(AtomicCell::new(None)),
Arc::new(AtomicCell::new(None)),
Arc::new(NamespaceCapabilities::resolved(Capabilities {
poller_autoscaling: true,
..Default::default()
Expand Down Expand Up @@ -1352,6 +1394,7 @@ mod tests {
wft_poller_shared: None,
},
Arc::new(AtomicCell::new(None)),
Arc::new(AtomicCell::new(None)),
Arc::new(NamespaceCapabilities::resolved(Capabilities {
worker_poll_complete_on_shutdown: graceful,
..Default::default()
Expand Down Expand Up @@ -1420,6 +1463,7 @@ mod tests {
ingested_last_period: Default::default(),
scale_up_allowed: AtomicBool::new(true),
last_successful_poll_time: Arc::new(AtomicCell::new(None)),
reported_target: Arc::new(AtomicCell::new(None)),
persistent_error_warning_state: Default::default(),
exponential_backoff: parking_lot::Mutex::new(TASK_POLL_BACKOFF.build()),
resource_exhausted_backoff: parking_lot::Mutex::new(THROTTLE_POLL_BACKOFF.build()),
Expand Down
2 changes: 2 additions & 0 deletions crates/sdk-core/src/worker/activities.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1059,6 +1059,7 @@ mod tests {
max_tps: None,
},
Arc::new(AtomicCell::new(None)),
Arc::new(AtomicCell::new(None)),
Arc::new(NamespaceCapabilities::default()),
);
WorkerActivityTasks::new(
Expand Down Expand Up @@ -1146,6 +1147,7 @@ mod tests {
max_tps: None,
},
Arc::new(AtomicCell::new(None)),
Arc::new(AtomicCell::new(None)),
Arc::new(NamespaceCapabilities::default()),
);
let atm = WorkerActivityTasks::new(
Expand Down
17 changes: 16 additions & 1 deletion crates/sdk-core/src/worker/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,9 @@ use temporalio_common::protos::{
protocol::v1::Message as ProtocolMessage,
query::v1::WorkflowQueryResult,
sdk::v1::WorkflowTaskCompletedMetadata,
taskqueue::v1::{StickyExecutionAttributes, TaskQueue, TaskQueueMetadata},
taskqueue::v1::{
PollerScalingInfo, StickyExecutionAttributes, TaskQueue, TaskQueueMetadata,
},
worker::v1::{WorkerHeartbeat, WorkerSlotsInfo},
workflowservice::v1::{get_system_info_response::Capabilities, *},
},
Expand Down Expand Up @@ -444,6 +446,11 @@ pub struct PollOptions {
pub no_retry: Option<NoRetryOnMatching>,
/// Overrides the default RPC timeout for the poll request
pub timeout_override: Option<Duration>,
/// The poller pool's current scaling target, reported to the server so it can spread
/// poller-scaling suggestions across a fleet by share rather than letting them land on
/// whichever worker happens to receive the most tasks. `None` for a fixed-size pool,
/// which has no target to balance; the server then leaves this worker out of fairness.
Comment on lines +449 to +452

@Sushisource Sushisource Oct 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/// The poller pool's current scaling target, reported to the server so it can spread
/// poller-scaling suggestions across a fleet by share rather than letting them land on
/// whichever worker happens to receive the most tasks. `None` for a fixed-size pool,
/// which has no target to balance; the server then leaves this worker out of fairness.
/// The pool's current scaling target. `None` for a fixed-size pool.

pub poller_target: Option<i32>,
}
/// Additional options specific to workflow task polling
#[derive(Debug, Clone)]
Expand Down Expand Up @@ -496,6 +503,9 @@ impl WorkerClient for WorkerClientBag {
worker_instance_key: self.worker_instance_key.to_string(),
worker_control_task_queue: self.worker_control_task_queue(),
poller_group_id: Default::default(),
poller_scaling_info: poll_options
.poller_target
.map(|poller_target| PollerScalingInfo { poller_target }),
}
.into_request();
request.extensions_mut().insert(IsWorkerTaskLongPoll);
Expand Down Expand Up @@ -536,6 +546,9 @@ impl WorkerClient for WorkerClientBag {
worker_instance_key: self.worker_instance_key.to_string(),
worker_control_task_queue: self.worker_control_task_queue(),
poller_group_id: Default::default(),
poller_scaling_info: poll_options
.poller_target
.map(|poller_target| PollerScalingInfo { poller_target }),
}
.into_request();
request.extensions_mut().insert(IsWorkerTaskLongPoll);
Expand Down Expand Up @@ -1383,6 +1396,7 @@ mod tests {
task_queue: "application-queue".to_string(),
no_retry: None,
timeout_override: None,
poller_target: None,
},
PollNexusOptions {
worker_commands_queue: false,
Expand All @@ -1396,6 +1410,7 @@ mod tests {
task_queue: "worker-command-queue".to_string(),
no_retry: None,
timeout_override: None,
poller_target: None,
},
PollNexusOptions {
worker_commands_queue: true,
Expand Down
Loading
Loading