use std::collections::HashMap;
use std::sync::atomic::{AtomicI32, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};
use dex_protocol::dex::poll_worker_task::Request as TaskRequest;
use dex_protocol::dex::poll_worker_task_frame::Output;
use dex_protocol::dex::poll_worker_task_request::Frame as RequestFrame;
use dex_protocol::dex::poll_worker_task_response::Frame as ResponseFrame;
use dex_protocol::dex::worker_task_service_client::WorkerTaskServiceClient;
use dex_protocol::dex::{
PollWorkerTask, PollWorkerTaskFrame, PollWorkerTaskPong, PollWorkerTaskRegister,
PollWorkerTaskReject, PollWorkerTaskRejectReason, PollWorkerTaskRequest, WorkerErrorResponse,
};
use tokio::sync::{Notify, mpsc, watch};
use tokio::task::{JoinHandle, JoinSet};
use tokio::time::{Instant, sleep, sleep_until, timeout_at};
use tokio_stream::{StreamExt, wrappers::ReceiverStream};
use tonic::transport::Endpoint;
use tonic::{Code, Status};
use crate::Credential;
use crate::context::InvocationCancellation;
use crate::credential::{CredentialChannel, CredentialInterceptor};
use crate::pull_registration::PullRegistration;
use crate::worker::worker_error_response;
use crate::worker_dispatcher::WorkerDispatcher;
use crate::{PullWorkerOptions, SdkError, SdkResult};
mod reconnect;
const DRAINING: u8 = 1;
const FORCED: u8 = 2;
#[derive(Clone, Copy)]
struct StopSignal {
phase: u8,
deadline: Option<Instant>,
}
pub(crate) fn outbound_task_channel(
endpoint: Endpoint,
credential: Option<Arc<dyn Credential>>,
) -> CredentialChannel {
CredentialChannel::new(
endpoint.connect_lazy(),
CredentialInterceptor::new(credential),
)
}
pub(crate) struct PullWorker {
dispatcher: WorkerDispatcher,
client: WorkerTaskServiceClient<CredentialChannel>,
credential: Option<Arc<dyn Credential>>,
app_worker_id: String,
flow_types: Vec<String>,
queues: Vec<String>,
max_concurrent_tasks: usize,
reserved_rpc_tasks: usize,
reserved_rpc_declines: Arc<AtomicU64>,
idle_streams_per_queue: usize,
stop_drain_timeout: Duration,
pub(crate) registration: Arc<PullRegistration>,
should_warn_rpc_capacity: bool,
stop_on_permanent_stream_rejection: bool,
}
struct PullState {
streams: HashMap<u64, StreamSlot>,
busy_tasks: usize,
next_stream_id: u64,
is_stopping: bool,
logged_unsupported: bool,
}
struct StreamSlot {
queue: String,
is_busy: bool,
is_registered: bool,
is_retiring: bool,
is_draining: bool,
retirement: Arc<Notify>,
}
struct PullShared {
registration: Arc<PullRegistration>,
state: Mutex<PullState>,
notify: Notify,
max_concurrent_tasks: usize,
reserved_rpc_tasks: usize,
reserved_rpc_declines: Arc<AtomicU64>,
idle_streams_per_queue: usize,
}
struct ActiveTask {
task_id: String,
cancellation: InvocationCancellation,
transport_cancellation: InvocationCancellation,
cancel_code: Arc<AtomicI32>,
handler: JoinHandle<()>,
deadline: Option<Instant>,
}
struct StreamRegistration {
shared: Arc<PullShared>,
stream_id: u64,
}
impl Drop for StreamRegistration {
fn drop(&mut self) {
let mut state = self.shared.state.lock().expect("pull state lock");
if let Some(slot) = state.streams.remove(&self.stream_id)
&& slot.is_busy
{
state.busy_tasks -= 1;
}
self.shared.notify.notify_one();
}
}
impl PullWorker {
pub(crate) fn new(
options: &PullWorkerOptions,
dispatcher: &WorkerDispatcher,
channel: CredentialChannel,
) -> SdkResult<Self> {
let max_concurrent_tasks = match options.max_concurrent_tasks {
Some(value) => value,
None => std::thread::available_parallelism()
.map(std::num::NonZeroUsize::get)
.map_err(|error| SdkError::Service {
service: crate::ServiceError::local(
"worker",
format!("read CPU count: {error}"),
),
})?,
};
let reserved_rpc_tasks = options
.reserved_rpc_tasks
.unwrap_or(usize::from(max_concurrent_tasks >= 2));
if reserved_rpc_tasks >= max_concurrent_tasks {
return Err(invalid(
"reserved_rpc_tasks must be less than max_concurrent_tasks",
));
}
let idle_streams_per_queue = options.idle_streams_per_queue.unwrap_or(1);
let stop_drain_timeout = options
.stop_drain_timeout
.unwrap_or(Duration::from_secs(30));
if max_concurrent_tasks == 0 || idle_streams_per_queue == 0 || stop_drain_timeout.is_zero()
{
return Err(invalid(
"pull Worker task limit, idle stream count, and drain timeout must be positive",
));
}
let app_worker_id = match &options.app_worker_id {
Some(value) => value.clone(),
None => generate_app_worker_id(),
};
validate_app_worker_id(&app_worker_id)?;
let flow_types = dispatcher.flow_types();
if flow_types
.iter()
.any(|flow_type| flow_type == &app_worker_id)
{
return Err(invalid(format!(
"app Worker ID conflicts with Flow type {app_worker_id:?}"
)));
}
let queues: Vec<String> = std::iter::once(app_worker_id.clone())
.chain(flow_types.iter().cloned())
.collect();
Ok(Self {
dispatcher: dispatcher.clone(),
credential: channel.credential(),
client: WorkerTaskServiceClient::new(channel),
app_worker_id,
flow_types,
registration: Arc::new(PullRegistration::new(&queues)),
stop_on_permanent_stream_rejection: options.stop_on_permanent_stream_rejection,
queues,
max_concurrent_tasks,
reserved_rpc_tasks,
reserved_rpc_declines: Arc::new(AtomicU64::new(0)),
should_warn_rpc_capacity: options.reserved_rpc_tasks.is_none()
&& max_concurrent_tasks == 1,
idle_streams_per_queue,
stop_drain_timeout,
})
}
pub(crate) fn reserved_rpc_declines(&self) -> u64 {
self.reserved_rpc_declines.load(Ordering::Relaxed)
}
pub(crate) async fn run(&self, mut shutdown: watch::Receiver<bool>) -> SdkResult<()> {
if self.should_warn_rpc_capacity {
tracing::warn!(
"RPCs may wait behind long-running Steps; use max_concurrent_tasks >= 2 or set reserved_rpc_tasks explicitly"
);
}
let shared = Arc::new(PullShared {
registration: Arc::clone(&self.registration),
state: Mutex::new(PullState {
streams: HashMap::new(),
busy_tasks: 0,
next_stream_id: 0,
is_stopping: false,
logged_unsupported: false,
}),
notify: Notify::new(),
max_concurrent_tasks: self.max_concurrent_tasks,
reserved_rpc_tasks: self.reserved_rpc_tasks,
reserved_rpc_declines: Arc::clone(&self.reserved_rpc_declines),
idle_streams_per_queue: self.idle_streams_per_queue,
});
let (stop, stop_receiver) = watch::channel(StopSignal {
phase: 0,
deadline: None,
});
let mut streams = JoinSet::new();
loop {
if *shutdown.borrow()
|| (self.stop_on_permanent_stream_rejection
&& self.registration.permanent_error().is_some())
{
break;
}
self.reconcile(&shared, &stop_receiver, &mut streams);
tokio::select! {
change = shutdown.changed() => {
if change.is_err() || *shutdown.borrow() { break; }
}
() = shared.notify.notified() => {}
result = streams.join_next(), if !streams.is_empty() => {
if let Some(Err(error)) = result { tracing::error!(%error, "pull Worker stream failed"); }
}
}
}
shared.state.lock().expect("pull state lock").is_stopping = true;
self.registration.finish(None);
let stop_deadline = Instant::now() + self.stop_drain_timeout;
stop.send_replace(StopSignal {
phase: DRAINING,
deadline: Some(stop_deadline),
});
if timeout_at(stop_deadline, async {
while let Some(result) = streams.join_next().await {
if let Err(error) = result {
tracing::error!(%error, "pull Worker stream failed");
}
}
})
.await
.is_err()
{
stop.send_replace(StopSignal {
phase: FORCED,
deadline: Some(stop_deadline),
});
while let Some(result) = streams.join_next().await {
if let Err(error) = result {
tracing::error!(%error, "pull Worker stream failed");
}
}
}
if self.stop_on_permanent_stream_rejection
&& let Some(rejection) = self.registration.permanent_error()
{
return Err(SdkError::from_pull_registration(rejection));
}
Ok(())
}
fn reconcile(
&self,
shared: &Arc<PullShared>,
stop: &watch::Receiver<StopSignal>,
streams: &mut JoinSet<()>,
) {
let mut state = shared.state.lock().expect("pull state lock");
if state.is_stopping {
return;
}
for queue in &self.queues {
let registered_idle = state
.streams
.values()
.filter(|slot| {
slot.queue == *queue
&& slot.is_registered
&& !slot.is_busy
&& !slot.is_retiring
&& !slot.is_draining
})
.count();
let mut surplus = registered_idle.saturating_sub(self.idle_streams_per_queue);
for slot in state.streams.values_mut().filter(|slot| {
slot.queue == *queue
&& slot.is_registered
&& !slot.is_busy
&& !slot.is_retiring
&& !slot.is_draining
}) {
if surplus == 0 {
break;
}
slot.is_retiring = true;
slot.retirement.notify_one();
surplus -= 1;
}
let idle = state
.streams
.values()
.filter(|slot| slot.queue == *queue && !slot.is_busy && !slot.is_retiring)
.count();
for _ in idle..self.idle_streams_per_queue {
let stream_id = state.next_stream_id;
state.next_stream_id += 1;
let retirement = Arc::new(Notify::new());
state.streams.insert(
stream_id,
StreamSlot {
queue: queue.clone(),
is_busy: false,
is_registered: false,
is_retiring: false,
is_draining: false,
retirement: Arc::clone(&retirement),
},
);
let worker = PullStream {
client: self.client.clone(),
credential: self.credential.clone(),
dispatcher: self.dispatcher.clone(),
app_worker_id: self.app_worker_id.clone(),
flow_types: self.flow_types.clone(),
queue: queue.clone(),
stream_id,
shared: Arc::clone(shared),
stop: stop.clone(),
stop_drain_timeout: self.stop_drain_timeout,
is_registered: false,
has_recorded_result: false,
retirement,
};
streams.spawn(async move { worker.run().await });
}
}
}
}
struct PullStream {
client: WorkerTaskServiceClient<CredentialChannel>,
credential: Option<Arc<dyn Credential>>,
dispatcher: WorkerDispatcher,
app_worker_id: String,
flow_types: Vec<String>,
queue: String,
stream_id: u64,
shared: Arc<PullShared>,
stop: watch::Receiver<StopSignal>,
stop_drain_timeout: Duration,
is_registered: bool,
has_recorded_result: bool,
retirement: Arc<Notify>,
}
impl PullStream {
async fn poll(&mut self) -> (bool, bool, Result<(), Status>) {
self.is_registered = false;
self.has_recorded_result = false;
self.shared
.state
.lock()
.expect("pull state lock")
.streams
.get_mut(&self.stream_id)
.expect("registered pull stream")
.is_draining = false;
let result = self.poll_stream().await;
if !result
.2
.as_ref()
.is_err_and(|status| self.should_defer_stream_rejection(status))
{
self.record_stream_result(result.2.as_ref().err());
}
result
}
fn should_defer_stream_rejection(&self, status: &Status) -> bool {
status.code() == Code::Unauthenticated
&& self
.credential
.as_ref()
.is_some_and(|credential| credential.token_after_rejection().is_some())
}
fn record_stream_result(&mut self, failure: Option<&Status>) {
if self.has_recorded_result {
return;
}
self.has_recorded_result = true;
self.shared.registration.ended(
&self.queue,
self.is_registered,
failure.filter(|status| is_permanent(status.code())),
);
self.is_registered = false;
self.shared
.state
.lock()
.expect("pull state lock")
.streams
.get_mut(&self.stream_id)
.expect("registered pull stream")
.is_registered = false;
self.shared.notify.notify_one();
}
async fn poll_stream(&mut self) -> (bool, bool, Result<(), Status>) {
let (outbound_sender, receiver) = mpsc::channel(16);
let mut sender = Some(outbound_sender);
let register = PollWorkerTaskRegister {
queue: self.queue.clone(),
app_worker_id: self.app_worker_id.clone(),
flow_types: self.flow_types.clone(),
sdk_name: "rust".to_string(),
sdk_version: env!("CARGO_PKG_VERSION").to_string(),
};
if sender
.as_ref()
.expect("open Worker stream")
.send(PollWorkerTaskRequest {
frame: Some(RequestFrame::Register(register)),
})
.await
.is_err()
{
return (
false,
false,
Err(Status::internal("Worker register channel closed")),
);
}
let opening = tokio::select! {
result = self.client.poll_worker_task(ReceiverStream::new(receiver)) => result,
change = self.stop.changed() => {
if change.is_err() || self.stop.borrow().phase != 0 {
return (false, false, Ok(()));
}
return (false, false, Err(Status::unavailable("Worker poll interrupted")));
}
};
let mut response_stream = match opening {
Ok(response) => response.into_inner(),
Err(status) => return (false, false, Err(status)),
};
let mut was_registered = false;
let mut supports_unknown_name_rejection = false;
let mut is_draining = false;
let mut active: Option<ActiveTask> = None;
loop {
let deadline = active.as_ref().and_then(|task| task.deadline);
tokio::select! {
response = response_stream.message() => {
match response {
Ok(Some(response)) => {
if !was_registered {
match response.frame {
Some(ResponseFrame::Registered(registered)) => { supports_unknown_name_rejection = registered.supports_unknown_name_rejection; was_registered = true; self.is_registered = true; self.shared.state.lock().expect("pull state lock").streams.get_mut(&self.stream_id).expect("registered pull stream").is_registered = true; self.shared.registration.registered(&self.queue); self.shared.notify.notify_one(); }
Some(ResponseFrame::Reconnect(_)) => return (false, true, Ok(())),
_ => return (false, false, Err(Status::internal("Worker task stream did not register"))),
}
continue;
}
match response.frame {
Some(ResponseFrame::Task(task)) => {
let unknown_reason = self.dispatcher.unknown_task_reason(&task);
let rejection = unknown_reason.filter(|_| supports_unknown_name_rejection).or_else(|| {
let reserve_task = if unknown_reason.is_none() { Some(&task) } else { None };
(active.is_some() || is_draining || !self.shared.try_accept(self.stream_id, reserve_task))
.then_some(PollWorkerTaskRejectReason::NoCapacity)
});
if let Some(reason) = rejection {
let reject = PollWorkerTaskReject { task_id: task.task_id, reason: reason as i32 };
if let Err(status) = send_control_frame(sender.as_ref().expect("open Worker stream"), RequestFrame::Reject(reject)).await {
if !self.should_defer_stream_rejection(&status) {
self.record_stream_result(Some(&status));
}
self.cancel_and_wait(&mut active).await;
return (was_registered, false, Err(status));
}
} else {
let cancellation = InvocationCancellation::new();
let transport_cancellation = InvocationCancellation::new();
let cancel_code = Arc::new(AtomicI32::new(0));
let task_id = task.task_id.clone();
let deadline = task.deadline.as_ref().map(task_deadline);
let handler = tokio::spawn(execute_task(self.dispatcher.clone(), task, cancellation.clone(), transport_cancellation.clone(), Arc::clone(&cancel_code), sender.as_ref().expect("open Worker stream").clone()));
active = Some(ActiveTask { task_id, cancellation, transport_cancellation, cancel_code, handler, deadline });
self.shared.notify.notify_one();
}
}
Some(ResponseFrame::CancelTask(cancel)) => {
if let Some(task) = &active && task.task_id == cancel.task_id {
task.cancel_code.store(Code::Cancelled as i32, Ordering::Release);
task.cancellation.cancel();
}
}
Some(ResponseFrame::Ping(_)) => {
if let Err(status) = send_control_frame(sender.as_ref().expect("open Worker stream"), RequestFrame::Pong(PollWorkerTaskPong {})).await {
if !self.should_defer_stream_rejection(&status) {
self.record_stream_result(Some(&status));
}
self.cancel_and_wait(&mut active).await;
return (was_registered, false, Err(status));
}
}
Some(ResponseFrame::Reconnect(_)) => {
self.shared.mark_draining(self.stream_id);
is_draining = true;
if active.is_none() { return (was_registered, true, Ok(())); }
self.shared.notify.notify_one();
}
_ => {}
}
}
Ok(None) => {
self.record_stream_result(None);
self.cancel_and_wait(&mut active).await;
return (was_registered, false, Ok(()));
}
Err(status) => {
if !self.should_defer_stream_rejection(&status) {
self.record_stream_result(Some(&status));
}
self.cancel_and_wait(&mut active).await;
return (was_registered, false, Err(status));
}
}
}
result = async { (&mut active.as_mut().expect("active task").handler).await }, if active.is_some() => {
if let Err(error) = result { tracing::error!(%error, "pull Worker handler failed"); }
active = None;
let should_retire = self.shared.release(self.stream_id);
if should_retire || is_draining || self.stop.borrow().phase != 0 {
let is_reconnecting = !should_retire && is_draining && self.stop.borrow().phase == 0;
let result = close_poll_send_and_wait(&mut sender, &mut response_stream, &mut self.stop, self.stop_drain_timeout).await;
return (was_registered, is_reconnecting, result);
}
}
() = async { sleep_until(deadline.expect("task deadline")).await }, if deadline.is_some() => {
if let Some(task) = &mut active {
task.cancel_code.store(Code::DeadlineExceeded as i32, Ordering::Release);
task.cancellation.cancel();
task.deadline = None;
}
}
() = self.retirement.notified(), if active.is_none() => {
let result = close_poll_send_and_wait(&mut sender, &mut response_stream, &mut self.stop, self.stop_drain_timeout).await;
return (was_registered, false, result);
}
change = self.stop.changed() => {
if change.is_err() || self.stop.borrow().phase != 0 {
is_draining = true;
if self.stop.borrow().phase >= FORCED && let Some(task) = &active {
task.cancel_code.store(Code::Cancelled as i32, Ordering::Release);
task.cancellation.cancel();
task.transport_cancellation.cancel();
}
if active.is_none() { return (was_registered, false, Ok(())); }
}
}
}
}
}
async fn cancel_and_wait(&self, active: &mut Option<ActiveTask>) {
if let Some(task) = active.take() {
task.cancel_code
.store(Code::Cancelled as i32, Ordering::Release);
task.cancellation.cancel();
task.transport_cancellation.cancel();
if let Err(error) = task.handler.await {
tracing::error!(%error, "pull Worker handler failed");
}
self.shared.release(self.stream_id);
}
}
}
impl PullShared {
fn try_accept(&self, stream_id: u64, task: Option<&PollWorkerTask>) -> bool {
let mut state = self.state.lock().expect("pull state lock");
if state.is_stopping || state.busy_tasks >= self.max_concurrent_tasks {
return false;
}
let busy_tasks = state.busy_tasks;
let Some(slot) = state.streams.get_mut(&stream_id) else {
return false;
};
if slot.is_busy || slot.is_retiring || slot.is_draining {
return false;
}
if self.reserved_rpc_tasks > 0
&& busy_tasks >= self.max_concurrent_tasks - self.reserved_rpc_tasks
&& let Some(task) = task
{
let flow_type = match &task.request {
Some(TaskRequest::WaitFor(request)) => Some(&request.flow_type),
Some(TaskRequest::Execute(request)) => Some(&request.flow_type),
_ => None,
};
if let Some(flow_type) = flow_type {
self.reserved_rpc_declines.fetch_add(1, Ordering::Relaxed);
tracing::debug!(%flow_type, "decline Step task for reserved RPC capacity");
return false;
}
}
slot.is_busy = true;
state.busy_tasks += 1;
true
}
fn release(&self, stream_id: u64) -> bool {
let mut state = self.state.lock().expect("pull state lock");
let slot = state
.streams
.get_mut(&stream_id)
.expect("registered pull stream");
assert!(slot.is_busy, "pull stream must have an active task");
slot.is_busy = false;
let queue = slot.queue.clone();
state.busy_tasks -= 1;
let idle_replacements = state
.streams
.iter()
.filter(|(candidate_id, slot)| {
**candidate_id != stream_id
&& slot.queue == queue
&& slot.is_registered
&& !slot.is_busy
&& !slot.is_retiring
&& !slot.is_draining
})
.count();
let should_retire = idle_replacements >= self.idle_streams_per_queue;
if should_retire {
state
.streams
.get_mut(&stream_id)
.expect("registered pull stream")
.is_retiring = true;
}
self.notify.notify_one();
should_retire
}
fn mark_draining(&self, stream_id: u64) {
self.state
.lock()
.expect("pull state lock")
.streams
.get_mut(&stream_id)
.expect("registered pull stream")
.is_draining = true;
}
fn is_retiring(&self, stream_id: u64) -> bool {
self.state
.lock()
.expect("pull state lock")
.streams
.get(&stream_id)
.expect("registered pull stream")
.is_retiring
}
}
async fn execute_task(
dispatcher: WorkerDispatcher,
task: PollWorkerTask,
cancellation: InvocationCancellation,
transport_cancellation: InvocationCancellation,
cancel_code: Arc<AtomicI32>,
sender: mpsc::Sender<PollWorkerTaskRequest>,
) {
let task_id = task.task_id;
match task.request {
Some(TaskRequest::WaitFor(request)) => {
let mut output = dispatcher
.invoke_wait_for_with_cancellation(request, cancellation.clone())
.into_stream(crate::worker::worker_status);
while let Some(result) = output.next().await {
let frame = match result {
Ok(value) => Output::WaitFor(value),
Err(status) => Output::WorkerError(with_cancel_status(
worker_error_from_status(status),
&cancel_code,
)),
};
if send_task_frame(&sender, &task_id, frame, &transport_cancellation)
.await
.is_err()
{
cancellation.cancel();
transport_cancellation.cancel();
}
}
}
Some(TaskRequest::Execute(request)) => {
let mut output = dispatcher
.invoke_execute_with_cancellation(request, cancellation.clone())
.into_stream(crate::worker::worker_status);
while let Some(result) = output.next().await {
let frame = match result {
Ok(value) => Output::Execute(value),
Err(status) => Output::WorkerError(with_cancel_status(
worker_error_from_status(status),
&cancel_code,
)),
};
if send_task_frame(&sender, &task_id, frame, &transport_cancellation)
.await
.is_err()
{
cancellation.cancel();
transport_cancellation.cancel();
}
}
}
Some(TaskRequest::Rpc(request)) => {
let frame = match dispatcher
.invoke_rpc_with_cancellation(request, cancellation)
.await
{
Ok(value) => Output::Rpc(value),
Err(error) => Output::WorkerError(with_cancel_status(
worker_error_response(error),
&cancel_code,
)),
};
if let Err(error) =
send_task_frame(&sender, &task_id, frame, &transport_cancellation).await
{
tracing::warn!(%error, "send Worker RPC result");
}
}
None => {
let error = WorkerErrorResponse {
detail: "Worker task request is missing".to_string(),
grpc_status: Code::InvalidArgument as i32,
..Default::default()
};
if let Err(error) = send_task_frame(
&sender,
&task_id,
Output::WorkerError(error),
&transport_cancellation,
)
.await
{
tracing::warn!(%error, "send Worker task error");
}
}
}
}
async fn send_control_frame(
sender: &mpsc::Sender<PollWorkerTaskRequest>,
frame: RequestFrame,
) -> Result<(), Status> {
sender
.send(PollWorkerTaskRequest { frame: Some(frame) })
.await
.map_err(|error| {
Status::unavailable(format!(
"Worker task stream cannot send control frame: {error}"
))
})
}
async fn close_poll_send_and_wait(
sender: &mut Option<mpsc::Sender<PollWorkerTaskRequest>>,
responses: &mut tonic::Streaming<dex_protocol::dex::PollWorkerTaskResponse>,
stop: &mut watch::Receiver<StopSignal>,
stop_drain_timeout: Duration,
) -> Result<(), Status> {
sender.take();
let mut deadline = stop
.borrow()
.deadline
.unwrap_or_else(|| Instant::now() + stop_drain_timeout);
loop {
tokio::select! {
response = responses.message() => {
if response?.is_none() { return Ok(()); }
}
() = sleep_until(deadline) => {
return Err(Status::deadline_exceeded("Worker task stream did not close after drain"));
}
change = stop.changed() => {
if change.is_err() {
return Err(Status::cancelled("Worker stop signal closed"));
}
if let Some(stop_deadline) = stop.borrow_and_update().deadline {
deadline = stop_deadline;
}
}
}
}
}
fn with_cancel_status(
mut error: WorkerErrorResponse,
cancel_code: &AtomicI32,
) -> WorkerErrorResponse {
let code = cancel_code.load(Ordering::Acquire);
if code != 0 {
error.grpc_status = code;
}
error
}
async fn send_task_frame(
sender: &mpsc::Sender<PollWorkerTaskRequest>,
task_id: &str,
output: Output,
cancellation: &InvocationCancellation,
) -> Result<(), Status> {
let request = PollWorkerTaskRequest {
frame: Some(RequestFrame::TaskFrame(PollWorkerTaskFrame {
task_id: task_id.to_string(),
output: Some(output),
})),
};
tokio::select! {
biased;
result = sender.send(request) => result.map_err(|error| Status::unavailable(format!("Worker task stream cannot send task frame: {error}"))),
() = cancellation.cancelled() => Err(Status::cancelled("Worker task was cancelled")),
}
}
fn worker_error_from_status(status: Status) -> WorkerErrorResponse {
#[derive(prost::Message)]
struct RpcStatus {
#[prost(int32, tag = "1")]
code: i32,
#[prost(string, tag = "2")]
message: String,
#[prost(message, repeated, tag = "3")]
details: Vec<prost_types::Any>,
}
use prost::Message;
if let Ok(details) = RpcStatus::decode(status.details())
&& let Some(worker) = details
.details
.iter()
.find(|detail| detail.type_url.ends_with("/dex.WorkerErrorResponse"))
&& let Ok(mut error) = WorkerErrorResponse::decode(worker.value.as_slice())
{
error.grpc_status = status.code() as i32;
return error;
}
WorkerErrorResponse {
detail: status.message().to_string(),
grpc_status: status.code() as i32,
..Default::default()
}
}
fn task_deadline(timestamp: &prost_types::Timestamp) -> Instant {
let current = Instant::now();
if timestamp.seconds < 0 || !(0..1_000_000_000).contains(×tamp.nanos) {
return current;
}
let Some(deadline) = SystemTime::UNIX_EPOCH
.checked_add(Duration::from_secs(timestamp.seconds as u64))
.and_then(|time| time.checked_add(Duration::from_nanos(timestamp.nanos as u64)))
else {
return current;
};
current
.checked_add(
deadline
.duration_since(SystemTime::now())
.unwrap_or_default(),
)
.unwrap_or(current)
}
fn is_permanent(code: Code) -> bool {
matches!(
code,
Code::Unauthenticated
| Code::PermissionDenied
| Code::InvalidArgument
| Code::Unimplemented
| Code::ResourceExhausted
)
}
fn generate_app_worker_id() -> String {
let identity = super::worker_options::WorkerEnvironmentConfig::from_environment().identity;
let mut cleaned = String::new();
for character in identity.to_ascii_lowercase().chars() {
if cleaned.len() >= 54 {
break;
}
if character.is_ascii_lowercase() || character.is_ascii_digit() || character == '-' {
cleaned.push(character);
} else if !cleaned.is_empty() {
cleaned.push('-');
}
}
let cleaned = cleaned.trim_matches('-');
let identity = if cleaned.is_empty() {
"worker"
} else {
cleaned
};
let random = uuid::Uuid::new_v4();
let mut bits = 0u64;
for byte in &random.as_bytes()[..5] {
bits = (bits << 8) | u64::from(*byte);
}
let alphabet = b"abcdefghijklmnopqrstuvwxyz234567";
let suffix: String = (0..8)
.rev()
.map(|index| alphabet[((bits >> (index * 5)) & 31) as usize] as char)
.collect();
format!("{identity}-{suffix}")
}
fn validate_app_worker_id(value: &str) -> SdkResult<()> {
if value.is_empty() || value.len() > 63 {
return Err(invalid("app Worker ID must contain 1-63 characters"));
}
if let Some(character) = value.chars().find(|character| {
!character.is_ascii_alphanumeric() && !matches!(character, '.' | '_' | '-')
}) {
return Err(invalid(format!(
"app Worker ID contains invalid character {character:?}"
)));
}
Ok(())
}
fn invalid(message: impl Into<String>) -> SdkError {
SdkError::FlowDefinition {
message: message.into(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::value_hydrator::ValueHydrator;
use crate::{
BlobCache, BlobCacheConfig, Context, FileCredential, Flow, HandlerError, HandlerResult,
Registry, Rpc, RpcList, RpcResult, Worker, WorkerOptions, WorkerTarget,
};
use dex_protocol::dex::flow_service_client::FlowServiceClient;
use dex_protocol::dex::worker_task_service_server::{
WorkerTaskService, WorkerTaskServiceServer,
};
use dex_protocol::dex::{
Context as ProtoContext, InvokeWaitForMethodRequest, InvokeWorkerRpcRequest,
PollWorkerTaskCancelTask, PollWorkerTaskPing, PollWorkerTaskReconnect,
PollWorkerTaskRegistered, PollWorkerTaskResponse,
};
use prost::Message;
use std::collections::VecDeque;
use std::sync::Condvar;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tokio::net::TcpListener;
use tokio::time::timeout;
use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
use tonic::transport::Server;
use tonic::{Request, Response};
struct ReserveCapacityFlow;
impl Flow for ReserveCapacityFlow {
type StartInput = ();
fn steps(&self) -> crate::StepList<'_, ()> {
crate::StepList::start(&ReserveCapacityStep)
}
}
struct ReserveCapacityStep;
impl crate::Step for ReserveCapacityStep {
type Input = ();
fn execute(&self, _: &mut Context, _: ()) -> HandlerResult<crate::StepDecision> {
Ok(crate::StepDecision::graceful_complete(()))
}
}
#[derive(Clone)]
struct RegistrationService {
opened: mpsc::Sender<mpsc::Sender<Result<PollWorkerTaskResponse, Status>>>,
}
#[tonic::async_trait]
impl WorkerTaskService for RegistrationService {
type PollWorkerTaskStream = ReceiverStream<Result<PollWorkerTaskResponse, Status>>;
async fn poll_worker_task(
&self,
request: Request<tonic::Streaming<PollWorkerTaskRequest>>,
) -> Result<Response<Self::PollWorkerTaskStream>, Status> {
let mut incoming = request.into_inner();
let frame = incoming
.message()
.await?
.ok_or_else(|| Status::invalid_argument("missing Register"))?;
assert!(matches!(frame.frame, Some(RequestFrame::Register(_))));
let (sender, receiver) = mpsc::channel(8);
self.opened
.send(sender)
.await
.map_err(|error| Status::internal(error.to_string()))?;
let _stream_reader = tokio::spawn(async move {
loop {
match incoming.message().await {
Ok(Some(_)) => {}
Ok(None) => break,
Err(error) => {
eprintln!("registration test stream closed: {error}");
break;
}
}
}
});
Ok(Response::new(ReceiverStream::new(receiver)))
}
}
async fn send_registered(sender: &mpsc::Sender<Result<PollWorkerTaskResponse, Status>>) {
sender
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Registered(
PollWorkerTaskRegistered::default(),
)),
}))
.await
.expect("send Registered");
}
#[tokio::test(flavor = "multi_thread")]
async fn registration_transport_handles_permanent_rejections_and_stop() {
for code in [
Code::InvalidArgument,
Code::ResourceExhausted,
Code::Unimplemented,
Code::PermissionDenied,
Code::Unauthenticated,
] {
for has_registered in [false, true] {
for should_stop in [false, true] {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind registration service");
let address = listener.local_addr().expect("service address");
let (opened, mut incoming) = mpsc::channel(8);
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(RegistrationService {
opened,
}))
.serve_with_incoming_shutdown(
TcpListenerStream::new(listener),
async move {
while !*server_shutdown.borrow()
&& server_shutdown.changed().await.is_ok()
{
}
},
),
);
let channel = outbound_task_channel(
Endpoint::from_shared(format!("http://{address}")).expect("endpoint"),
None,
);
let (cache, _directory) = test_cache();
let dispatcher = WorkerDispatcher::new(
Registry::new(),
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker = PullWorker::new(
&PullWorkerOptions::new().stop_on_permanent_stream_rejection(should_stop),
&dispatcher,
channel,
)
.expect("pull Worker");
let registration = Arc::clone(&worker.registration);
let (stop, shutdown) = watch::channel(false);
let running = tokio::spawn(async move { worker.run(shutdown).await });
let sender = timeout(Duration::from_secs(5), incoming.recv())
.await
.expect("stream opened")
.expect("response sender");
if has_registered {
send_registered(&sender).await;
timeout(Duration::from_secs(5), registration.wait())
.await
.expect("readiness deadline")
.expect("ready");
assert!(registration.snapshot().is_ready);
}
sender
.send(Err(Status::new(code, "permanent rejection")))
.await
.expect("send rejection");
if !has_registered {
let error = timeout(Duration::from_secs(5), registration.wait())
.await
.expect("rejection deadline")
.expect_err("rejected");
assert!(
error
.pull_registration_error()
.is_some_and(|rejection| rejection.status.code() == code)
);
}
timeout(Duration::from_secs(5), async {
while registration.snapshot().queues[0]
.last_permanent_error
.is_none()
{
tokio::task::yield_now().await;
}
})
.await
.expect("rejection recorded");
assert!(!registration.snapshot().is_ready);
if !should_stop {
assert!(!running.is_finished());
}
if !should_stop {
stop.send_replace(true);
}
let result = timeout(Duration::from_secs(5), running)
.await
.expect("Worker stopped")
.expect("join Worker");
if should_stop {
assert!(
result
.expect_err("fatal rejection")
.pull_registration_error()
.is_some_and(|rejection| rejection.status.code() == code)
);
} else {
result.expect("ordinary Stop");
}
stop.send_replace(true);
if has_registered {
registration.wait().await.expect("latched registration");
}
drop(sender);
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("server stopped")
.expect("join server")
.expect("server result");
}
}
}
}
#[tokio::test(flavor = "multi_thread")]
async fn registered_spare_stream_prevents_fatal_shutdown_and_recovers_readiness() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind service");
let address = listener.local_addr().expect("service address");
let (opened, mut incoming) = mpsc::channel(8);
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(RegistrationService { opened }))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok() {}
}),
);
let channel = outbound_task_channel(
Endpoint::from_shared(format!("http://{address}")).expect("endpoint"),
None,
);
let (cache, _directory) = test_cache();
let dispatcher = WorkerDispatcher::new(
Registry::new(),
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker = PullWorker::new(
&PullWorkerOptions::new()
.idle_streams_per_queue(2)
.stop_on_permanent_stream_rejection(true),
&dispatcher,
channel,
)
.expect("Worker");
let registration = Arc::clone(&worker.registration);
let (stop, shutdown) = watch::channel(false);
let running = tokio::spawn(async move { worker.run(shutdown).await });
let first = timeout(Duration::from_secs(5), incoming.recv())
.await
.expect("first stream")
.expect("first sender");
let second = timeout(Duration::from_secs(5), incoming.recv())
.await
.expect("second stream")
.expect("second sender");
send_registered(&first).await;
timeout(Duration::from_secs(5), registration.wait())
.await
.expect("readiness deadline")
.expect("ready");
second
.send(Err(Status::permission_denied("rejected")))
.await
.expect("reject stream");
timeout(Duration::from_secs(5), async {
while registration.snapshot().queues[0]
.last_permanent_error
.is_none()
{
tokio::task::yield_now().await;
}
})
.await
.expect("rejection recorded");
assert!(registration.snapshot().is_ready);
assert!(registration.permanent_error().is_none());
first
.send(Err(Status::unavailable("transient")))
.await
.expect("disconnect registered stream");
timeout(Duration::from_secs(5), async {
while registration.snapshot().queues[0].has_registered_stream {
tokio::task::yield_now().await;
}
})
.await
.expect("disconnect recorded");
assert!(!registration.snapshot().is_ready);
assert!(registration.permanent_error().is_none());
let replacement = timeout(Duration::from_secs(5), incoming.recv())
.await
.expect("replacement stream")
.expect("replacement sender");
send_registered(&replacement).await;
timeout(Duration::from_secs(5), async {
while !registration.snapshot().is_ready {
tokio::task::yield_now().await;
}
})
.await
.expect("readiness restored");
assert!(!running.is_finished());
assert!(
registration.snapshot().queues[0]
.last_permanent_error
.is_none()
);
stop.send_replace(true);
timeout(Duration::from_secs(5), running)
.await
.expect("Worker stopped")
.expect("join Worker")
.expect("ordinary shutdown");
drop((first, second, replacement));
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("server stopped")
.expect("join server")
.expect("server result");
}
#[tokio::test(flavor = "multi_thread")]
async fn permanent_rejection_updates_readiness_before_cancelled_handler_returns() {
for rejection_code in [Code::PermissionDenied, Code::Unauthenticated] {
assert_rejection_updates_readiness_before_cancelled_handler_returns(rejection_code)
.await;
}
}
async fn assert_rejection_updates_readiness_before_cancelled_handler_returns(
rejection_code: Code,
) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind service");
let address = listener.local_addr().expect("service address");
let (opened, mut incoming) = mpsc::channel(8);
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(RegistrationService { opened }))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok() {}
}),
);
let credential_directory = tempfile::tempdir().expect("credential directory");
let credential_path = credential_directory.path().join("access-token");
std::fs::write(&credential_path, "file-access-token").expect("write credential");
let credential: Arc<dyn Credential> =
Arc::new(FileCredential::new(&credential_path).expect("file credential"));
let channel = outbound_task_channel(
Endpoint::from_shared(format!("http://{address}")).expect("endpoint"),
Some(credential),
);
let (cache, _directory) = test_cache();
let started = Arc::new(Notify::new());
let release = Arc::new((Mutex::new(false), Condvar::new()));
let registry = Registry::new()
.register(DrainFlow {
started: Arc::clone(&started),
release: Arc::clone(&release),
wait_for_cancel: false,
result_size: None,
})
.expect("register Flow");
let dispatcher = WorkerDispatcher::new(
registry,
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker = PullWorker::new(
&PullWorkerOptions::new()
.stop_drain_timeout(Duration::from_millis(100))
.stop_on_permanent_stream_rejection(true),
&dispatcher,
channel,
)
.expect("Worker");
let registration = Arc::clone(&worker.registration);
let (stop, shutdown) = watch::channel(false);
let running = tokio::spawn(async move { worker.run(shutdown).await });
let first = timeout(Duration::from_secs(5), incoming.recv())
.await
.expect("first stream")
.expect("first sender");
let second = timeout(Duration::from_secs(5), incoming.recv())
.await
.expect("second stream")
.expect("second sender");
send_registered(&first).await;
send_registered(&second).await;
timeout(Duration::from_secs(5), registration.wait())
.await
.expect("readiness deadline")
.expect("ready");
first
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Task(PollWorkerTask {
task_id: "rejected-busy-task".to_string(),
deadline: None,
request: Some(TaskRequest::Rpc(InvokeWorkerRpcRequest {
context: Some(ProtoContext {
flow_id: "flow-1".to_string(),
..Default::default()
}),
flow_type: "DrainFlow".to_string(),
rpc_name: "get_drain_result".to_string(),
input: Some(crate::value_mapper::encode(&()).expect("encode input")),
..Default::default()
})),
})),
}))
.await
.expect("send RPC task");
timeout(Duration::from_secs(5), started.notified())
.await
.expect("handler started");
first
.send(Err(Status::new(rejection_code, "rejected while busy")))
.await
.expect("reject stream");
let recorded = timeout(Duration::from_secs(5), async {
while registration.snapshot().is_ready {
tokio::task::yield_now().await;
}
})
.await;
let is_waiting_for_handler = !running.is_finished();
let (lock, changed) = &*release;
*lock.lock().expect("handler gate lock") = true;
changed.notify_all();
let result = timeout(Duration::from_secs(5), running)
.await
.expect("Worker stopped")
.expect("join Worker");
stop.send_replace(true);
drop((first, second));
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("server stopped")
.expect("join server")
.expect("server result");
recorded.expect("rejection recorded before handler release");
assert!(is_waiting_for_handler);
assert!(
result
.expect_err("fatal rejection")
.pull_registration_error()
.is_some()
);
}
#[test]
fn registration_wait_reports_start_failure_and_non_pull_worker() {
let (cache, _directory) = test_cache();
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("wait runtime");
let worker = Worker::try_new(
Registry::new(),
Arc::clone(&cache),
WorkerOptions::new()
.server_address("127.0.0.1:1")
.attribute_index_sync_timeout(Duration::from_millis(100))
.pull(PullWorkerOptions::new()),
)
.expect("Worker");
let failure = worker.start().expect_err("startup failed");
let waiting = runtime
.block_on(worker.wait_pull_registered())
.expect_err("wait failed");
assert_eq!(failure.to_string(), waiting.to_string());
worker.stop();
let listener = Worker::try_new(Registry::new(), cache, WorkerOptions::new())
.expect("listening Worker");
assert!(runtime.block_on(listener.wait_pull_registered()).is_err());
assert!(!listener.pull_registration_state().is_ready);
assert!(listener.pull_registration_state().queues.is_empty());
listener.stop();
}
struct RegisteredFlow;
impl Flow for RegisteredFlow {
type StartInput = ();
}
fn test_cache() -> (Arc<BlobCache>, tempfile::TempDir) {
let directory = tempfile::tempdir().expect("cache directory");
let cache = Arc::new(
BlobCache::open(
BlobCacheConfig::new(directory.path(), 1024 * 1024, 0).expect("cache config"),
)
.expect("open cache"),
);
(cache, directory)
}
#[test]
fn validates_pull_options_and_target_conflicts() {
let (cache, _directory) = test_cache();
let registry = Registry::new()
.register(RegisteredFlow)
.expect("register Flow");
for options in [
WorkerOptions::new().pull(PullWorkerOptions::new().max_concurrent_tasks(0)),
WorkerOptions::new().pull(
PullWorkerOptions::new()
.max_concurrent_tasks(2)
.reserved_rpc_tasks(2),
),
WorkerOptions::new().pull(
PullWorkerOptions::new()
.max_concurrent_tasks(2)
.reserved_rpc_tasks(3),
),
WorkerOptions::new().pull(PullWorkerOptions::new().idle_streams_per_queue(0)),
WorkerOptions::new().pull(PullWorkerOptions::new().stop_drain_timeout(Duration::ZERO)),
WorkerOptions::new().pull(PullWorkerOptions::new().app_worker_id("bad/id")),
WorkerOptions::new().pull(PullWorkerOptions::new().app_worker_id("RegisteredFlow")),
WorkerOptions::new()
.pull(PullWorkerOptions::new())
.bind_address("127.0.0.1:8803"),
WorkerOptions::new()
.pull(PullWorkerOptions::new())
.worker_target(WorkerTarget::pull()),
WorkerOptions::new().worker_target(WorkerTarget::pull()),
] {
assert!(matches!(
Worker::try_new(registry.clone(), Arc::clone(&cache), options),
Err(SdkError::FlowDefinition { .. })
));
}
assert!(
Worker::try_new(
registry,
cache,
WorkerOptions::new().pull(PullWorkerOptions::new())
)
.is_ok()
);
assert!(WorkerTarget::pull().headless(true).validate().is_err());
}
struct ReserveWarningSubscriber(Arc<AtomicU64>);
impl tracing::Subscriber for ReserveWarningSubscriber {
fn enabled(&self, metadata: &tracing::Metadata<'_>) -> bool {
*metadata.level() == tracing::Level::WARN
}
fn new_span(&self, _: &tracing::span::Attributes<'_>) -> tracing::span::Id {
tracing::span::Id::from_u64(1)
}
fn record(&self, _: &tracing::span::Id, _: &tracing::span::Record<'_>) {}
fn record_follows_from(&self, _: &tracing::span::Id, _: &tracing::span::Id) {}
fn event(&self, event: &tracing::Event<'_>) {
if *event.metadata().level() == tracing::Level::WARN {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
fn enter(&self, _: &tracing::span::Id) {}
fn exit(&self, _: &tracing::span::Id) {}
}
#[tokio::test]
async fn defaults_reserve_one_and_warns_for_single_capacity() {
let warnings = Arc::new(AtomicU64::new(0));
let _subscriber =
tracing::subscriber::set_default(ReserveWarningSubscriber(Arc::clone(&warnings)));
for capacity in [1, 4] {
let channel = outbound_task_channel(Endpoint::from_static("http://127.0.0.1:1"), None);
let (cache, _directory) = test_cache();
let dispatcher = WorkerDispatcher::new(
Registry::new(),
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker = PullWorker::new(
&PullWorkerOptions::new().max_concurrent_tasks(capacity),
&dispatcher,
channel,
)
.expect("Worker");
assert_eq!(usize::from(capacity >= 2), worker.reserved_rpc_tasks);
assert_eq!(0, worker.reserved_rpc_declines());
let (_stop, shutdown) = watch::channel(true);
worker.run(shutdown).await.expect("stop Worker");
}
assert_eq!(1, warnings.load(Ordering::Relaxed));
}
#[test]
fn generates_distinct_valid_app_worker_ids() {
let first = generate_app_worker_id();
let second = generate_app_worker_id();
assert_ne!(first, second);
validate_app_worker_id(&first).expect("valid first ID");
assert!(first.len() <= 63);
let suffix = first.rsplit('-').next().expect("ID suffix");
assert_eq!(8, suffix.len());
assert!(
suffix
.bytes()
.all(|byte| byte.is_ascii_lowercase() || (b'2'..=b'7').contains(&byte))
);
assert!(validate_app_worker_id("").is_err());
assert!(validate_app_worker_id(&"a".repeat(64)).is_err());
assert!(validate_app_worker_id("has space").is_err());
}
fn reserved_rpc_shared(reserved_rpc_tasks: usize) -> PullShared {
PullShared {
registration: Arc::default(),
state: Mutex::new(PullState {
streams: (1..=20)
.map(|stream_id| {
(
stream_id,
StreamSlot {
queue: "queue".to_string(),
is_busy: false,
is_registered: true,
is_retiring: false,
is_draining: false,
retirement: Arc::new(Notify::new()),
},
)
})
.collect(),
busy_tasks: 0,
next_stream_id: 21,
is_stopping: false,
logged_unsupported: false,
}),
notify: Notify::new(),
max_concurrent_tasks: 2,
reserved_rpc_tasks,
reserved_rpc_declines: Arc::new(AtomicU64::new(0)),
idle_streams_per_queue: 1,
}
}
#[test]
fn reserves_rpc_capacity_across_all_busy_handlers() {
let step = PollWorkerTask {
request: Some(TaskRequest::Execute(Default::default())),
..Default::default()
};
let wait_for = PollWorkerTask {
request: Some(TaskRequest::WaitFor(Default::default())),
..Default::default()
};
let rpc = PollWorkerTask {
request: Some(TaskRequest::Rpc(Default::default())),
..Default::default()
};
for reserve in [0, 1] {
for first in [&step, &wait_for, &rpc] {
let shared = reserved_rpc_shared(reserve);
assert!(shared.try_accept(1, Some(first)));
if reserve == 1 {
assert!(!shared.try_accept(2, Some(&step)));
assert!(!shared.try_accept(2, Some(&wait_for)));
assert!(shared.try_accept(2, Some(&rpc)));
} else {
assert!(shared.try_accept(2, Some(&step)));
}
assert!(!shared.try_accept(3, Some(&step)));
assert!(!shared.try_accept(3, Some(&rpc)));
let state = shared.state.lock().expect("state lock");
assert_eq!(2, state.busy_tasks);
assert_eq!(
if reserve == 1 { 2 } else { 0 },
shared.reserved_rpc_declines.load(Ordering::Relaxed)
);
assert!(!state.streams[&3].is_retiring);
}
}
}
#[test]
fn concurrent_step_admission_preserves_rpc_reserve() {
let shared = Arc::new(reserved_rpc_shared(1));
let barrier = Arc::new(std::sync::Barrier::new(20));
let contenders: Vec<_> = (1..=20)
.map(|stream_id| {
let shared = Arc::clone(&shared);
let barrier = Arc::clone(&barrier);
std::thread::spawn(move || {
let task = PollWorkerTask {
request: Some(TaskRequest::Execute(Default::default())),
..Default::default()
};
barrier.wait();
shared.try_accept(stream_id, Some(&task))
})
})
.collect();
assert_eq!(
1,
contenders
.into_iter()
.map(|contender| usize::from(contender.join().expect("admission thread")))
.sum::<usize>()
);
assert_eq!(1, shared.state.lock().expect("state lock").busy_tasks);
}
#[test]
fn rejects_an_assignment_when_global_capacity_is_full() {
let shared = PullShared {
registration: Arc::default(),
state: Mutex::new(PullState {
streams: HashMap::from([
(
1,
StreamSlot {
queue: "queue".to_string(),
is_busy: false,
is_registered: false,
is_retiring: false,
is_draining: false,
retirement: Arc::new(Notify::new()),
},
),
(
2,
StreamSlot {
queue: "queue".to_string(),
is_busy: false,
is_registered: false,
is_retiring: false,
is_draining: false,
retirement: Arc::new(Notify::new()),
},
),
]),
busy_tasks: 0,
next_stream_id: 3,
is_stopping: false,
logged_unsupported: false,
}),
notify: Notify::new(),
max_concurrent_tasks: 1,
reserved_rpc_tasks: 0,
reserved_rpc_declines: Arc::new(AtomicU64::new(0)),
idle_streams_per_queue: 1,
};
assert!(shared.try_accept(1, None));
assert!(!shared.try_accept(2, None));
assert_eq!(PollWorkerTaskRejectReason::NoCapacity as i32, 1);
shared.release(1);
assert!(shared.try_accept(2, None));
}
#[test]
fn draining_stream_is_not_an_idle_replacement() {
let shared = PullShared {
registration: Arc::default(),
state: Mutex::new(PullState {
streams: HashMap::from([
(
1,
StreamSlot {
queue: "queue".to_string(),
is_busy: true,
is_registered: true,
is_retiring: false,
is_draining: false,
retirement: Arc::new(Notify::new()),
},
),
(
2,
StreamSlot {
queue: "queue".to_string(),
is_busy: false,
is_registered: true,
is_retiring: false,
is_draining: false,
retirement: Arc::new(Notify::new()),
},
),
]),
busy_tasks: 1,
next_stream_id: 3,
is_stopping: false,
logged_unsupported: false,
}),
notify: Notify::new(),
max_concurrent_tasks: 1,
reserved_rpc_tasks: 0,
reserved_rpc_declines: Arc::new(AtomicU64::new(0)),
idle_streams_per_queue: 1,
};
shared.mark_draining(2);
assert!(!shared.release(1));
assert!(!shared.is_retiring(1));
assert!(shared.try_accept(1, None));
assert!(!shared.try_accept(2, None));
}
#[test]
fn maps_push_worker_failure_fields_into_pull_frame() {
let status = crate::worker::worker_status(HandlerError::retry_after(
7,
"PaymentDeclined",
"card declined",
));
let error = worker_error_from_status(status);
assert_eq!(Code::Unknown as i32, error.grpc_status);
assert_eq!("card declined", error.detail);
assert_eq!("PaymentDeclined", error.error_type);
assert_eq!(7, error.retry_after_seconds);
assert!(
error
.stack_trace
.contains("maps_push_worker_failure_fields")
);
let cancelled_reason = AtomicI32::new(Code::Cancelled as i32);
let cancelled_error = with_cancel_status(error.clone(), &cancelled_reason);
assert_eq!(Code::Cancelled as i32, cancelled_error.grpc_status);
assert_eq!(error.detail, cancelled_error.detail);
assert_eq!(error.error_type, cancelled_error.error_type);
assert_eq!(error.stack_trace, cancelled_error.stack_trace);
assert_eq!(
error.retry_after_seconds,
cancelled_error.retry_after_seconds
);
let deadline_reason = AtomicI32::new(Code::DeadlineExceeded as i32);
assert_eq!(
Code::DeadlineExceeded as i32,
with_cancel_status(error, &deadline_reason).grpc_status
);
let cancelled = worker_error_from_status(Status::cancelled("cancelled"));
assert_eq!(Code::Cancelled as i32, cancelled.grpc_status);
}
#[test]
fn expired_task_deadline_is_immediate() {
let deadline = task_deadline(&prost_types::Timestamp {
seconds: 0,
nanos: 0,
});
assert!(deadline <= Instant::now());
let malformed = task_deadline(&prost_types::Timestamp {
seconds: i64::MAX,
nanos: -1,
});
assert!(malformed <= Instant::now());
}
#[tokio::test]
async fn cancellation_waits_for_the_handler_to_return() {
let cancellation = InvocationCancellation::new();
let waiting = cancellation.clone();
let handler = tokio::spawn(async move { waiting.cancelled().await });
cancellation.cancel();
timeout(Duration::from_secs(1), handler)
.await
.expect("handler returned")
.expect("join handler");
}
#[tokio::test]
async fn cancelled_task_does_not_block_on_a_full_outbound_stream() {
let (sender, _receiver) = mpsc::channel(1);
sender
.try_send(PollWorkerTaskRequest {
frame: Some(RequestFrame::Pong(PollWorkerTaskPong {})),
})
.expect("fill outbound stream");
let cancellation = InvocationCancellation::new();
cancellation.cancel();
let result = timeout(
Duration::from_secs(1),
send_task_frame(
&sender,
"task-1",
Output::WorkerError(WorkerErrorResponse::default()),
&cancellation,
),
)
.await
.expect("send stopped");
assert_eq!(Code::Cancelled, result.expect_err("cancelled send").code());
}
#[derive(Clone)]
struct ReconnectService {
registrations: Arc<AtomicUsize>,
pongs: Arc<AtomicUsize>,
notify: Arc<Notify>,
headers: Arc<Mutex<Vec<Option<String>>>>,
credential_token: Arc<Mutex<String>>,
}
struct RotatingCredential(Arc<Mutex<String>>);
impl Credential for RotatingCredential {
fn token(&self) -> Result<String, crate::CredentialError> {
Ok(self.0.lock().expect("credential token lock").clone())
}
}
#[tonic::async_trait]
impl WorkerTaskService for ReconnectService {
type PollWorkerTaskStream = ReceiverStream<Result<PollWorkerTaskResponse, Status>>;
async fn poll_worker_task(
&self,
request: Request<tonic::Streaming<PollWorkerTaskRequest>>,
) -> Result<Response<Self::PollWorkerTaskStream>, Status> {
self.headers.lock().expect("headers lock").push(
request
.metadata()
.get("authorization")
.map(|value| value.to_str().expect("ASCII bearer header").to_string()),
);
let mut incoming = request.into_inner();
let first = incoming
.message()
.await?
.ok_or_else(|| Status::invalid_argument("missing Register"))?;
let Some(RequestFrame::Register(register)) = first.frame else {
return Err(Status::invalid_argument("Register must be first"));
};
assert_eq!("rust", register.sdk_name);
assert_eq!(register.queue, register.app_worker_id);
if self.registrations.fetch_add(1, Ordering::SeqCst) == 0 {
*self.credential_token.lock().expect("credential token lock") =
"rotated".to_string();
}
self.notify.notify_one();
let (sender, receiver) = mpsc::channel(4);
let pongs = Arc::clone(&self.pongs);
let _stream_reader = tokio::spawn(async move {
let registered = PollWorkerTaskResponse {
frame: Some(ResponseFrame::Registered(PollWorkerTaskRegistered {
supports_unknown_name_rejection: false,
})),
};
if sender.send(Ok(registered)).await.is_err() {
return;
}
let ping = PollWorkerTaskResponse {
frame: Some(ResponseFrame::Ping(PollWorkerTaskPing {})),
};
if sender.send(Ok(ping)).await.is_err() {
return;
}
match incoming.message().await {
Ok(Some(PollWorkerTaskRequest {
frame: Some(RequestFrame::Pong(_)),
})) => {
pongs.fetch_add(1, Ordering::SeqCst);
}
Ok(other) => {
eprintln!("expected Pong frame, received {other:?}");
return;
}
Err(error) => {
eprintln!("read Pong frame: {error}");
return;
}
}
let reconnect = PollWorkerTaskResponse {
frame: Some(ResponseFrame::Reconnect(PollWorkerTaskReconnect {
reason: "test".to_string(),
})),
};
if let Err(error) = sender.send(Ok(reconnect)).await {
eprintln!("send reconnect frame: {error}");
}
});
Ok(Response::new(ReceiverStream::new(receiver)))
}
}
#[tokio::test(flavor = "multi_thread")]
async fn reconnects_after_ping_and_drains_idle_stream_on_stop() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind fake task service");
let address = listener.local_addr().expect("fake service address");
let service = ReconnectService {
registrations: Arc::new(AtomicUsize::new(0)),
pongs: Arc::new(AtomicUsize::new(0)),
notify: Arc::new(Notify::new()),
headers: Arc::new(Mutex::new(Vec::new())),
credential_token: Arc::new(Mutex::new("initial".to_string())),
};
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(service.clone()))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok() {}
}),
);
let channel = Endpoint::from_shared(format!("http://{address}"))
.expect("endpoint")
.connect_lazy();
let channel = CredentialChannel::new(
channel,
CredentialInterceptor::new(Some(Arc::new(RotatingCredential(Arc::clone(
&service.credential_token,
))))),
);
let (cache, _directory) = test_cache();
let dispatcher = WorkerDispatcher::new(
Registry::new(),
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker =
PullWorker::new(&PullWorkerOptions::new(), &dispatcher, channel).expect("pull Worker");
let (stop, shutdown) = watch::channel(false);
let running = tokio::spawn(async move { worker.run(shutdown).await });
timeout(Duration::from_secs(5), async {
while service.registrations.load(Ordering::SeqCst) < 2 {
service.notify.notified().await;
}
})
.await
.expect("second registration");
assert!(service.pongs.load(Ordering::SeqCst) >= 1);
assert_eq!(
&service.headers.lock().expect("headers lock")[..2],
&[
Some("Bearer initial".to_string()),
Some("Bearer rotated".to_string())
]
);
stop.send_replace(true);
timeout(Duration::from_secs(5), running)
.await
.expect("pull Worker stopped")
.expect("join Worker")
.expect("Worker result");
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("fake server stopped")
.expect("join server")
.expect("server result");
}
#[derive(Clone)]
struct RejectService {
rejected: mpsc::Sender<PollWorkerTaskReject>,
}
#[tonic::async_trait]
impl WorkerTaskService for RejectService {
type PollWorkerTaskStream = ReceiverStream<Result<PollWorkerTaskResponse, Status>>;
async fn poll_worker_task(
&self,
request: Request<tonic::Streaming<PollWorkerTaskRequest>>,
) -> Result<Response<Self::PollWorkerTaskStream>, Status> {
let mut incoming = request.into_inner();
let first = incoming
.message()
.await?
.ok_or_else(|| Status::invalid_argument("missing Register"))?;
if !matches!(first.frame, Some(RequestFrame::Register(_))) {
return Err(Status::invalid_argument("Register must be first"));
}
let (sender, receiver) = mpsc::channel(2);
let rejected = self.rejected.clone();
let _stream_reader = tokio::spawn(async move {
let registered = PollWorkerTaskResponse {
frame: Some(ResponseFrame::Registered(PollWorkerTaskRegistered {
supports_unknown_name_rejection: false,
})),
};
if sender.send(Ok(registered)).await.is_err() {
return;
}
let assigned = PollWorkerTaskResponse {
frame: Some(ResponseFrame::Task(PollWorkerTask {
task_id: "assigned-1".to_string(),
deadline: None,
request: Some(TaskRequest::Execute(
dex_protocol::dex::InvokeExecuteMethodRequest {
flow_type: "ReserveCapacityFlow".to_string(),
step_type: "ReserveCapacityStep".to_string(),
..Default::default()
},
)),
})),
};
if sender.send(Ok(assigned)).await.is_err() {
return;
}
match incoming.message().await {
Ok(Some(PollWorkerTaskRequest {
frame: Some(RequestFrame::Reject(frame)),
})) => {
if let Err(error) = rejected.send(frame).await {
eprintln!("record Reject frame: {error}");
}
}
Ok(other) => eprintln!("expected Reject frame, received {other:?}"),
Err(error) => eprintln!("read Reject frame: {error}"),
}
});
Ok(Response::new(ReceiverStream::new(receiver)))
}
}
#[tokio::test(flavor = "multi_thread")]
async fn sends_no_capacity_reject_frame() {
for reserve in [0, 1] {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind fake task service");
let address = listener.local_addr().expect("fake service address");
let (rejected, mut rejected_frames) = mpsc::channel(1);
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(RejectService { rejected }))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok()
{
}
}),
);
let channel = Endpoint::from_shared(format!("http://{address}"))
.expect("endpoint")
.connect_lazy();
let channel = CredentialChannel::new(channel, CredentialInterceptor::new(None));
let (cache, _directory) = test_cache();
let dispatcher = WorkerDispatcher::new(
Registry::new()
.register(ReserveCapacityFlow)
.expect("register Flow"),
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let shared = Arc::new(PullShared {
registration: Arc::default(),
state: Mutex::new(PullState {
streams: HashMap::from([(
1,
StreamSlot {
queue: "worker-1".to_string(),
is_busy: false,
is_registered: false,
is_retiring: false,
is_draining: false,
retirement: Arc::new(Notify::new()),
},
)]),
busy_tasks: 1,
next_stream_id: 2,
is_stopping: false,
logged_unsupported: false,
}),
notify: Notify::new(),
max_concurrent_tasks: 1 + reserve,
reserved_rpc_tasks: reserve,
reserved_rpc_declines: Arc::new(AtomicU64::new(0)),
idle_streams_per_queue: 1,
});
let (_stop, stop_receiver) = watch::channel(StopSignal {
phase: 0,
deadline: None,
});
let mut stream = PullStream {
client: WorkerTaskServiceClient::new(channel),
credential: None,
dispatcher,
app_worker_id: "worker-1".to_string(),
flow_types: Vec::new(),
queue: "worker-1".to_string(),
stream_id: 1,
shared,
stop: stop_receiver,
stop_drain_timeout: Duration::from_secs(30),
is_registered: false,
has_recorded_result: false,
retirement: Arc::new(Notify::new()),
};
let (was_registered, _, result) = timeout(Duration::from_secs(5), stream.poll())
.await
.expect("poll returned");
assert!(was_registered);
result.expect("poll result");
let reject = rejected_frames.recv().await.expect("Reject frame");
assert_eq!("assigned-1", reject.task_id);
assert_eq!(PollWorkerTaskRejectReason::NoCapacity as i32, reject.reason);
assert_eq!(
reserve as u64,
stream.shared.reserved_rpc_declines.load(Ordering::Relaxed)
);
assert!(!stream.shared.state.lock().expect("state lock").streams[&1].is_retiring);
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("fake server stopped")
.expect("join server")
.expect("server result");
}
}
#[derive(Clone)]
struct UnknownTaskService {
capabilities: Arc<Mutex<VecDeque<bool>>>,
received: mpsc::Sender<PollWorkerTaskRequest>,
}
#[tonic::async_trait]
impl WorkerTaskService for UnknownTaskService {
type PollWorkerTaskStream = ReceiverStream<Result<PollWorkerTaskResponse, Status>>;
async fn poll_worker_task(
&self,
request: Request<tonic::Streaming<PollWorkerTaskRequest>>,
) -> Result<Response<Self::PollWorkerTaskStream>, Status> {
let mut incoming = request.into_inner();
let register = incoming
.message()
.await?
.ok_or_else(|| Status::invalid_argument("missing Register"))?;
if !matches!(register.frame, Some(RequestFrame::Register(_))) {
return Err(Status::invalid_argument("Register must be first"));
}
let supports_unknown_name_rejection = self
.capabilities
.lock()
.expect("capabilities lock")
.pop_front()
.ok_or_else(|| Status::internal("unexpected registration"))?;
let (sender, receiver) = mpsc::channel(4);
sender
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Registered(PollWorkerTaskRegistered {
supports_unknown_name_rejection,
})),
}))
.await
.map_err(|_| Status::internal("registration response closed"))?;
send_unknown_task(&sender, "unknown-step", true).await?;
let received = self.received.clone();
let _stream_reader = tokio::spawn(async move {
let mut completed_tasks = 0;
while let Ok(Some(frame)) = incoming.message().await {
if matches!(
frame.frame,
Some(RequestFrame::TaskFrame(_)) | Some(RequestFrame::Reject(_))
) {
if received.send(frame).await.is_err() {
return;
}
completed_tasks += 1;
if completed_tasks == 1 {
if send_unknown_task(&sender, "unknown-rpc", false)
.await
.is_err()
{
return;
}
} else if sender
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Ping(PollWorkerTaskPing {})),
}))
.await
.is_err()
{
return;
}
} else if matches!(frame.frame, Some(RequestFrame::Pong(_))) {
received
.send(frame)
.await
.expect("the test retains the Pong receiver until the frame arrives");
return;
}
}
});
Ok(Response::new(ReceiverStream::new(receiver)))
}
}
async fn send_unknown_task(
sender: &mpsc::Sender<Result<PollWorkerTaskResponse, Status>>,
task_id: &str,
is_step: bool,
) -> Result<(), Status> {
let value = crate::value_mapper::encode(&"unknown".to_string())
.map_err(|error| Status::internal(error.to_string()))?;
let context = Some(ProtoContext {
flow_id: "flow-1".to_string(),
..Default::default()
});
let request = if is_step {
TaskRequest::WaitFor(InvokeWaitForMethodRequest {
context,
flow_type: "TestFlow".to_string(),
step_type: "MissingStep".to_string(),
step_input: Some(value),
..Default::default()
})
} else {
TaskRequest::Rpc(InvokeWorkerRpcRequest {
context,
flow_type: "TestFlow".to_string(),
rpc_name: "MissingRPC".to_string(),
input: Some(value),
..Default::default()
})
};
sender
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Task(PollWorkerTask {
task_id: task_id.to_string(),
deadline: None,
request: Some(request),
})),
}))
.await
.map_err(|_| Status::internal("task response closed"))
}
#[tokio::test(flavor = "multi_thread")]
async fn unknown_task_rejection_follows_registered_capability_on_reconnect() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind fake task service");
let address = listener.local_addr().expect("fake service address");
let (received, mut frames) = mpsc::channel(9);
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(UnknownTaskService {
capabilities: Arc::new(Mutex::new(VecDeque::from([false, true, false]))),
received,
}))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok() {}
}),
);
let channel = CredentialChannel::new(
Endpoint::from_shared(format!("http://{address}"))
.expect("endpoint")
.connect_lazy(),
CredentialInterceptor::new(None),
);
let (cache, _directory) = test_cache();
let dispatcher = WorkerDispatcher::new(
Registry::new(),
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let shared = Arc::new(PullShared {
registration: Arc::default(),
state: Mutex::new(PullState {
streams: HashMap::from([(
1,
StreamSlot {
queue: "worker".to_string(),
is_busy: false,
is_registered: false,
is_retiring: false,
is_draining: false,
retirement: Arc::new(Notify::new()),
},
)]),
busy_tasks: 0,
next_stream_id: 2,
is_stopping: false,
logged_unsupported: false,
}),
notify: Notify::new(),
max_concurrent_tasks: 1,
reserved_rpc_tasks: 0,
reserved_rpc_declines: Arc::new(AtomicU64::new(0)),
idle_streams_per_queue: 1,
});
let (_stop, stop_receiver) = watch::channel(StopSignal {
phase: 0,
deadline: None,
});
let mut stream = PullStream {
client: WorkerTaskServiceClient::new(channel),
credential: None,
dispatcher,
app_worker_id: "worker".to_string(),
flow_types: Vec::new(),
queue: "worker".to_string(),
stream_id: 1,
shared,
stop: stop_receiver,
stop_drain_timeout: Duration::from_secs(30),
is_registered: false,
has_recorded_result: false,
retirement: Arc::new(Notify::new()),
};
for _ in 0..3 {
let (was_registered, _, result) = timeout(Duration::from_secs(5), stream.poll())
.await
.expect("poll returned");
assert!(was_registered);
result.expect("poll result");
}
let mut received = Vec::new();
for _ in 0..9 {
received.push(
timeout(Duration::from_secs(5), frames.recv())
.await
.expect("Worker frame timeout")
.expect("Worker frame channel"),
);
}
assert_legacy_unknown_frames(&received[0..3]);
assert_unknown_rejections(&received[3..6]);
assert_legacy_unknown_frames(&received[6..9]);
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("fake server stopped")
.expect("join server")
.expect("server result");
}
fn assert_legacy_unknown_frames(frames: &[PollWorkerTaskRequest]) {
for frame in &frames[0..2] {
let Some(RequestFrame::TaskFrame(task)) = frame.frame.as_ref() else {
panic!("expected legacy Worker error frame");
};
assert!(matches!(task.output, Some(Output::WorkerError(_))));
}
assert!(matches!(frames[2].frame, Some(RequestFrame::Pong(_))));
}
fn assert_unknown_rejections(frames: &[PollWorkerTaskRequest]) {
let Some(RequestFrame::Reject(step)) = frames[0].frame.as_ref() else {
panic!("expected unknown Step rejection");
};
assert_eq!(PollWorkerTaskRejectReason::UnknownStep as i32, step.reason);
let Some(RequestFrame::Reject(rpc)) = frames[1].frame.as_ref() else {
panic!("expected unknown RPC rejection");
};
assert_eq!(PollWorkerTaskRejectReason::UnknownRpc as i32, rpc.reason);
assert!(matches!(frames[2].frame, Some(RequestFrame::Pong(_))));
}
const GET_DRAIN_RESULT: Rpc<(), String> = Rpc::new("get_drain_result");
struct DrainFlow {
started: Arc<Notify>,
release: Arc<(Mutex<bool>, Condvar)>,
wait_for_cancel: bool,
result_size: Option<usize>,
}
impl DrainFlow {
fn get_drain_result(&self, context: &mut Context) -> HandlerResult<RpcResult<String>> {
self.started.notify_one();
if self.wait_for_cancel {
context.wait_for_cancellation();
return Err(HandlerError::new("Cancelled", "handler was cancelled"));
}
let (lock, changed) = &*self.release;
let mut is_released = lock.lock().expect("drain gate lock");
while !*is_released {
is_released = changed.wait(is_released).expect("drain gate lock");
}
Ok(RpcResult::new(match self.result_size {
Some(result_size) => "x".repeat(result_size),
None => "done".to_string(),
}))
}
}
impl Flow for DrainFlow {
type StartInput = ();
fn rpcs(&self) -> RpcList<Self> {
RpcList::new().function_without_input(GET_DRAIN_RESULT, Self::get_drain_result)
}
}
struct BurstRegistration {
queue: String,
responses: mpsc::Sender<Result<PollWorkerTaskResponse, Status>>,
is_busy: Arc<AtomicBool>,
}
#[derive(Clone)]
struct BurstService {
registrations: mpsc::Sender<BurstRegistration>,
active_streams: Arc<AtomicUsize>,
closed_while_busy: Arc<AtomicBool>,
changed: Arc<Notify>,
}
#[tonic::async_trait]
impl WorkerTaskService for BurstService {
type PollWorkerTaskStream = ReceiverStream<Result<PollWorkerTaskResponse, Status>>;
async fn poll_worker_task(
&self,
request: Request<tonic::Streaming<PollWorkerTaskRequest>>,
) -> Result<Response<Self::PollWorkerTaskStream>, Status> {
let mut incoming = request.into_inner();
let first = incoming
.message()
.await?
.ok_or_else(|| Status::invalid_argument("missing Register"))?;
let Some(RequestFrame::Register(register)) = first.frame else {
return Err(Status::invalid_argument("Register must be first"));
};
let (sender, receiver) = mpsc::channel(8);
sender
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Registered(PollWorkerTaskRegistered {
supports_unknown_name_rejection: false,
})),
}))
.await
.map_err(|_| Status::internal("registered response closed"))?;
let is_busy = Arc::new(AtomicBool::new(false));
self.registrations
.send(BurstRegistration {
queue: register.queue,
responses: sender.clone(),
is_busy: Arc::clone(&is_busy),
})
.await
.map_err(|_| Status::internal("registration receiver closed"))?;
self.active_streams.fetch_add(1, Ordering::AcqRel);
self.changed.notify_one();
let service = self.clone();
let _stream_reader = tokio::spawn(async move {
loop {
match incoming.message().await {
Ok(Some(PollWorkerTaskRequest {
frame: Some(RequestFrame::TaskFrame(_)),
})) => {
is_busy.store(false, Ordering::Release);
}
Ok(Some(_)) => {}
Ok(None) => break,
Err(error) => {
eprintln!("burst stream disconnected: {error}");
break;
}
}
}
if is_busy.load(Ordering::Acquire) {
service.closed_while_busy.store(true, Ordering::Release);
}
service.active_streams.fetch_sub(1, Ordering::AcqRel);
service.changed.notify_one();
drop(sender);
});
Ok(Response::new(ReceiverStream::new(receiver)))
}
}
#[tokio::test(flavor = "multi_thread")]
async fn surplus_streams_shrink_after_burst_without_losing_readiness() {
for idle_streams in [1, 2] {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind fake task service");
let address = listener.local_addr().expect("fake service address");
let (registrations, mut received) = mpsc::channel(16);
let service = BurstService {
registrations,
active_streams: Arc::new(AtomicUsize::new(0)),
closed_while_busy: Arc::new(AtomicBool::new(false)),
changed: Arc::new(Notify::new()),
};
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(service.clone()))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok()
{
}
}),
);
let channel = CredentialChannel::new(
Endpoint::from_shared(format!("http://{address}"))
.expect("endpoint")
.connect_lazy(),
CredentialInterceptor::new(None),
);
let (cache, _directory) = test_cache();
let started = Arc::new(Notify::new());
let release = Arc::new((Mutex::new(false), Condvar::new()));
let registry = Registry::new()
.register(DrainFlow {
started: Arc::clone(&started),
release: Arc::clone(&release),
wait_for_cancel: false,
result_size: None,
})
.expect("register Drain Flow");
let dispatcher = WorkerDispatcher::new(
registry,
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker = PullWorker::new(
&PullWorkerOptions::new()
.app_worker_id("burst-worker")
.max_concurrent_tasks(3)
.reserved_rpc_tasks(0)
.idle_streams_per_queue(idle_streams),
&dispatcher,
channel,
)
.expect("pull Worker");
let registration = Arc::clone(&worker.registration);
let (stop, shutdown) = watch::channel(false);
let running = tokio::spawn(async move { worker.run(shutdown).await });
timeout(Duration::from_secs(5), registration.wait())
.await
.expect("registration deadline")
.expect("registered");
for task_number in 0..3 {
let stream = timeout(Duration::from_secs(5), async {
loop {
let stream = received.recv().await.expect("registration");
if stream.queue == "DrainFlow" {
break stream;
}
}
})
.await
.expect("task stream registered");
stream.is_busy.store(true, Ordering::Release);
stream
.responses
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Task(PollWorkerTask {
task_id: format!("burst-{task_number}"),
deadline: None,
request: Some(TaskRequest::Rpc(InvokeWorkerRpcRequest {
context: Some(ProtoContext {
flow_id: "burst-flow".to_string(),
run_id: "burst-run".to_string(),
..Default::default()
}),
flow_type: "DrainFlow".to_string(),
rpc_name: "get_drain_result".to_string(),
input: Some(
crate::value_mapper::encode(&()).expect("encode input"),
),
..Default::default()
})),
})),
}))
.await
.expect("send task");
timeout(Duration::from_secs(5), started.notified())
.await
.expect("RPC started");
}
await_burst_stream_count(&service, ®istration, 3 + 2 * idle_streams).await;
assert!(!service.closed_while_busy.load(Ordering::Acquire));
loop {
match received.try_recv() {
Ok(pending) => drop(pending),
Err(mpsc::error::TryRecvError::Empty) => break,
Err(mpsc::error::TryRecvError::Disconnected) => {
panic!("registration sender closed")
}
}
}
let (lock, changed) = &*release;
*lock.lock().expect("burst gate lock") = true;
changed.notify_all();
await_burst_stream_count(&service, ®istration, 2 * idle_streams).await;
assert!(!service.closed_while_busy.load(Ordering::Acquire));
assert!(registration.snapshot().is_ready);
registration
.wait()
.await
.expect("readiness remains latched");
stop.send_replace(true);
timeout(Duration::from_secs(5), running)
.await
.expect("Worker stopped")
.expect("join Worker")
.expect("Worker result");
drop(received);
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("server stopped")
.expect("join server")
.expect("server result");
}
}
async fn await_burst_stream_count(
service: &BurstService,
registration: &PullRegistration,
expected: usize,
) {
timeout(Duration::from_secs(5), async {
loop {
assert!(registration.snapshot().is_ready);
if service.active_streams.load(Ordering::Acquire) == expected {
break;
}
service.changed.notified().await;
}
})
.await
.expect("stream count converged");
}
enum DrainEvent {
FinalFrame(Box<PollWorkerTaskFrame>),
HalfClosed,
ReadError(Status),
}
#[derive(Clone)]
struct DrainService {
events: mpsc::Sender<DrainEvent>,
assigned: Arc<AtomicBool>,
cancel_kind: Option<Code>,
read_delay: Duration,
is_reconnecting: bool,
ready: Option<mpsc::Sender<()>>,
}
#[tonic::async_trait]
impl WorkerTaskService for DrainService {
type PollWorkerTaskStream = ReceiverStream<Result<PollWorkerTaskResponse, Status>>;
async fn poll_worker_task(
&self,
request: Request<tonic::Streaming<PollWorkerTaskRequest>>,
) -> Result<Response<Self::PollWorkerTaskStream>, Status> {
let mut incoming = request.into_inner();
let first = incoming
.message()
.await?
.ok_or_else(|| Status::invalid_argument("missing Register"))?;
let Some(RequestFrame::Register(register)) = first.frame else {
return Err(Status::invalid_argument("Register must be first"));
};
let (sender, receiver) = mpsc::channel(4);
let registered = PollWorkerTaskResponse {
frame: Some(ResponseFrame::Registered(PollWorkerTaskRegistered {
supports_unknown_name_rejection: false,
})),
};
sender
.send(Ok(registered))
.await
.map_err(|_| Status::internal("registration response closed"))?;
let is_task_stream =
register.queue == "DrainFlow" && !self.assigned.swap(true, Ordering::AcqRel);
if is_task_stream {
let request = InvokeWorkerRpcRequest {
context: Some(ProtoContext {
flow_id: "flow-1".to_string(),
run_id: "run-1".to_string(),
..Default::default()
}),
flow_type: "DrainFlow".to_string(),
rpc_name: "get_drain_result".to_string(),
input: Some(
crate::value_mapper::encode(&())
.map_err(|error| Status::internal(error.to_string()))?,
),
..Default::default()
};
let task = PollWorkerTaskResponse {
frame: Some(ResponseFrame::Task(PollWorkerTask {
task_id: "drain-task".to_string(),
deadline: (self.cancel_kind == Some(Code::DeadlineExceeded)).then_some(
prost_types::Timestamp {
seconds: 0,
nanos: 0,
},
),
request: Some(TaskRequest::Rpc(request)),
})),
};
sender
.send(Ok(task))
.await
.map_err(|_| Status::internal("task response closed"))?;
if self.cancel_kind == Some(Code::Cancelled) {
sender
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::CancelTask(PollWorkerTaskCancelTask {
task_id: "drain-task".to_string(),
})),
}))
.await
.map_err(|_| Status::internal("cancel response closed"))?;
}
if self.is_reconnecting {
sender
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Reconnect(PollWorkerTaskReconnect {
reason: "test".to_string(),
})),
}))
.await
.map_err(|_| Status::internal("reconnect response closed"))?;
}
if !self.read_delay.is_zero() {
sender
.send(Ok(PollWorkerTaskResponse {
frame: Some(ResponseFrame::Ping(PollWorkerTaskPing {})),
}))
.await
.map_err(|_| Status::internal("ping response closed"))?;
}
}
let events = self.events.clone();
let read_delay = self.read_delay;
let ready = self.ready.clone();
let _stream_reader = tokio::spawn(async move {
if is_task_stream && !read_delay.is_zero() {
match incoming.message().await {
Ok(Some(PollWorkerTaskRequest {
frame: Some(RequestFrame::Pong(_)),
})) => {}
Ok(other) => {
eprintln!("expected Pong before delayed read, received {other:?}");
return;
}
Err(error) => {
eprintln!("read Pong before delayed read: {error}");
return;
}
}
if let Some(ready) = ready
&& let Err(error) = ready.send(()).await
{
eprintln!("record delayed read readiness: {error}");
return;
}
sleep(read_delay).await;
}
loop {
match incoming.message().await {
Ok(Some(PollWorkerTaskRequest {
frame: Some(RequestFrame::TaskFrame(frame)),
})) if is_task_stream => {
if let Err(error) =
events.send(DrainEvent::FinalFrame(Box::new(frame))).await
{
eprintln!("record final frame: {error}");
break;
}
}
Ok(Some(_)) => {}
Ok(None) => {
if is_task_stream
&& let Err(error) = events.send(DrainEvent::HalfClosed).await
{
eprintln!("record half-close: {error}");
}
break;
}
Err(status) => {
if is_task_stream
&& let Err(error) = events.send(DrainEvent::ReadError(status)).await
{
eprintln!("record stream error: {error}");
}
break;
}
}
}
drop(sender);
});
Ok(Response::new(ReceiverStream::new(receiver)))
}
}
#[tokio::test(flavor = "multi_thread")]
async fn stop_drain_delivers_final_frame_before_half_close() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind fake task service");
let address = listener.local_addr().expect("fake service address");
let (events, mut received) = mpsc::channel(4);
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(DrainService {
events,
assigned: Arc::new(AtomicBool::new(false)),
cancel_kind: None,
read_delay: Duration::ZERO,
is_reconnecting: false,
ready: None,
}))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok() {}
}),
);
let channel = Endpoint::from_shared(format!("http://{address}"))
.expect("endpoint")
.connect_lazy();
let channel = CredentialChannel::new(channel, CredentialInterceptor::new(None));
let (cache, _directory) = test_cache();
let started = Arc::new(Notify::new());
let release = Arc::new((Mutex::new(false), Condvar::new()));
let registry = Registry::new()
.register(DrainFlow {
started: Arc::clone(&started),
release: Arc::clone(&release),
wait_for_cancel: false,
result_size: None,
})
.expect("register Drain Flow");
let dispatcher = WorkerDispatcher::new(
registry,
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker = PullWorker::new(
&PullWorkerOptions::new().app_worker_id("worker-drain"),
&dispatcher,
channel,
)
.expect("pull Worker");
let (stop, shutdown) = watch::channel(false);
let running = tokio::spawn(async move { worker.run(shutdown).await });
timeout(Duration::from_secs(5), started.notified())
.await
.expect("RPC started");
stop.send_replace(true);
assert!(!running.is_finished());
let (lock, changed) = &*release;
*lock.lock().expect("drain gate lock") = true;
changed.notify_all();
timeout(Duration::from_secs(5), running)
.await
.expect("Worker stopped")
.expect("join Worker")
.expect("Worker result");
let final_frame = timeout(Duration::from_secs(5), received.recv())
.await
.expect("final frame received")
.expect("final frame event");
match final_frame {
DrainEvent::FinalFrame(frame) => {
assert_eq!("drain-task", frame.task_id);
assert!(matches!(frame.output, Some(Output::Rpc(_))));
}
DrainEvent::HalfClosed => panic!("stream closed before final frame"),
DrainEvent::ReadError(error) => panic!("stream failed before final frame: {error}"),
}
let closed = timeout(Duration::from_secs(5), received.recv())
.await
.expect("half-close received")
.expect("half-close event");
assert!(matches!(closed, DrainEvent::HalfClosed));
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("fake server stopped")
.expect("join server")
.expect("server result");
}
#[tokio::test(flavor = "multi_thread")]
async fn stop_drain_delivers_large_result_after_delayed_read() {
delayed_read_delivers_large_result(false, false).await;
}
#[tokio::test(flavor = "multi_thread")]
async fn busy_reconnect_delivers_large_result_after_delayed_read() {
delayed_read_delivers_large_result(true, false).await;
}
#[tokio::test(flavor = "multi_thread")]
async fn stop_rebounds_busy_reconnect_half_close() {
delayed_read_delivers_large_result(true, true).await;
}
async fn delayed_read_delivers_large_result(
is_reconnecting: bool,
should_stop_during_reconnect: bool,
) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind fake task service");
let address = listener.local_addr().expect("fake service address");
let (events, mut received) = mpsc::channel(4);
let (ready, mut ready_receiver) = mpsc::channel(1);
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(DrainService {
events,
assigned: Arc::new(AtomicBool::new(false)),
cancel_kind: None,
read_delay: if should_stop_during_reconnect {
Duration::from_secs(8)
} else {
Duration::from_millis(6_100)
},
is_reconnecting,
ready: Some(ready),
}))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok() {}
}),
);
let channel = Endpoint::from_shared(format!("http://{address}"))
.expect("endpoint")
.connect_lazy();
let channel = CredentialChannel::new(channel, CredentialInterceptor::new(None));
let (cache, _directory) = test_cache();
let started = Arc::new(Notify::new());
let release = Arc::new((Mutex::new(false), Condvar::new()));
let registry = Registry::new()
.register(DrainFlow {
started: Arc::clone(&started),
release: Arc::clone(&release),
wait_for_cancel: false,
result_size: Some(3_000_000),
})
.expect("register Drain Flow");
let dispatcher = WorkerDispatcher::new(
registry,
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker = PullWorker::new(
&PullWorkerOptions::new()
.app_worker_id("worker-drain")
.stop_drain_timeout(Duration::from_secs(7)),
&dispatcher,
channel,
)
.expect("pull Worker");
let (stop, shutdown) = watch::channel(false);
let running = tokio::spawn(async move { worker.run(shutdown).await });
timeout(Duration::from_secs(5), started.notified())
.await
.expect("RPC started");
timeout(Duration::from_secs(5), ready_receiver.recv())
.await
.expect("delayed read ready")
.expect("readiness event");
if !is_reconnecting {
stop.send_replace(true);
}
let (lock, changed) = &*release;
*lock.lock().expect("drain gate lock") = true;
changed.notify_all();
if should_stop_during_reconnect {
sleep(Duration::from_secs(3)).await;
stop.send_replace(true);
}
let final_frame = timeout(Duration::from_secs(10), received.recv())
.await
.expect("final frame received")
.expect("final frame event");
match final_frame {
DrainEvent::FinalFrame(frame) => {
assert_eq!("drain-task", frame.task_id);
assert!(matches!(frame.output, Some(Output::Rpc(_))));
assert!(frame.encoded_len() > 3_000_000);
}
DrainEvent::HalfClosed => panic!("stream closed before final frame"),
DrainEvent::ReadError(error) => panic!("stream failed before final frame: {error}"),
}
let closed = timeout(Duration::from_secs(5), received.recv())
.await
.expect("half-close received")
.expect("half-close event");
assert!(matches!(closed, DrainEvent::HalfClosed));
stop.send_replace(true);
timeout(Duration::from_secs(5), running)
.await
.expect("Worker stopped")
.expect("join Worker")
.expect("Worker result");
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("fake server stopped")
.expect("join server")
.expect("server result");
}
async fn cancellation_frame_has_status(cancel_kind: Code) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind fake task service");
let address = listener.local_addr().expect("fake service address");
let (events, mut received) = mpsc::channel(4);
let (server_stop, mut server_shutdown) = watch::channel(false);
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(DrainService {
events,
assigned: Arc::new(AtomicBool::new(false)),
cancel_kind: Some(cancel_kind),
read_delay: Duration::ZERO,
is_reconnecting: false,
ready: None,
}))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move {
while !*server_shutdown.borrow() && server_shutdown.changed().await.is_ok() {}
}),
);
let channel = Endpoint::from_shared(format!("http://{address}"))
.expect("endpoint")
.connect_lazy();
let channel = CredentialChannel::new(channel, CredentialInterceptor::new(None));
let (cache, _directory) = test_cache();
let started = Arc::new(Notify::new());
let release = Arc::new((Mutex::new(false), Condvar::new()));
let registry = Registry::new()
.register(DrainFlow {
started,
release,
wait_for_cancel: true,
result_size: None,
})
.expect("register cancellation Flow");
let dispatcher = WorkerDispatcher::new(
registry,
ValueHydrator::new(FlowServiceClient::new(channel.clone()), cache),
crate::server_protocol::NegotiatedServerProtocolVersion::default(),
);
let worker = PullWorker::new(
&PullWorkerOptions::new().app_worker_id("worker-cancel"),
&dispatcher,
channel,
)
.expect("pull Worker");
let (stop, shutdown) = watch::channel(false);
let running = tokio::spawn(async move { worker.run(shutdown).await });
let event = timeout(Duration::from_secs(5), received.recv())
.await
.expect("error frame received")
.expect("error frame event");
match event {
DrainEvent::FinalFrame(frame) => {
let Some(Output::WorkerError(error)) = frame.output else {
panic!("expected Worker error frame");
};
assert_eq!(cancel_kind as i32, error.grpc_status);
assert_eq!("Cancelled", error.error_type);
}
DrainEvent::HalfClosed => panic!("stream closed before error frame"),
DrainEvent::ReadError(error) => panic!("stream failed before error frame: {error}"),
}
stop.send_replace(true);
timeout(Duration::from_secs(5), running)
.await
.expect("Worker stopped")
.expect("join Worker")
.expect("Worker result");
server_stop.send_replace(true);
timeout(Duration::from_secs(5), server)
.await
.expect("fake server stopped")
.expect("join server")
.expect("server result");
}
#[tokio::test(flavor = "multi_thread")]
async fn cancel_task_reports_cancelled_status() {
cancellation_frame_has_status(Code::Cancelled).await;
}
#[tokio::test(flavor = "multi_thread")]
async fn deadline_reports_deadline_exceeded_status() {
cancellation_frame_has_status(Code::DeadlineExceeded).await;
}
}