aster_forge_tasks/
dispatch.rs

1//! Shared dispatcher mechanics and result counters.
2
3use std::future::Future;
4use std::time::Duration;
5
6use chrono::{DateTime, Utc};
7use futures::stream::{self, StreamExt};
8
9use crate::{TaskCoreError, TaskLease, TaskRecord, task_lease_expires_at};
10
11/// Product-owned dispatch lane configuration.
12///
13/// The lane enum and task kind enum stay in product crates, but the scheduler shape is shared:
14/// every lane has a bounded concurrency limit, a fixed list of task kinds, an optional
15/// fast-continue mode for immediately filling freed slots, and a product lock key used by the
16/// store layer to serialize lane capacity checks.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct TaskLaneConfig<Kind: 'static, Lane> {
19    /// Product lane identifier.
20    pub lane: Lane,
21    /// Product task kinds assigned to this lane.
22    pub kinds: &'static [Kind],
23    /// Maximum active tasks for this lane.
24    pub limit: usize,
25    /// Whether the dispatcher should immediately claim another batch after finishing this one.
26    pub fast_continue: bool,
27    /// Product configuration key used as the lane-level transaction lock.
28    pub lock_key: &'static str,
29}
30
31/// Candidate task selected for compare-and-swap claiming.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct TaskClaimCandidate {
34    /// Index of the task in the originally listed due-task batch.
35    pub index: usize,
36    /// Persisted task identifier.
37    pub task_id: i64,
38    /// Processing token that must still be present for the claim to succeed.
39    pub expected_processing_token: i64,
40    /// Processing token to store when the claim succeeds.
41    pub next_processing_token: i64,
42}
43
44/// Successfully claimed task metadata.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct ClaimedTask {
47    /// Index of the task in the originally listed due-task batch.
48    pub index: usize,
49    /// Persisted task identifier.
50    pub task_id: i64,
51    /// Processing token assigned to the successful claim.
52    pub processing_token: i64,
53}
54
55/// Minimal task row view required by generic claim logic.
56pub trait ClaimableTaskRecord<Kind>: TaskRecord<Kind> {
57    /// Current processing token stored on the task row.
58    fn processing_token(&self) -> i64;
59}
60
61/// Product storage adapter required by generic lane claiming.
62///
63/// Implementations keep ownership of database transactions, row locks, product statuses, and ORM
64/// models. Forge only provides the shared claim algorithm around capacity checks, token overflow
65/// handling, lane validation, and conversion into [`TaskLease`] values.
66#[async_trait::async_trait]
67pub trait TaskClaimStore<Task, Kind: 'static, Lane>: Sync {
68    /// Product error type returned by storage operations.
69    type Error: From<TaskCoreError> + Send;
70
71    /// Lists due tasks that are claimable for the provided task kinds.
72    async fn list_claimable_by_kinds(
73        &self,
74        now: DateTime<Utc>,
75        stale_before: DateTime<Utc>,
76        kinds: &'static [Kind],
77        limit: u64,
78    ) -> std::result::Result<Vec<Task>, Self::Error>;
79
80    /// Counts active processing tasks for the provided kinds.
81    async fn count_active_processing_by_kinds(
82        &self,
83        now: DateTime<Utc>,
84        kinds: &'static [Kind],
85    ) -> std::result::Result<u64, Self::Error>;
86
87    /// Atomically claims candidate tasks after repeating the capacity check in product storage.
88    async fn claim_candidates_for_lane(
89        &self,
90        lane_config: TaskLaneConfig<Kind, Lane>,
91        candidates: &[TaskClaimCandidate],
92        stale_before: DateTime<Utc>,
93        claimed_at: DateTime<Utc>,
94        lease_expires_at: DateTime<Utc>,
95    ) -> std::result::Result<Vec<ClaimedTask>, Self::Error>;
96}
97
98/// Claims due tasks for a single lane using a product storage adapter.
99pub async fn claim_due_for_lane<Store, Task, Kind, Lane, LaneFn>(
100    store: &Store,
101    lane_config: TaskLaneConfig<Kind, Lane>,
102    processing_stale_secs: i64,
103    task_lane: LaneFn,
104) -> std::result::Result<Vec<(Task, TaskLease)>, Store::Error>
105where
106    Store: TaskClaimStore<Task, Kind, Lane>,
107    Task: ClaimableTaskRecord<Kind> + Clone + Send + Sync,
108    Kind: Copy + Eq + std::fmt::Debug + std::fmt::Display + Send + Sync + 'static,
109    Lane: Copy + Eq + std::fmt::Debug + Send + Sync + 'static,
110    LaneFn: Fn(Kind) -> Lane,
111{
112    if lane_config.limit == 0 {
113        return Ok(Vec::new());
114    }
115
116    let now = Utc::now();
117    let stale_before = now - chrono::Duration::seconds(processing_stale_secs.max(1));
118    let due = store
119        .list_claimable_by_kinds(
120            now,
121            stale_before,
122            lane_config.kinds,
123            claim_limit_to_u64(lane_config.limit),
124        )
125        .await?;
126    if due.is_empty() {
127        return Ok(Vec::new());
128    }
129
130    let active = store
131        .count_active_processing_by_kinds(now, lane_config.kinds)
132        .await?;
133    let available = available_lane_capacity(lane_config.limit, active);
134    if available == 0 {
135        tracing::debug!(
136            lane = ?lane_config.lane,
137            active,
138            limit = lane_config.limit,
139            "background task lane is at capacity; skipping claim"
140        );
141        return Ok(Vec::new());
142    }
143
144    let mut candidates = Vec::with_capacity(due.len());
145    for (index, task) in due.iter().enumerate() {
146        if task_lane(task.kind()) != lane_config.lane {
147            tracing::warn!(
148                task_id = task.id(),
149                kind = %task.kind(),
150                lane = ?lane_config.lane,
151                "claimable task kind does not match lane config; skipping"
152            );
153            continue;
154        }
155        let next_processing_token = task.processing_token().checked_add(1).ok_or_else(|| {
156            TaskCoreError::invalid_value("background task processing token overflow")
157        })?;
158
159        candidates.push(TaskClaimCandidate {
160            index,
161            task_id: task.id(),
162            expected_processing_token: task.processing_token(),
163            next_processing_token,
164        });
165    }
166    if candidates.is_empty() {
167        return Ok(Vec::new());
168    }
169
170    let claimed_at = Utc::now();
171    let claimed = store
172        .claim_candidates_for_lane(
173            lane_config,
174            &candidates,
175            stale_before,
176            claimed_at,
177            task_lease_expires_at(claimed_at, processing_stale_secs),
178        )
179        .await?;
180    let mut claimed_tasks = Vec::with_capacity(claimed.len());
181    for claim in claimed {
182        claimed_tasks.push((
183            due[claim.index].clone(),
184            TaskLease::new(claim.task_id, claim.processing_token),
185        ));
186    }
187
188    Ok(claimed_tasks)
189}
190
191/// Returns remaining lane capacity after subtracting currently active tasks.
192pub fn available_lane_capacity(limit: usize, active: u64) -> usize {
193    let active = usize::try_from(active).unwrap_or(usize::MAX);
194    limit.saturating_sub(active)
195}
196
197/// Converts a lane limit into the `u64` shape used by database query limits.
198pub fn claim_limit_to_u64(limit: usize) -> u64 {
199    u64::try_from(limit).unwrap_or(u64::MAX)
200}
201
202/// Runs item handlers with a bounded number of in-flight futures.
203///
204/// Claimed task batches are already capacity-checked by lane claiming. This helper only preserves
205/// the product's requested upper bound while letting independent task futures finish out of order.
206pub async fn run_with_concurrency_limit<T, O, F, Fut>(
207    items: Vec<T>,
208    limit: usize,
209    handler: F,
210) -> Vec<O>
211where
212    F: FnMut(T) -> Fut,
213    Fut: Future<Output = O>,
214{
215    stream::iter(items.into_iter().map(handler))
216        .buffer_unordered(limit.max(1))
217        .collect()
218        .await
219}
220
221/// Runs lane dispatchers concurrently and aggregates lane statistics.
222///
223/// If one or more lanes fail, the first error is returned after all lanes have completed. This
224/// mirrors the existing Yggdrasil and Drive behavior: independent lanes are allowed to finish, then
225/// the dispatch pass reports the first lane failure to the caller.
226pub async fn dispatch_lanes<Kind, Lane, Error, F, Fut>(
227    lane_configs: Vec<TaskLaneConfig<Kind, Lane>>,
228    lane_parallelism: usize,
229    dispatch_lane: F,
230) -> std::result::Result<DispatchStats, Error>
231where
232    Kind: Send + Sync + 'static,
233    Lane: Send + Sync,
234    F: FnMut(TaskLaneConfig<Kind, Lane>) -> Fut,
235    Fut: Future<Output = std::result::Result<DispatchStats, Error>>,
236{
237    let lane_results = stream::iter(lane_configs.into_iter().map(dispatch_lane))
238        .buffer_unordered(lane_parallelism.max(1))
239        .collect::<Vec<_>>()
240        .await;
241    let mut stats = DispatchStats::default();
242    let mut first_error = None;
243
244    for result in lane_results {
245        match result {
246            Ok(lane_stats) => stats.add(lane_stats),
247            Err(error) => {
248                if first_error.is_none() {
249                    first_error = Some(error);
250                }
251            }
252        }
253    }
254
255    if let Some(first_error) = first_error {
256        return Err(first_error);
257    }
258
259    Ok(stats)
260}
261
262/// Runs an already claimed task batch and aggregates task execution outcomes.
263///
264/// The batch is sorted before execution so products can preserve stable created-at/id ordering
265/// while still allowing individual task futures to finish out of order.
266pub async fn run_claimed_task_batch<T, SortKey, Error, SortFn, HandlerFn, HandlerFut>(
267    mut claimed_tasks: Vec<T>,
268    sort_key: SortFn,
269    handler: HandlerFn,
270) -> std::result::Result<DispatchStats, Error>
271where
272    SortKey: Ord,
273    SortFn: FnMut(&T) -> SortKey,
274    HandlerFn: FnMut(T) -> HandlerFut,
275    HandlerFut: Future<Output = std::result::Result<TaskDispatchOutcome, Error>>,
276{
277    let concurrency = claimed_tasks.len().max(1);
278    claimed_tasks.sort_by_key(sort_key);
279
280    let results = run_with_concurrency_limit(claimed_tasks, concurrency, handler).await;
281    let mut stats = DispatchStats::default();
282    let mut first_error = None;
283
284    for result in results {
285        match result {
286            Ok(outcome) => stats.add_outcome(outcome),
287            Err(error) => {
288                if first_error.is_none() {
289                    first_error = Some(error);
290                }
291            }
292        }
293    }
294
295    if let Some(first_error) = first_error {
296        return Err(first_error);
297    }
298
299    Ok(stats)
300}
301
302/// Drains a dispatcher until it has no claimed work and no active processing tasks.
303///
304/// Product crates provide the real dispatch function and processing-count query. The helper keeps
305/// the shared bounded retry loop and aggregate statistics in Forge while leaving storage and status
306/// semantics outside this crate.
307pub async fn drain_dispatcher<DispatchFn, DispatchFut, CountFn, CountFut, Error>(
308    max_rounds: usize,
309    processing_poll_interval: Duration,
310    mut dispatch_due: DispatchFn,
311    mut count_processing: CountFn,
312) -> std::result::Result<DispatchStats, Error>
313where
314    DispatchFn: FnMut() -> DispatchFut,
315    DispatchFut: Future<Output = std::result::Result<DispatchStats, Error>>,
316    CountFn: FnMut() -> CountFut,
317    CountFut: Future<Output = std::result::Result<u64, Error>>,
318{
319    let mut total = DispatchStats::default();
320    tracing::debug!("draining background task dispatcher");
321
322    for _ in 0..max_rounds {
323        let stats = dispatch_due().await?;
324        let claimed = stats.claimed;
325        total.add(stats);
326        if claimed > 0 {
327            continue;
328        }
329
330        if count_processing().await? == 0 {
331            tracing::debug!("background task drain finished because no tasks are processing");
332            break;
333        }
334
335        tokio::time::sleep(processing_poll_interval).await;
336    }
337
338    tracing::debug!(
339        claimed = total.claimed,
340        succeeded = total.succeeded,
341        retried = total.retried,
342        failed = total.failed,
343        "background task dispatcher drain completed"
344    );
345    Ok(total)
346}
347
348/// Aggregate counters returned by a background task dispatch pass.
349#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
350pub struct DispatchStats {
351    /// Number of tasks claimed for execution.
352    pub claimed: usize,
353    /// Number of tasks completed successfully.
354    pub succeeded: usize,
355    /// Number of tasks scheduled for retry.
356    pub retried: usize,
357    /// Number of tasks permanently failed.
358    pub failed: usize,
359}
360
361impl DispatchStats {
362    /// Adds another dispatch counter set into this one.
363    pub fn add(&mut self, other: Self) {
364        self.claimed += other.claimed;
365        self.succeeded += other.succeeded;
366        self.retried += other.retried;
367        self.failed += other.failed;
368    }
369
370    /// Returns whether any dispatch activity happened.
371    pub const fn has_activity(&self) -> bool {
372        self.claimed > 0 || self.succeeded > 0 || self.retried > 0 || self.failed > 0
373    }
374
375    /// Adds a task execution outcome to the aggregate counters.
376    pub fn add_outcome(&mut self, outcome: TaskDispatchOutcome) {
377        self.succeeded += outcome.succeeded;
378        self.retried += outcome.retried;
379        self.failed += outcome.failed;
380    }
381}
382
383/// Counters returned by one claimed task execution.
384#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
385pub struct TaskDispatchOutcome {
386    /// Number of tasks completed successfully.
387    pub succeeded: usize,
388    /// Number of tasks scheduled for retry.
389    pub retried: usize,
390    /// Number of tasks permanently failed.
391    pub failed: usize,
392}
393
394impl TaskDispatchOutcome {
395    /// Creates a successful task outcome.
396    pub const fn succeeded() -> Self {
397        Self {
398            succeeded: 1,
399            retried: 0,
400            failed: 0,
401        }
402    }
403
404    /// Creates a retried task outcome.
405    pub const fn retried() -> Self {
406        Self {
407            succeeded: 0,
408            retried: 1,
409            failed: 0,
410        }
411    }
412
413    /// Creates a permanently failed task outcome.
414    pub const fn failed() -> Self {
415        Self {
416            succeeded: 0,
417            retried: 0,
418            failed: 1,
419        }
420    }
421}
422
423#[cfg(test)]
424mod tests {
425    use std::sync::Arc;
426    use std::sync::atomic::{AtomicUsize, Ordering};
427    use std::time::Duration;
428
429    use super::{
430        ClaimableTaskRecord, ClaimedTask, DispatchStats, TaskClaimCandidate, TaskClaimStore,
431        TaskDispatchOutcome, TaskLaneConfig, available_lane_capacity, claim_due_for_lane,
432        claim_limit_to_u64, dispatch_lanes, drain_dispatcher, run_claimed_task_batch,
433        run_with_concurrency_limit,
434    };
435    use crate::TaskRecord;
436
437    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
438    enum TestLane {
439        Default,
440    }
441
442    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
443    enum TestKind {
444        Example,
445    }
446
447    impl std::fmt::Display for TestKind {
448        fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
449            match self {
450                Self::Example => formatter.write_str("example"),
451            }
452        }
453    }
454
455    #[derive(Debug, Clone)]
456    struct TestTask {
457        id: i64,
458        kind: TestKind,
459        processing_token: i64,
460    }
461
462    impl TaskRecord<TestKind> for TestTask {
463        fn id(&self) -> i64 {
464            self.id
465        }
466
467        fn kind(&self) -> TestKind {
468            self.kind
469        }
470
471        fn payload_json(&self) -> &str {
472            "{}"
473        }
474
475        fn result_json(&self) -> Option<&str> {
476            None
477        }
478    }
479
480    impl ClaimableTaskRecord<TestKind> for TestTask {
481        fn processing_token(&self) -> i64 {
482            self.processing_token
483        }
484    }
485
486    struct TestClaimStore {
487        due: Vec<TestTask>,
488        active: u64,
489    }
490
491    #[async_trait::async_trait]
492    impl TaskClaimStore<TestTask, TestKind, TestLane> for TestClaimStore {
493        type Error = crate::TaskCoreError;
494
495        async fn list_claimable_by_kinds(
496            &self,
497            _now: chrono::DateTime<chrono::Utc>,
498            _stale_before: chrono::DateTime<chrono::Utc>,
499            _kinds: &'static [TestKind],
500            _limit: u64,
501        ) -> std::result::Result<Vec<TestTask>, Self::Error> {
502            Ok(self.due.clone())
503        }
504
505        async fn count_active_processing_by_kinds(
506            &self,
507            _now: chrono::DateTime<chrono::Utc>,
508            _kinds: &'static [TestKind],
509        ) -> std::result::Result<u64, Self::Error> {
510            Ok(self.active)
511        }
512
513        async fn claim_candidates_for_lane(
514            &self,
515            _lane_config: TaskLaneConfig<TestKind, TestLane>,
516            candidates: &[TaskClaimCandidate],
517            _stale_before: chrono::DateTime<chrono::Utc>,
518            _claimed_at: chrono::DateTime<chrono::Utc>,
519            _lease_expires_at: chrono::DateTime<chrono::Utc>,
520        ) -> std::result::Result<Vec<ClaimedTask>, Self::Error> {
521            Ok(candidates
522                .iter()
523                .map(|candidate| ClaimedTask {
524                    index: candidate.index,
525                    task_id: candidate.task_id,
526                    processing_token: candidate.next_processing_token,
527                })
528                .collect())
529        }
530    }
531
532    const TEST_KINDS: &[TestKind] = &[TestKind::Example];
533
534    #[test]
535    fn available_lane_capacity_saturates_at_zero() {
536        assert_eq!(available_lane_capacity(4, 0), 4);
537        assert_eq!(available_lane_capacity(4, 2), 2);
538        assert_eq!(available_lane_capacity(4, 4), 0);
539        assert_eq!(available_lane_capacity(4, 9), 0);
540        assert_eq!(available_lane_capacity(4, u64::MAX), 0);
541    }
542
543    #[test]
544    fn claim_limit_to_u64_accepts_common_limits() {
545        assert_eq!(claim_limit_to_u64(0), 0);
546        assert_eq!(claim_limit_to_u64(16), 16);
547    }
548
549    #[tokio::test]
550    async fn claim_due_for_lane_returns_claimed_tasks_with_leases() {
551        let store = TestClaimStore {
552            due: vec![TestTask {
553                id: 42,
554                kind: TestKind::Example,
555                processing_token: 7,
556            }],
557            active: 0,
558        };
559        let lane_config = TaskLaneConfig {
560            lane: TestLane::Default,
561            kinds: TEST_KINDS,
562            limit: 2,
563            fast_continue: false,
564            lock_key: "test",
565        };
566
567        let claimed = claim_due_for_lane(&store, lane_config, 60, |_| TestLane::Default)
568            .await
569            .expect("claim should succeed");
570
571        assert_eq!(claimed.len(), 1);
572        assert_eq!(claimed[0].0.id, 42);
573        assert_eq!(claimed[0].1.task_id, 42);
574        assert_eq!(claimed[0].1.processing_token, 8);
575    }
576
577    #[tokio::test]
578    async fn claim_due_for_lane_skips_when_lane_is_at_capacity() {
579        let store = TestClaimStore {
580            due: vec![TestTask {
581                id: 42,
582                kind: TestKind::Example,
583                processing_token: 7,
584            }],
585            active: 2,
586        };
587        let lane_config = TaskLaneConfig {
588            lane: TestLane::Default,
589            kinds: TEST_KINDS,
590            limit: 2,
591            fast_continue: false,
592            lock_key: "test",
593        };
594
595        let claimed = claim_due_for_lane(&store, lane_config, 60, |_| TestLane::Default)
596            .await
597            .expect("claim should succeed");
598
599        assert!(claimed.is_empty());
600    }
601
602    #[test]
603    fn dispatch_stats_tracks_activity_and_adds_outcomes() {
604        let mut stats = DispatchStats::default();
605        assert!(!stats.has_activity());
606
607        stats.claimed = 2;
608        assert!(stats.has_activity());
609
610        stats.add_outcome(TaskDispatchOutcome {
611            succeeded: 1,
612            retried: 2,
613            failed: 3,
614        });
615        assert_eq!(stats.succeeded, 1);
616        assert_eq!(stats.retried, 2);
617        assert_eq!(stats.failed, 3);
618
619        stats.add(DispatchStats {
620            claimed: 4,
621            succeeded: 5,
622            retried: 6,
623            failed: 7,
624        });
625        assert_eq!(stats.claimed, 6);
626        assert_eq!(stats.succeeded, 6);
627        assert_eq!(stats.retried, 8);
628        assert_eq!(stats.failed, 10);
629    }
630
631    #[tokio::test]
632    async fn run_with_concurrency_limit_caps_parallelism() {
633        let active = Arc::new(AtomicUsize::new(0));
634        let peak = Arc::new(AtomicUsize::new(0));
635
636        let mut results = run_with_concurrency_limit(vec![1, 2, 3, 4, 5], 2, {
637            let active = active.clone();
638            let peak = peak.clone();
639            move |value| {
640                let active = active.clone();
641                let peak = peak.clone();
642                async move {
643                    let current = active.fetch_add(1, Ordering::SeqCst) + 1;
644                    peak.fetch_max(current, Ordering::SeqCst);
645                    tokio::time::sleep(Duration::from_millis(1)).await;
646                    active.fetch_sub(1, Ordering::SeqCst);
647                    value * 2
648                }
649            }
650        })
651        .await;
652        results.sort_unstable();
653
654        assert_eq!(results, vec![2, 4, 6, 8, 10]);
655        assert_eq!(peak.load(Ordering::SeqCst), 2);
656    }
657
658    #[tokio::test]
659    async fn dispatch_lanes_aggregates_successful_lane_stats() {
660        let lanes = vec![
661            TaskLaneConfig {
662                lane: TestLane::Default,
663                kinds: TEST_KINDS,
664                limit: 1,
665                fast_continue: false,
666                lock_key: "one",
667            },
668            TaskLaneConfig {
669                lane: TestLane::Default,
670                kinds: TEST_KINDS,
671                limit: 2,
672                fast_continue: false,
673                lock_key: "two",
674            },
675        ];
676
677        let stats = dispatch_lanes(lanes, 2, |lane| async move {
678            Ok::<_, ()>(DispatchStats {
679                claimed: lane.limit,
680                succeeded: lane.limit,
681                retried: 0,
682                failed: 0,
683            })
684        })
685        .await
686        .expect("lane dispatch should succeed");
687
688        assert_eq!(stats.claimed, 3);
689        assert_eq!(stats.succeeded, 3);
690    }
691
692    #[tokio::test]
693    async fn dispatch_lanes_returns_first_error_after_all_lanes_finish() {
694        let lanes = vec![
695            TaskLaneConfig {
696                lane: TestLane::Default,
697                kinds: TEST_KINDS,
698                limit: 1,
699                fast_continue: false,
700                lock_key: "one",
701            },
702            TaskLaneConfig {
703                lane: TestLane::Default,
704                kinds: TEST_KINDS,
705                limit: 2,
706                fast_continue: false,
707                lock_key: "two",
708            },
709        ];
710        let calls = Arc::new(AtomicUsize::new(0));
711
712        let error = dispatch_lanes(lanes, 2, {
713            let calls = calls.clone();
714            move |lane| {
715                let calls = calls.clone();
716                async move {
717                    calls.fetch_add(1, Ordering::SeqCst);
718                    if lane.limit == 1 {
719                        Err("first")
720                    } else {
721                        Ok(DispatchStats::default())
722                    }
723                }
724            }
725        })
726        .await
727        .expect_err("lane dispatch should return first error");
728
729        assert_eq!(error, "first");
730        assert_eq!(calls.load(Ordering::SeqCst), 2);
731    }
732
733    #[tokio::test]
734    async fn run_claimed_task_batch_sorts_and_aggregates_outcomes() {
735        let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
736
737        let stats = run_claimed_task_batch(
738            vec![
739                (2, TaskDispatchOutcome::retried()),
740                (1, TaskDispatchOutcome::succeeded()),
741            ],
742            |(order, _)| *order,
743            {
744                let seen = seen.clone();
745                move |(order, outcome)| {
746                    let seen = seen.clone();
747                    async move {
748                        match seen.lock() {
749                            Ok(mut seen) => seen.push(order),
750                            Err(poisoned) => poisoned.into_inner().push(order),
751                        }
752                        Ok::<_, ()>(outcome)
753                    }
754                }
755            },
756        )
757        .await
758        .expect("claimed task batch should succeed");
759
760        assert_eq!(stats.succeeded, 1);
761        assert_eq!(stats.retried, 1);
762        let seen = match seen.lock() {
763            Ok(seen) => seen.clone(),
764            Err(poisoned) => poisoned.into_inner().clone(),
765        };
766        assert_eq!(seen, vec![1, 2]);
767    }
768
769    #[tokio::test]
770    async fn drain_dispatcher_accumulates_until_queue_and_processing_are_empty() {
771        let dispatch_calls = Arc::new(AtomicUsize::new(0));
772        let count_calls = Arc::new(AtomicUsize::new(0));
773
774        let stats = drain_dispatcher(
775            4,
776            Duration::from_millis(1),
777            {
778                let dispatch_calls = dispatch_calls.clone();
779                move || {
780                    let dispatch_calls = dispatch_calls.clone();
781                    async move {
782                        let call = dispatch_calls.fetch_add(1, Ordering::SeqCst);
783                        let stats = match call {
784                            0 => DispatchStats {
785                                claimed: 2,
786                                succeeded: 1,
787                                retried: 0,
788                                failed: 0,
789                            },
790                            _ => DispatchStats::default(),
791                        };
792                        Ok::<_, ()>(stats)
793                    }
794                }
795            },
796            {
797                let count_calls = count_calls.clone();
798                move || {
799                    let count_calls = count_calls.clone();
800                    async move {
801                        count_calls.fetch_add(1, Ordering::SeqCst);
802                        Ok::<_, ()>(0)
803                    }
804                }
805            },
806        )
807        .await
808        .expect("drain should succeed");
809
810        assert_eq!(stats.claimed, 2);
811        assert_eq!(stats.succeeded, 1);
812        assert_eq!(dispatch_calls.load(Ordering::SeqCst), 2);
813        assert_eq!(count_calls.load(Ordering::SeqCst), 1);
814    }
815
816    #[tokio::test]
817    async fn drain_dispatcher_stops_at_max_rounds_when_processing_stays_active() {
818        let dispatch_calls = Arc::new(AtomicUsize::new(0));
819        let count_calls = Arc::new(AtomicUsize::new(0));
820
821        let stats = drain_dispatcher(
822            3,
823            Duration::from_millis(1),
824            {
825                let dispatch_calls = dispatch_calls.clone();
826                move || {
827                    let dispatch_calls = dispatch_calls.clone();
828                    async move {
829                        dispatch_calls.fetch_add(1, Ordering::SeqCst);
830                        Ok::<_, ()>(DispatchStats::default())
831                    }
832                }
833            },
834            {
835                let count_calls = count_calls.clone();
836                move || {
837                    let count_calls = count_calls.clone();
838                    async move {
839                        count_calls.fetch_add(1, Ordering::SeqCst);
840                        Ok::<_, ()>(1)
841                    }
842                }
843            },
844        )
845        .await
846        .expect("drain should succeed");
847
848        assert_eq!(stats, DispatchStats::default());
849        assert_eq!(dispatch_calls.load(Ordering::SeqCst), 3);
850        assert_eq!(count_calls.load(Ordering::SeqCst), 3);
851    }
852}
aster_forge_tasks/dispatch.rs

aster_forge_tasks/
dispatch.rs
