moirai-executor 0.8.2

#![cfg_attr(test, allow(clippy::unwrap_used, reason = "test scope"))]

#[cfg(test)]
#[allow(clippy::module_inception)]
mod tests {
    use std::{sync::Arc, time::Instant};

    use moirai_core::Priority;

    use super::super::directory::BlockLookup;
    use super::super::registry::{CancelOutcome, TaskRegistry};
    use super::super::state::{PRIORITY_FROM_INDEX, task_location};

    #[test]
    fn lifecycle_token_records_started_and_completed_metadata() {
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();

        let running = lifecycle.start(3);
        let started = registry.get_metadata(task_id).unwrap();
        assert_eq!(started.id, task_id);
        assert_eq!(started.worker_id, Some(3));
        let started_at = started.started_at.expect("start() must stamp started_at");
        assert!(started.completed_at.is_none());

        let execution_time = running.complete();

        let completed = registry.get_metadata(task_id).unwrap();
        let completed_at = completed
            .completed_at
            .expect("complete() must stamp completed_at");
        assert!(
            completed_at >= started_at,
            "completion must not precede start"
        );
        assert_eq!(completed.started_at, Some(started_at));
        assert_eq!(completed.execution_duration(), Some(execution_time));
        assert!(registry.is_completed(task_id));
    }

    #[test]
    fn register_next_task_returns_id_and_lifecycle_token() {
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();

        let running = lifecycle.start(2);
        let execution_time = running.complete();

        let metadata = registry.get_metadata(task_id).unwrap();
        assert_eq!(metadata.id, task_id);
        assert_eq!(metadata.worker_id, Some(2));
        let started_at = metadata.started_at.expect("start() must stamp started_at");
        let completed_at = metadata
            .completed_at
            .expect("complete() must stamp completed_at");
        assert!(
            completed_at >= started_at,
            "completion must not precede start"
        );
        assert_eq!(metadata.execution_duration(), Some(execution_time));
    }

    #[test]
    fn lifecycle_token_keeps_its_dense_block_alive_after_registry_drop() {
        let (lifecycle, block) = {
            let registry = TaskRegistry::new();
            let (task_id, lifecycle) = registry.register_next_task();
            let (block_index, _) = task_location(task_id);
            let block = match registry
                .blocks
                .read()
                .expect("test registry lock")
                .lookup(block_index)
            {
                BlockLookup::Live(block) => Arc::downgrade(block),
                _ => panic!("the registered task block must be resident"),
            };
            assert_eq!(block.strong_count(), 2);
            (lifecycle, block)
        };

        assert_eq!(block.strong_count(), 1);
        lifecycle.start(5).complete();
        assert!(
            block.upgrade().is_none(),
            "the token must release its sole block owner after completion"
        );
    }

    /// A token dropped before `start` means the job was discarded unrun, and
    /// its result sender drops with it, so the handle resolves to `Cancelled`.
    /// The registry records the same outcome rather than a completion.
    #[test]
    fn unstarted_lifecycle_token_drop_records_the_task_cancelled() {
        let before = Instant::now();
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();

        drop(lifecycle);

        let metadata = registry
            .get_metadata(task_id)
            .expect("registered task metadata must remain readable");
        assert_eq!(metadata.started_at, None);
        assert!(metadata.cancelled);
        // Dropping an unstarted token still closes the task out.
        assert!(
            metadata.completed_at >= Some(before),
            "the drop must stamp a completion time from this test's window"
        );
        assert_eq!(registry.active_count(), 0);
        assert_eq!(registry.completed_count(), 1);
    }

    #[test]
    fn running_lifecycle_token_dropped_unfinished_records_the_task_cancelled() {
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();

        drop(lifecycle.start(1));

        let metadata = registry.get_metadata(task_id).unwrap();
        assert!(registry.is_completed(task_id));
        assert!(metadata.cancelled);
        assert_eq!(metadata.worker_id, Some(1));
    }

    #[test]
    fn running_lifecycle_token_cancel_records_cancelled_and_completed() {
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();

        lifecycle.start(2).cancel();

        let metadata = registry.get_metadata(task_id).unwrap();
        assert!(metadata.cancelled);
        assert!(registry.is_completed(task_id));
        let started_at = metadata.started_at.expect("start() stamps started_at");
        let completed_at = metadata.completed_at.expect("cancel stamps completed_at");
        assert!(completed_at >= started_at);
        assert_eq!(metadata.worker_id, Some(2));
    }

