aion_worker/runtime/loop_.rs
1//! receive->dispatch->report worker loop + bounded concurrency
2
3use std::collections::{BTreeMap, BTreeSet, HashMap};
4use std::sync::Arc;
5
6use aion_core::{ActivityError, ActivityId, Payload, WorkflowId};
7use async_trait::async_trait;
8use futures::StreamExt;
9use futures::future;
10use tokio::sync::{Semaphore, mpsc};
11use tracing::{debug, info};
12
13use crate::config::WorkerConfig;
14use crate::context::{ActivityContext, HeartbeatRequest};
15use crate::error::WorkerError;
16use crate::protocol::reconnect::UnackedResultTracker;
17use crate::protocol::{
18 ActivityExecutionKey, ActivityTask, HeartbeatBookkeeper, WorkerSession, WorkerSessionEvent,
19};
20use crate::runtime::report::{
21 DispatchFinished, InFlightActivity, RuntimeChannels, drain_remaining, record_first_error,
22 report_finished,
23};
24
25/// Dispatch seam used by the receive loop to execute decoded activity tasks.
26#[async_trait]
27pub trait ActivityDispatcher: Send + Sync + 'static {
28 /// Executes one decoded activity task with the provided handler context.
29 async fn dispatch(
30 &self,
31 task: ActivityTask,
32 context: ActivityContext,
33 ) -> Result<DispatchOutcome, WorkerError>;
34
35 /// Activity type names this dispatcher can serve.
36 fn activity_types(&self) -> BTreeSet<String>;
37}
38
39/// Activity execution outcome returned by the dispatch seam.
40#[derive(Clone, Debug, PartialEq, Eq)]
41pub enum DispatchOutcome {
42 /// Activity completed with an output payload.
43 Completed {
44 /// Opaque output payload.
45 output: Payload,
46 },
47 /// Activity failed with explicit classification.
48 Failed {
49 /// Classified activity failure.
50 failure: ActivityError,
51 },
52}
53
54/// Future that never resolves, used by the default serve entrypoint.
55pub type NoShutdown = future::Pending<()>;
56
57/// Why the serve loop ended without an error.
58#[derive(Clone, Copy, Debug, PartialEq, Eq)]
59pub enum ServeEnd {
60 /// The caller's shutdown future fired; in-flight work was drained.
61 Shutdown,
62 /// The server ended the task stream cleanly without announcing a drain.
63 /// The reconnect-aware run loop treats this unannounced close as a
64 /// budgeted retryable session drop — never as a run end.
65 StreamClosed,
66 /// The server announced a drain: in-flight work was finished and
67 /// reported, and the run loop reconnects after the schedule's initial
68 /// backoff without consuming any drop budget.
69 Drained,
70}
71
72/// Per-session health accounting written by the serve loop for the
73/// reconnect-aware caller's drop-budget reset decision.
74#[derive(Debug, Default)]
75pub struct SessionHealth {
76 /// Activity tasks whose outcome report was sent on this session.
77 pub tasks_reported: usize,
78 /// When the receive stream ended or dropped, captured before in-flight
79 /// handlers are drained — so post-drop draining never extends the
80 /// session's measured connected lifetime.
81 pub stream_ended_at: Option<tokio::time::Instant>,
82 /// Latched when a drain frame is observed on this session: the eventual
83 /// stream end — clean OR abrupt — is then drain-class (the server
84 /// announced it was going away), so the drop consumes no budget even if
85 /// the post-drain reporting fails. Survives an error return because this
86 /// is an out-parameter.
87 pub drain_received: bool,
88}
89
90/// Runs the worker receive loop until the session's task stream completes.
91///
92/// The RUNTIME owns liveness: for a session that carries a server-assigned
93/// heartbeat window ([`WorkerSession::heartbeat_window`]), the loop
94/// automatically heartbeats every in-flight activity at a quarter-window
95/// cadence, so a healthy worker running a legitimately long activity is never
96/// expired by the server's heartbeat sweeper. Explicit handler heartbeats
97/// remain the way to attach PROGRESS payloads; they are forwarded as they
98/// arrive. The loop never enforces heartbeat timeouts locally and never
99/// aborts running handler tasks on cancellation.
100///
101/// Every computed dispatch outcome is recorded in `tracker` before its report
102/// is sent, so a caller that reconnects after a transport drop can re-report
103/// the backlog; the server acks each consumed report (`ResultAck`), and only
104/// that explicit acknowledgement clears a tracker entry.
105///
106/// # Errors
107///
108/// Returns [`WorkerError`] when task decode, dispatch, heartbeat send, or result
109/// reporting fails.
110pub async fn serve_activity_tasks<S, D>(
111 config: &WorkerConfig,
112 session: &mut S,
113 dispatcher: Arc<D>,
114 tracker: &mut UnackedResultTracker,
115) -> Result<ServeEnd, WorkerError>
116where
117 S: WorkerSession,
118 D: ActivityDispatcher,
119{
120 let mut health = SessionHealth::default();
121 serve_activity_tasks_until(
122 config,
123 session,
124 dispatcher,
125 tracker,
126 &mut health,
127 future::pending(),
128 )
129 .await
130}
131
132/// Runs the worker receive loop until the session's task stream completes.
133///
134/// The RUNTIME owns liveness (#176): when the session carries a
135/// server-assigned heartbeat window ([`WorkerSession::heartbeat_window`],
136/// from the `RegisterAck`), the loop automatically sends a liveness heartbeat
137/// for EVERY in-flight activity at a quarter-window cadence
138/// ([`liveness_pump_interval`]). The server's heartbeat sweeper expires any
139/// worker whose in-flight task exceeds the window without a heartbeat — that
140/// is dead/wedged-PROCESS detection, and a healthy process running a
141/// multi-minute handler must never trip it, so keeping tasks beating is the
142/// runtime's job, not each handler's. A wedged process (deadlocked loop,
143/// stopped host) stops pumping and is correctly expired. Explicit handler
144/// heartbeats remain the way to attach PROGRESS payloads and are forwarded as
145/// they arrive; the loop never enforces heartbeat timeouts locally and never
146/// aborts running handler tasks on cancellation.
147///
148/// Every computed dispatch outcome is recorded in `tracker` before its report
149/// is sent, so a caller that reconnects after a transport drop can re-report
150/// the backlog; the server ingests reports idempotently and acks each one
151/// with a `ResultAck` frame. Only that explicit acknowledgement clears a
152/// tracker entry — a successful send proves nothing on its own.
153///
154/// `health` accumulates session-health accounting: the activity tasks whose
155/// outcome report was sent on this session, and the instant the receive
156/// stream ended (captured before in-flight handlers are drained). It is an
157/// out-parameter (rather than part of the return value) so the accounting
158/// survives an error return: the reconnect-aware caller uses it for the
159/// drop-budget reset decision — a session that served at least one task, or
160/// that stayed connected longer than the maximum backoff delay measured to
161/// the recorded stream end (never to the end of the post-drop drain), resets
162/// the cumulative drop budget even when it later drops.
163///
164/// On a clean end this returns [`ServeEnd`] distinguishing a caller-driven
165/// shutdown from a server-side stream close, so the caller can treat the
166/// latter as a retryable drop.
167///
168/// # Errors
169///
170/// Returns [`WorkerError`] when task decode, dispatch, heartbeat send, or result
171/// reporting fails.
172pub async fn serve_activity_tasks_until<S, D, Shutdown>(
173 config: &WorkerConfig,
174 session: &mut S,
175 dispatcher: Arc<D>,
176 tracker: &mut UnackedResultTracker,
177 health: &mut SessionHealth,
178 shutdown: Shutdown,
179) -> Result<ServeEnd, WorkerError>
180where
181 S: WorkerSession,
182 D: ActivityDispatcher,
183 Shutdown: Future<Output = ()> + Send,
184{
185 ensure_max_concurrency(config)?;
186 let semaphore = Arc::new(Semaphore::new(config.max_concurrency));
187 let (result_sender, heartbeat_sender, mut channels) = runtime_channels();
188 let bookkeeper = HeartbeatBookkeeper::default();
189 let mut liveness_pump = liveness_pump_for(session);
190 let mut stream = session.receive_tasks();
191 let mut in_flight = HashMap::<ActivityExecutionKey, InFlightActivity>::new();
192 let mut pending_error = None;
193 // Overridden at the shutdown break sites; every other clean exit is the
194 // server ending the stream.
195 let mut end = ServeEnd::StreamClosed;
196 // Every handle a spawned activity needs is stable for the loop's whole
197 // life, so it is assembled ONCE here rather than rebuilt at every task
198 // frame. It borrows the two senders, so it is dropped explicitly below
199 // BEFORE they are — the drain depends on those channels closing.
200 let spawn = SpawnContext::new(dispatcher, &result_sender, &heartbeat_sender, &bookkeeper);
201 tokio::pin!(shutdown);
202
203 // No batching preamble: the select arms below consume queued dispatch
204 // outcomes and heartbeats directly, so nothing waits for a stream event.
205 while pending_error.is_none() {
206 tokio::select! {
207 biased;
208 () = &mut shutdown => {
209 cancel_all_in_flight(&in_flight);
210 end = ServeEnd::Shutdown;
211 break;
212 }
213 // Dispatch outcomes are reported the moment they complete — the
214 // loop must not sit in `stream.next()` while a finished result
215 // waits, or a single dispatched task on an otherwise idle stream
216 // is only reported when the stream ends (the server-side dispatch
217 // would time out against a healthy worker).
218 finished = channels.results.recv() => {
219 consume_finished(
220 session,
221 &bookkeeper,
222 finished,
223 &mut in_flight,
224 tracker,
225 health,
226 &mut pending_error,
227 )
228 .await;
229 }
230 // Handler heartbeats are forwarded as they arrive for the same
231 // reason: the server's liveness window must be beatable while the
232 // stream is idle.
233 request = channels.heartbeats.recv() => {
234 forward_heartbeat(session, &bookkeeper, request, &mut pending_error).await;
235 }
236 // Automatic connection lease beat plus per-task liveness beats.
237 // This arm remains active while idle so an open but wedged runtime
238 // becomes detectable by the server's connection lease.
239 () = tick_liveness_pump(&mut liveness_pump) => {
240 pump_liveness(session, &bookkeeper, &in_flight, &mut pending_error).await;
241 }
242 event = stream.next() => {
243 let Some(event) = event else { break; };
244 match event {
245 Ok(WorkerSessionEvent::Cancel { workflow_id, activity_id }) => {
246 deliver_cancellation(workflow_id, &activity_id, &in_flight);
247 }
248 // Acks are bookkeeping, not work: consumed without a
249 // concurrency permit, like cancellation delivery.
250 Ok(WorkerSessionEvent::ResultAck { workflow_id, activity_id }) => {
251 acknowledge_result(&workflow_id, &activity_id, tracker);
252 }
253 Ok(WorkerSessionEvent::LivenessPing { sequence, silence_window }) => {
254 answer_ping(session, sequence, silence_window, &mut pending_error).await;
255 }
256 Ok(WorkerSessionEvent::Drain) => {
257 info!("server drain received; finishing in-flight work before reconnect");
258 health.drain_received = true;
259 end = ServeEnd::Drained;
260 break;
261 }
262 Err(error) => {
263 pending_error = Some(error);
264 break;
265 }
266 Ok(WorkerSessionEvent::Task(proto_task)) => {
267 let permit =
268 acquire_permit_or_shutdown(shutdown.as_mut(), &semaphore).await?;
269 let Some(permit) = permit else {
270 cancel_all_in_flight(&in_flight);
271 end = ServeEnd::Shutdown;
272 break;
273 };
274 if !spawn.admit(
275 *proto_task,
276 permit,
277 &mut in_flight,
278 &mut pending_error,
279 )? {
280 break;
281 }
282 }
283 }
284 }
285 }
286 }
287
288 // The stream just ended — cleanly, by error, or by shutdown. Capture the
289 // moment before draining in-flight handlers so the caller's drop-budget
290 // reset decision measures connected time, never drain time.
291 health.stream_ended_at = Some(tokio::time::Instant::now());
292
293 // The spawn context borrows both senders, so it goes first: `drain_remaining`
294 // reads the runtime channels to completion, and they only end once every
295 // sender is gone.
296 drop(spawn);
297 drop((result_sender, heartbeat_sender));
298 drain_remaining(
299 session,
300 &bookkeeper,
301 &mut channels,
302 &mut in_flight,
303 tracker,
304 &mut health.tasks_reported,
305 &mut pending_error,
306 )
307 .await;
308
309 pending_error.map_or(Ok(end), Err)
310}
311
312/// Builds the runtime's dispatch-outcome and heartbeat channels.
313fn runtime_channels() -> (
314 mpsc::UnboundedSender<DispatchFinished>,
315 mpsc::UnboundedSender<HeartbeatRequest>,
316 RuntimeChannels,
317) {
318 let (result_sender, result_receiver) = mpsc::unbounded_channel();
319 let (heartbeat_sender, heartbeat_receiver) = mpsc::unbounded_channel();
320 let channels = RuntimeChannels {
321 heartbeats: heartbeat_receiver,
322 results: result_receiver,
323 };
324 (result_sender, heartbeat_sender, channels)
325}
326
327/// The handles every spawned activity needs, assembled once for the receive
328/// loop's whole life.
329///
330/// Held as one value rather than seven locals so the loop's task arm stays a
331/// call rather than a seven-field literal — the shape that pushed
332/// [`serve_activity_tasks_until`] past its length budget when the liveness-ping
333/// arm landed beside it.
334struct SpawnContext<'a, D> {
335 dispatcher: Arc<D>,
336 result_sender: &'a mpsc::UnboundedSender<DispatchFinished>,
337 heartbeat_sender: &'a mpsc::UnboundedSender<HeartbeatRequest>,
338 heartbeat_bookkeeper: &'a HeartbeatBookkeeper,
339}
340
341impl<'a, D> SpawnContext<'a, D>
342where
343 D: ActivityDispatcher,
344{
345 /// Bundle the loop's stable spawn handles.
346 const fn new(
347 dispatcher: Arc<D>,
348 result_sender: &'a mpsc::UnboundedSender<DispatchFinished>,
349 heartbeat_sender: &'a mpsc::UnboundedSender<HeartbeatRequest>,
350 heartbeat_bookkeeper: &'a HeartbeatBookkeeper,
351 ) -> Self {
352 Self {
353 dispatcher,
354 result_sender,
355 heartbeat_sender,
356 heartbeat_bookkeeper,
357 }
358 }
359
360 /// Decode one pushed task and spawn it against `permit`.
361 ///
362 /// Returns `false` when the task could not be decoded: the error is
363 /// recorded in `pending_error` and the caller must stop serving. The permit
364 /// is released on that path, because nothing was spawned to hold it.
365 fn admit(
366 &self,
367 proto_task: aion_proto::ProtoActivityTask,
368 permit: tokio::sync::OwnedSemaphorePermit,
369 in_flight: &mut HashMap<ActivityExecutionKey, InFlightActivity>,
370 pending_error: &mut Option<WorkerError>,
371 ) -> Result<bool, WorkerError> {
372 let task = match ActivityTask::try_from(proto_task) {
373 Ok(task) => task,
374 Err(error) => {
375 drop(permit);
376 *pending_error = Some(error);
377 return Ok(false);
378 }
379 };
380 spawn_activity(
381 task,
382 permit,
383 Arc::clone(&self.dispatcher),
384 self.result_sender.clone(),
385 self.heartbeat_sender.clone(),
386 self.heartbeat_bookkeeper,
387 in_flight,
388 )?;
389 Ok(true)
390 }
391}
392
393/// Rejects a zero `max_concurrency` before the serve loop starts.
394fn ensure_max_concurrency(config: &WorkerConfig) -> Result<(), WorkerError> {
395 if config.max_concurrency == 0 {
396 return Err(WorkerError::registration(InvalidMaxConcurrency));
397 }
398 Ok(())
399}
400
401/// Waits for a dispatch permit, racing the caller's shutdown future; returns
402/// `None` when shutdown won.
403async fn acquire_permit_or_shutdown<F>(
404 shutdown: std::pin::Pin<&mut F>,
405 semaphore: &Arc<Semaphore>,
406) -> Result<Option<tokio::sync::OwnedSemaphorePermit>, WorkerError>
407where
408 F: Future<Output = ()> + Send,
409{
410 tokio::select! {
411 biased;
412 () = shutdown => Ok(None),
413 permit = Arc::clone(semaphore).acquire_owned() => {
414 permit.map(Some).map_err(WorkerError::registration)
415 }
416 }
417}
418
419/// Build the automatic liveness pump for a session: sessions registered
420/// against a server heartbeat window ([`WorkerSession::heartbeat_window`])
421/// beat every in-flight activity at a quarter-window cadence so the server's
422/// expiry sweeper only ever fires on a genuinely dead/wedged process.
423/// Sessions without a window (fakes, tests) never pump — byte-identical to
424/// the pre-pump loop.
425fn liveness_pump_for<S>(session: &S) -> Option<tokio::time::Interval>
426where
427 S: WorkerSession,
428{
429 session.heartbeat_window().map(|window| {
430 let mut ticks = tokio::time::interval(liveness_pump_interval(window));
431 ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
432 ticks
433 })
434}
435
436/// Consume one queued dispatch outcome (a `None` channel read is a no-op)
437/// and report it through the session, mirroring the drain path's
438/// [`report_finished`].
439async fn consume_finished<S>(
440 session: &mut S,
441 heartbeat_bookkeeper: &HeartbeatBookkeeper,
442 finished: Option<DispatchFinished>,
443 in_flight: &mut HashMap<ActivityExecutionKey, InFlightActivity>,
444 tracker: &mut UnackedResultTracker,
445 health: &mut SessionHealth,
446 pending_error: &mut Option<WorkerError>,
447) where
448 S: WorkerSession,
449{
450 if let Some(finished) = finished {
451 report_finished(
452 session,
453 heartbeat_bookkeeper,
454 finished,
455 in_flight,
456 tracker,
457 &mut health.tasks_reported,
458 pending_error,
459 )
460 .await;
461 }
462}
463
464/// Automatic liveness-heartbeat cadence derived from the server-assigned
465/// heartbeat window: a quarter of the window, floored at one millisecond
466/// (`tokio::time::interval` rejects a zero period).
467///
468/// The server expires a task once it goes longer than the WHOLE window
469/// without a heartbeat, so a quarter-window pump gives roughly four beats per
470/// window — comfortably inside the contract even when an individual beat is
471/// delayed by a busy loop iteration. Deliberately derived rather than
472/// configurable: the window is the server operator's contract, and the pump
473/// cadence is an implementation detail of honouring it (mirroring the
474/// server's own derived sweep cadence).
475#[must_use]
476pub(crate) fn liveness_pump_interval(heartbeat_window: std::time::Duration) -> std::time::Duration {
477 (heartbeat_window / 4).max(std::time::Duration::from_millis(1))
478}
479
480/// Resolves on the next automatic liveness tick, or never for sessions
481/// without a server-assigned heartbeat window (fakes and unregistered
482/// sessions never pump).
483async fn tick_liveness_pump(pump: &mut Option<tokio::time::Interval>) {
484 match pump {
485 Some(ticks) => {
486 ticks.tick().await;
487 }
488 None => future::pending().await,
489 }
490}
491
492/// Sends one connection lease heartbeat followed by a task liveness heartbeat
493/// (with no progress payload) for every in-flight activity.
494async fn pump_liveness<S>(
495 session: &mut S,
496 heartbeat_bookkeeper: &HeartbeatBookkeeper,
497 in_flight: &HashMap<ActivityExecutionKey, InFlightActivity>,
498 pending_error: &mut Option<WorkerError>,
499) where
500 S: WorkerSession,
501{
502 record_first_error(pending_error, session.send_connection_heartbeat().await);
503 if pending_error.is_some() {
504 return;
505 }
506 for key in in_flight.keys() {
507 record_first_error(
508 pending_error,
509 crate::protocol::send_heartbeat(
510 session,
511 heartbeat_bookkeeper,
512 HeartbeatRequest {
513 workflow_id: key.workflow_id.clone(),
514 activity_id: key.activity_id.clone(),
515 detail: None,
516 },
517 )
518 .await,
519 );
520 if pending_error.is_some() {
521 // The session send path is broken; the loop is about to exit
522 // with this error, so further beats are pointless.
523 return;
524 }
525 }
526}
527
528/// Answer one server transport liveness ping (#197).
529///
530/// Two properties of WHERE this is called are load-bearing, and both are
531/// properties of the caller rather than of this function:
532///
533/// - it takes no concurrency permit, so a worker running at full concurrency
534/// still answers. The server is asking whether it can REACH this worker, and
535/// a worker that is busy is reachable;
536/// - no handler is consulted, so the answer cannot be delayed or refused by
537/// action code. Liveness is a RUNTIME obligation.
538///
539/// The declared silence window is logged rather than acted on. It is what the
540/// SERVER will judge this worker against, and an operator asking why a worker
541/// never clears its probation needs both halves of the exchange in one place. A
542/// worker-side dead-man switch off that window belongs to the liminal transport
543/// ([`runtime::liminal_liveness`](crate::runtime::liminal_liveness)).
544async fn answer_ping<S>(
545 session: &mut S,
546 sequence: u64,
547 silence_window: std::time::Duration,
548 pending_error: &mut Option<WorkerError>,
549) where
550 S: WorkerSession,
551{
552 tracing::trace!(
553 liveness_ping = sequence,
554 silence_window_ms = silence_window.as_millis(),
555 "answering the server's transport liveness ping"
556 );
557 record_first_error(pending_error, session.answer_liveness_ping(sequence).await);
558}
559
560/// Forwards one queued handler heartbeat (a `None` channel read is a no-op)
561/// to the session, recording the first error.
562async fn forward_heartbeat<S>(
563 session: &mut S,
564 heartbeat_bookkeeper: &HeartbeatBookkeeper,
565 request: Option<HeartbeatRequest>,
566 pending_error: &mut Option<WorkerError>,
567) where
568 S: WorkerSession,
569{
570 if let Some(request) = request {
571 record_first_error(
572 pending_error,
573 crate::protocol::send_heartbeat(session, heartbeat_bookkeeper, request).await,
574 );
575 }
576}
577
578/// Clears the acknowledged tracker entry; an unknown ack (already cleared on
579/// a previous session, or replaced by a re-record) is a logged no-op.
580fn acknowledge_result(
581 workflow_id: &WorkflowId,
582 activity_id: &ActivityId,
583 tracker: &mut UnackedResultTracker,
584) {
585 if tracker.acknowledge(workflow_id, activity_id).is_some() {
586 debug!(
587 workflow_id = %workflow_id,
588 activity_id = activity_id.sequence_position(),
589 "server acknowledged activity result; tracker entry cleared"
590 );
591 } else {
592 debug!(
593 workflow_id = %workflow_id,
594 activity_id = activity_id.sequence_position(),
595 "result ack for unknown tracker entry ignored"
596 );
597 }
598}
599
600/// Render an activity's display labels as a compact, log-friendly
601/// `key=value` list in stable key order (for example `brief=IP-001
602/// repo=ablative-io/yggdrasil`). Empty when the workflow attached none.
603fn render_labels(labels: &BTreeMap<String, String>) -> String {
604 labels
605 .iter()
606 .map(|(key, value)| format!("{key}={value}"))
607 .collect::<Vec<_>>()
608 .join(" ")
609}
610
611fn spawn_activity<D>(
612 task: ActivityTask,
613 permit: tokio::sync::OwnedSemaphorePermit,
614 dispatcher: Arc<D>,
615 result_sender: mpsc::UnboundedSender<DispatchFinished>,
616 heartbeat_sender: mpsc::UnboundedSender<HeartbeatRequest>,
617 heartbeat_bookkeeper: &HeartbeatBookkeeper,
618 in_flight: &mut HashMap<ActivityExecutionKey, InFlightActivity>,
619) -> Result<(), WorkerError>
620where
621 D: ActivityDispatcher,
622{
623 info!(
624 activity_type = %task.activity_type,
625 activity_id = task.activity_id.sequence_position(),
626 workflow_id = %task.workflow_id,
627 attempt = task.attempt,
628 labels = %render_labels(&task.labels),
629 "received activity task"
630 );
631 let key = ActivityExecutionKey::new(task.workflow_id.clone(), task.activity_id.clone());
632 heartbeat_bookkeeper.register(key.clone())?;
633 let (context, cancellation_handle) = ActivityContext::for_task(
634 task.workflow_id.clone(),
635 task.run_id.clone(),
636 task.activity_id.clone(),
637 task.attempt,
638 task.idempotency_key.clone(),
639 Some(heartbeat_sender),
640 );
641 let finished_key = key.clone();
642 let finished_run_id = task.run_id.clone();
643 let finished_completion_token = task.completion_token.clone();
644 let join_handle = tokio::spawn(async move {
645 let outcome = dispatcher.dispatch(task, context).await;
646 if result_sender
647 .send(DispatchFinished {
648 key: finished_key,
649 run_id: finished_run_id,
650 completion_token: finished_completion_token,
651 outcome,
652 })
653 .is_err()
654 {
655 debug!("worker loop stopped before dispatch outcome could be delivered");
656 }
657 drop(permit);
658 });
659 in_flight.insert(
660 key,
661 InFlightActivity {
662 cancellation_handle,
663 join_handle,
664 },
665 );
666 Ok(())
667}
668
669fn deliver_cancellation(
670 workflow_id: WorkflowId,
671 activity_id: &ActivityId,
672 in_flight: &HashMap<ActivityExecutionKey, InFlightActivity>,
673) {
674 let key = ActivityExecutionKey::new(workflow_id, activity_id.clone());
675 if let Some(in_flight_activity) = in_flight.get(&key) {
676 in_flight_activity.cancellation_handle.cancel();
677 info!(
678 activity_id = activity_id.sequence_position(),
679 "delivered cooperative activity cancellation"
680 );
681 }
682}
683
684fn cancel_all_in_flight(in_flight: &HashMap<ActivityExecutionKey, InFlightActivity>) {
685 for (key, in_flight_activity) in in_flight {
686 in_flight_activity.cancellation_handle.cancel();
687 info!(
688 activity_id = key.activity_id.sequence_position(),
689 workflow_id = %key.workflow_id,
690 "delivered cooperative activity cancellation during worker shutdown"
691 );
692 }
693}
694
695#[derive(Debug, thiserror::Error)]
696#[error("worker max_concurrency must be greater than zero")]
697struct InvalidMaxConcurrency;
698
699#[cfg(test)]
700#[path = "loop_tests.rs"]
701mod tests;
702
703/// Head-of-line measurement for the transport liveness ping (#197): what a
704/// ping's answer costs when a delivery is already in flight, and what it costs
705/// when the concurrency budget is exhausted.
706#[cfg(test)]
707#[path = "liveness_answer_tests.rs"]
708mod liveness_answer_tests;