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 1 commit
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
19 changes: 19 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,25 @@ message TaskQueueMetadata {
google.protobuf.DoubleValue max_tasks_per_second = 1;
}

// Worker-reported state about the poller pool issuing this poll, used by the server to
// distribute poller-scaling suggestions fairly across a fleet.
//
// Scaling suggestions ride on dispatched tasks, and the matcher hands tasks to waiting
// pollers in FIFO order, so a worker holding a large share of a task queue's pollers
// receives suggestions at a proportionally higher rate. That makes poller growth
// proportional to current poller count, which preserves and amplifies any imbalance
// between otherwise identical workers. Reporting this lets the server withhold
// suggestions that would push a worker further from its share.
message PollerScalingInfo {
// Number of concurrent poll requests this worker is currently targeting for this task
// queue and task type -- the poller pool's own scaling target, not the count
// momentarily in flight. Fleet-wide for the pool, not per partition.
//
// 0 means unknown (for example, a worker whose poller count is fixed rather than
// auto-scaled); the server applies no share-based fairness to such a worker.
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
10 changes: 10 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,11 @@ 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 for fair distribution of poller scaling
// suggestions across a fleet. Optional; when absent the server applies no share-based
// fairness to this worker.
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 +549,11 @@ 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 for fair distribution of poller scaling
// suggestions across a fleet. Optional; when absent the server applies no share-based
// fairness to this worker.
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
29 changes: 23 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 @@ -168,7 +168,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 +179,7 @@ impl LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind> {
task_queue,
no_retry,
timeout_override,
poller_target,
},
PollWorkflowOptions { sticky_queue_name },
)
Expand Down Expand Up @@ -226,7 +227,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 +237,7 @@ impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
task_queue,
no_retry,
timeout_override,
poller_target,
},
PollActivityOptions {
max_tasks_per_sec: options.max_tps,
Expand Down Expand Up @@ -284,7 +286,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 +296,7 @@ impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
task_queue,
no_retry,
timeout_override,
poller_target,
},
PollNexusOptions {
worker_commands_queue,
Expand Down Expand Up @@ -335,7 +338,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 +415,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 @@ -631,6 +635,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
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