Skip to main content

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;