    #[test]
    fn completed_lifecycle_token_is_not_recorded_cancelled() {
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();

        lifecycle.start(0).complete();

        let metadata = registry.get_metadata(task_id).unwrap();
        assert!(registry.is_completed(task_id));
        assert!(!metadata.cancelled);
    }

    /// `register_task` hands the caller an id to drive by hand; the task must
    /// not already be completed, or `mark_started` and `mark_completed` would
    /// be recording a lifecycle that had ended before it began.
    #[test]
    fn register_task_leaves_the_task_queued_for_the_caller_to_drive() {
        let registry = TaskRegistry::new();
        let task_id = registry.register_task();

        let queued = registry.get_metadata(task_id).unwrap();
        assert!(!registry.is_completed(task_id));
        assert_eq!(queued.started_at, None);
        assert_eq!(queued.completed_at, None);
        assert!(!queued.cancelled);
        assert_eq!(registry.active_count(), 1);

        registry.mark_started(task_id, 6);
        assert_eq!(registry.get_metadata(task_id).unwrap().worker_id, Some(6));
        assert!(!registry.is_completed(task_id));

        registry.mark_completed(task_id);
        let finished = registry.get_metadata(task_id).unwrap();
        assert!(registry.is_completed(task_id));
        let started_at = finished.started_at.expect("mark_started stamps started_at");
        let completed_at = finished
            .completed_at
            .expect("mark_completed stamps completed_at");
        assert!(completed_at >= started_at);
        assert!(!finished.cancelled);
        assert_eq!(registry.active_count(), 0);
    }

    /// The last task of a full block, registered by id, keeps the block
    /// resident until the caller completes it; a sweep that retired the block
    /// first would make the later `mark_started`/`mark_completed` no-ops.
    #[test]
    fn register_task_pins_its_block_until_the_caller_completes_it() {
        use std::time::Duration;

        use super::super::state::TASK_STATE_BLOCK_SIZE;

        let registry = TaskRegistry::new();
        for _ in 1..TASK_STATE_BLOCK_SIZE - 1 {
            let (_, lifecycle) = registry.register_next_task();
            lifecycle.start(0).complete();
        }
        let task_id = registry.register_task();
        assert_eq!(task_id, (TASK_STATE_BLOCK_SIZE - 1) as u64);

        assert_eq!(registry.cleanup_completed(Duration::ZERO), 0);
        registry.mark_started(task_id, 3);
        assert_eq!(registry.get_metadata(task_id).unwrap().worker_id, Some(3));

        registry.mark_completed(task_id);
        assert_eq!(registry.cleanup_completed(Duration::ZERO), 1);
        assert!(registry.is_completed(task_id));
    }

    #[test]
    fn priority_index_round_trips() {
        // PRIORITY_FROM_INDEX is the inverse of Priority::index (SSOT pair).
        for priority in [
            Priority::Low,
            Priority::Normal,
            Priority::High,
            Priority::Critical,
        ] {
            assert_eq!(PRIORITY_FROM_INDEX[priority.index()], priority);
        }
        assert_eq!(Priority::Low.index(), 0);
        assert_eq!(Priority::Normal.index(), 1);
        assert_eq!(Priority::High.index(), 2);
        assert_eq!(Priority::Critical.index(), 3);
    }

    #[test]
    fn lifecycle_token_records_spawn_priority() {
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();
        lifecycle.set_priority(Priority::Critical);

        assert_eq!(
            registry.get_metadata(task_id).unwrap().priority,
            Priority::Critical
        );

        lifecycle.start(0).complete();
        // The recorded priority survives completion.
        assert_eq!(
            registry.get_metadata(task_id).unwrap().priority,
            Priority::Critical
        );
    }

    #[test]
    fn cancel_before_start_skips_body_and_completes_as_cancelled() {
        let before = Instant::now();
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();

        assert_eq!(
            registry.request_cancel(task_id),
            Some(CancelOutcome::Requested)
        );

        // The job-start gate observes the request and never yields a running token.
        assert!(lifecycle.start_unless_cancelled(4).is_none());

        let metadata = registry.get_metadata(task_id).unwrap();
        assert!(metadata.cancelled);
        assert!(
            metadata.completed_at >= Some(before),
            "cancellation must stamp a completion time from this test's window"
        );
        assert!(metadata.started_at.is_none());
        assert!(registry.is_completed(task_id));
    }

    #[test]
    fn cancel_after_completion_reports_already_completed() {
        let registry = TaskRegistry::new();
        let (task_id, lifecycle) = registry.register_next_task();
        lifecycle.start(0).complete();

        assert_eq!(
            registry.request_cancel(task_id),
            Some(CancelOutcome::AlreadyCompleted)
        );
        assert!(!registry.get_metadata(task_id).unwrap().cancelled);
    }

