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

aster_forge_tasks/
dispatch.rs