    #[test]
    fn cancel_unknown_task_returns_none() {
        let registry = TaskRegistry::new();
        assert_eq!(registry.request_cancel(999), None);
    }

    /// Concurrent registration must hand out distinct ids and distinct slots.
    ///
    /// The registry is shared as `Arc<TaskRegistry>` without an outer mutex, so
    /// this is the contract that replaced it: many producers registering at
    /// once, across the 1024-id block boundary so the directory has to grow
    /// underneath them. A lost `fetch_add`, a torn directory growth, or two
    /// producers claiming one slot would show up as a duplicate id or a
    /// missing metadata entry.
    #[test]
    fn concurrent_registration_yields_distinct_ids_across_block_growth() {
        use std::collections::BTreeSet;
        use std::sync::Mutex;

        const PRODUCERS: usize = 8;
        const PER_PRODUCER: usize = 400; // 3200 ids: spans several 1024-id blocks

        let registry = Arc::new(TaskRegistry::new());
        let ids: Mutex<Vec<u64>> = Mutex::new(Vec::with_capacity(PRODUCERS * PER_PRODUCER));

        std::thread::scope(|scope| {
            for _ in 0..PRODUCERS {
                let registry = Arc::clone(&registry);
                let ids = &ids;
                scope.spawn(move || {
                    let mut local = Vec::with_capacity(PER_PRODUCER);
                    for _ in 0..PER_PRODUCER {
                        local.push(registry.register_task());
                    }
                    ids.lock().expect("test id collector").extend(local);
                });
            }
        });

        let ids = ids.into_inner().expect("test id collector");
        assert_eq!(ids.len(), PRODUCERS * PER_PRODUCER);
        let unique: BTreeSet<u64> = ids.iter().copied().collect();
        assert_eq!(
            unique.len(),
            ids.len(),
            "concurrent registration handed out a duplicate task id"
        );
        for id in ids {
            assert!(
                registry.get_metadata(id).is_some(),
                "registered task {id} has no slot"
            );
        }
    }

    /// A waker registered against an already-completed task must not be kept.
    ///
    /// `wait_for_task` documents the race this covers: it checks completion,
    /// registers a waker, then re-checks. If the task completes between the
    /// check and the registration, completion has already taken the (absent)
    /// waker, so nothing will ever take the one just stored — it and whatever
    /// the waker holds, typically an `Arc` to async task state, stay alive for
    /// the life of the registry slot, which is never reclaimed.
    #[test]
    fn waker_registered_after_completion_is_not_retained() {
        use std::sync::atomic::{AtomicUsize, Ordering};
        use std::task::{RawWaker, RawWakerVTable, Waker};

        static DROPPED: AtomicUsize = AtomicUsize::new(0);

        // A waker whose payload is an Arc, so retention is observable through
        // the strong count rather than inferred.
        unsafe fn clone_raw(data: *const ()) -> RawWaker {
            let arc = unsafe { Arc::from_raw(data.cast::<usize>()) };
            let cloned = Arc::clone(&arc);
            core::mem::forget(arc);
            RawWaker::new(Arc::into_raw(cloned).cast(), &VTABLE)
        }
        unsafe fn wake_raw(data: *const ()) {
            drop(unsafe { Arc::from_raw(data.cast::<usize>()) });
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
        unsafe fn wake_by_ref_raw(_: *const ()) {}
        unsafe fn drop_raw(data: *const ()) {
            drop(unsafe { Arc::from_raw(data.cast::<usize>()) });
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
        static VTABLE: RawWakerVTable =
            RawWakerVTable::new(clone_raw, wake_raw, wake_by_ref_raw, drop_raw);

        let payload = Arc::new(0usize);
        let probe = Arc::clone(&payload);
        let waker = unsafe { Waker::from_raw(RawWaker::new(Arc::into_raw(probe).cast(), &VTABLE)) };

        let registry = TaskRegistry::new();
        let id = registry.register_task();
        registry.mark_started(id, 0);
        registry.mark_completed(id);

        assert!(
            registry.register_waker(id, &waker),
            "the slot is still registered, so registration reports success"
        );
        drop(waker);

        assert_eq!(
            Arc::strong_count(&payload),
            1,
            "a waker registered after completion is still held by the registry slot;              nothing will ever take it, so it leaks for the life of the slot"
        );
    }
}