grafos_jobs/

coordinator.rs

//! Job coordinator — orchestrates chunk execution with retries and aggregation.

use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;
use core::fmt;

use grafos_std::error::FabricError;

use crate::output_store::JobOutputStore;
use crate::retry::{RetryPolicy, RetryableError};
use crate::work_chunk::{ChunkId, WorkChunk};

/// Status of a single chunk after execution.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChunkStatus {
    /// Chunk completed successfully.
    Success,
    /// Chunk failed permanently after exhausting retries.
    Failed(String),
    /// Chunk was skipped (output already in store from prior run).
    Skipped,
}

/// Per-chunk execution result.
#[derive(Debug, Clone)]
pub struct ChunkResult {
    /// Final status of this chunk.
    pub status: ChunkStatus,
    /// Number of retries consumed (0 = succeeded on first attempt).
    pub retries_used: u32,
}

/// Aggregate result of a job run.
pub struct JobResult {
    /// Per-chunk results keyed by `ChunkId`.
    pub chunk_results: BTreeMap<ChunkId, ChunkResult>,
    /// Aggregated output from the reducer.
    pub aggregate: Vec<u8>,
    /// Number of chunks that succeeded.
    pub succeeded: u32,
    /// Number of chunks that failed permanently.
    pub failed: u32,
    /// Number of chunks skipped (idempotent — already had output).
    pub skipped: u32,
}

impl fmt::Debug for JobResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("JobResult")
            .field("succeeded", &self.succeeded)
            .field("failed", &self.failed)
            .field("skipped", &self.skipped)
            .field("aggregate_len", &self.aggregate.len())
            .finish()
    }
}

/// Orchestrates chunk execution with retries and idempotent output capture.
///
/// The coordinator iterates over work chunks, executes each via a
/// caller-provided function, stores outputs in a [`JobOutputStore`],
/// and retries transient failures according to [`RetryPolicy`]. Once
/// all chunks are processed, the caller-provided reducer aggregates
/// the outputs.
///
/// Chunks whose output already exists in the store are skipped
/// (idempotent replay). This enables safe restart after partial failure.
pub struct JobCoordinator {
    policy: RetryPolicy,
}

impl JobCoordinator {
    /// Create a new coordinator with the given retry policy.
    pub fn new(policy: RetryPolicy) -> Self {
        JobCoordinator { policy }
    }

    /// Execute a job over the given chunks.
    ///
    /// # Parameters
    ///
    /// - `chunks`: Work units to process.
    /// - `store`: Output store for idempotent capture.
    /// - `execute_fn`: Called with serialized chunk bytes; returns output
    ///   bytes on success.
    /// - `reduce_fn`: Called with all `(ChunkId, output_bytes)` pairs to
    ///   produce the final aggregate.
    ///
    /// # Errors
    ///
    /// Returns `FabricError` only on store failures. Chunk-level failures
    /// are captured in [`JobResult::chunk_results`]; the job itself
    /// succeeds as long as the store is functional.
    pub fn run<F, R>(
        &mut self,
        chunks: &[Box<dyn WorkChunk>],
        store: &mut dyn JobOutputStore,
        mut execute_fn: F,
        reduce_fn: R,
    ) -> Result<JobResult, FabricError>
    where
        F: FnMut(&[u8]) -> Result<Vec<u8>, FabricError>,
        R: FnOnce(&[(ChunkId, Vec<u8>)]) -> Vec<u8>,
    {
        let mut chunk_results: BTreeMap<ChunkId, ChunkResult> = BTreeMap::new();
        let mut succeeded: u32 = 0;
        let mut failed: u32 = 0;
        let mut skipped: u32 = 0;

        for chunk in chunks {
            let cid = chunk.chunk_id();

            // Idempotent: skip if output already captured.
            if store.contains(cid) {
                chunk_results.insert(
                    cid,
                    ChunkResult {
                        status: ChunkStatus::Skipped,
                        retries_used: 0,
                    },
                );
                skipped += 1;
                continue;
            }

            let chunk_bytes = chunk.to_bytes();
            let max_attempts = self.policy.max_retries + 1;
            let mut last_err: Option<FabricError> = None;
            let mut retries_used: u32 = 0;
            let mut _backoff = self.policy.backoff();

            for attempt in 0..max_attempts {
                match execute_fn(&chunk_bytes) {
                    Ok(output) => {
                        store.put(cid, output)?;
                        retries_used = attempt;
                        last_err = None;
                        break;
                    }
                    Err(e) => {
                        retries_used = attempt;
                        // Permanent errors are not retried.
                        if self.policy.classify(&e) == RetryableError::Permanent {
                            last_err = Some(e);
                            break;
                        }
                        last_err = Some(e);
                        // Backoff: in a real system the coordinator would
                        // sleep for _backoff.next_delay() seconds. In this
                        // synchronous implementation we just advance the
                        // backoff state for bookkeeping.
                        let _delay = _backoff.next_delay();
                    }
                }
            }

            if let Some(e) = last_err {
                chunk_results.insert(
                    cid,
                    ChunkResult {
                        status: ChunkStatus::Failed(alloc::format!("{}", e)),
                        retries_used,
                    },
                );
                failed += 1;
            } else {
                chunk_results.insert(
                    cid,
                    ChunkResult {
                        status: ChunkStatus::Success,
                        retries_used,
                    },
                );
                succeeded += 1;
            }
        }

        // Gather all successful outputs for reduction.
        let mut outputs: Vec<(ChunkId, Vec<u8>)> = Vec::new();
        for chunk in chunks {
            let cid = chunk.chunk_id();
            if let Some(data) = store.get(cid)? {
                outputs.push((cid, data));
            }
        }

        let aggregate = reduce_fn(&outputs);

        Ok(JobResult {
            chunk_results,
            aggregate,
            succeeded,
            failed,
            skipped,
        })
    }

    /// Teardown: clear the output store.
    pub fn teardown(&self, store: &mut dyn JobOutputStore) {
        store.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{MemoryOutputStore, RetryPolicy};

    #[derive(Clone, serde::Serialize, serde::Deserialize)]
    struct TestChunk {
        id: u64,
        value: u64,
    }

    impl WorkChunk for TestChunk {
        fn chunk_id(&self) -> ChunkId {
            ChunkId(self.id)
        }
        fn to_bytes(&self) -> Vec<u8> {
            postcard::to_allocvec(self).unwrap()
        }
        fn from_bytes(bytes: &[u8]) -> Result<Self, FabricError> {
            postcard::from_bytes(bytes).map_err(|_| FabricError::IoError(-200))
        }
    }

    fn make_chunks(values: &[u64]) -> Vec<Box<dyn WorkChunk>> {
        values
            .iter()
            .enumerate()
            .map(|(i, &v)| -> Box<dyn WorkChunk> {
                Box::new(TestChunk {
                    id: i as u64,
                    value: v,
                })
            })
            .collect()
    }

    fn square_execute(bytes: &[u8]) -> Result<Vec<u8>, FabricError> {
        let chunk: TestChunk =
            postcard::from_bytes(bytes).map_err(|_| FabricError::IoError(-200))?;
        let result = chunk.value * chunk.value;
        Ok(postcard::to_allocvec(&result).unwrap())
    }

    fn sum_reduce(outputs: &[(ChunkId, Vec<u8>)]) -> Vec<u8> {
        let sum: u64 = outputs
            .iter()
            .map(|(_, v)| postcard::from_bytes::<u64>(v).unwrap_or(0))
            .sum();
        postcard::to_allocvec(&sum).unwrap()
    }

    #[test]
    fn basic_map_reduce() {
        let chunks = make_chunks(&[2, 3, 4]);
        let mut store = MemoryOutputStore::new();
        let mut coord = JobCoordinator::new(RetryPolicy::default());

        let result = coord
            .run(&chunks, &mut store, square_execute, sum_reduce)
            .unwrap();

        assert_eq!(result.succeeded, 3);
        assert_eq!(result.failed, 0);
        assert_eq!(result.skipped, 0);

        let total: u64 = postcard::from_bytes(&result.aggregate).unwrap();
        assert_eq!(total, 4 + 9 + 16); // 2^2 + 3^2 + 4^2
    }

    #[test]
    fn idempotent_skip_on_retry() {
        let chunks = make_chunks(&[5]);
        let mut store = MemoryOutputStore::new();

        // Pre-populate the store as if a prior run completed this chunk.
        let expected: u64 = 25;
        store
            .put(ChunkId(0), postcard::to_allocvec(&expected).unwrap())
            .unwrap();

        let mut coord = JobCoordinator::new(RetryPolicy::default());
        let result = coord
            .run(&chunks, &mut store, square_execute, sum_reduce)
            .unwrap();

        assert_eq!(result.succeeded, 0);
        assert_eq!(result.skipped, 1);
        assert_eq!(result.failed, 0);

        let total: u64 = postcard::from_bytes(&result.aggregate).unwrap();
        assert_eq!(total, 25);
    }

    #[test]
    fn retry_on_transient_failure_converges() {
        use std::cell::Cell;
        use std::rc::Rc;

        let fail_count = Rc::new(Cell::new(0u32));
        let fail_clone = fail_count.clone();

        let chunks = make_chunks(&[7]);
        let mut store = MemoryOutputStore::new();
        let mut coord = JobCoordinator::new(RetryPolicy {
            max_retries: 3,
            initial_backoff_secs: 1,
            max_backoff_secs: 4,
        });

        let result = coord
            .run(
                &chunks,
                &mut store,
                |bytes| {
                    let count = fail_clone.get();
                    fail_clone.set(count + 1);
                    if count < 2 {
                        // Fail with a transient error twice
                        Err(FabricError::Disconnected)
                    } else {
                        square_execute(bytes)
                    }
                },
                sum_reduce,
            )
            .unwrap();

        assert_eq!(result.succeeded, 1);
        assert_eq!(result.failed, 0);
        assert_eq!(fail_count.get(), 3); // 2 failures + 1 success

        let cr = result.chunk_results.get(&ChunkId(0)).unwrap();
        assert_eq!(cr.status, ChunkStatus::Success);
        assert_eq!(cr.retries_used, 2);

        let total: u64 = postcard::from_bytes(&result.aggregate).unwrap();
        assert_eq!(total, 49);
    }

    #[test]
    fn permanent_failure_not_retried() {
        use std::cell::Cell;
        use std::rc::Rc;

        let call_count = Rc::new(Cell::new(0u32));
        let call_clone = call_count.clone();

        let chunks = make_chunks(&[1]);
        let mut store = MemoryOutputStore::new();
        let mut coord = JobCoordinator::new(RetryPolicy {
            max_retries: 5,
            ..RetryPolicy::default()
        });

        let result = coord
            .run(
                &chunks,
                &mut store,
                |_bytes| {
                    call_clone.set(call_clone.get() + 1);
                    Err(FabricError::Fenced) // permanent
                },
                sum_reduce,
            )
            .unwrap();

        assert_eq!(result.failed, 1);
        assert_eq!(result.succeeded, 0);
        // Should be called exactly once — no retries for permanent errors.
        assert_eq!(call_count.get(), 1);

        let cr = result.chunk_results.get(&ChunkId(0)).unwrap();
        assert!(matches!(cr.status, ChunkStatus::Failed(_)));
    }

    #[test]
    fn retries_exhausted() {
        let chunks = make_chunks(&[1]);
        let mut store = MemoryOutputStore::new();
        let mut coord = JobCoordinator::new(RetryPolicy {
            max_retries: 2,
            initial_backoff_secs: 1,
            max_backoff_secs: 4,
        });

        let result = coord
            .run(
                &chunks,
                &mut store,
                |_bytes| Err(FabricError::LeaseExpired), // always transient
                sum_reduce,
            )
            .unwrap();

        assert_eq!(result.failed, 1);
        assert_eq!(result.succeeded, 0);

        let cr = result.chunk_results.get(&ChunkId(0)).unwrap();
        assert!(matches!(cr.status, ChunkStatus::Failed(_)));
        assert_eq!(cr.retries_used, 2); // last attempt index
    }

    #[test]
    fn determinism_same_inputs_same_outputs() {
        use std::cell::Cell;
        use std::rc::Rc;

        // Run twice with the same transient failure pattern.
        for _ in 0..2 {
            let fail_once = Rc::new(Cell::new(false));
            let fail_clone = fail_once.clone();

            let chunks = make_chunks(&[10, 20]);
            let mut store = MemoryOutputStore::new();
            let mut coord = JobCoordinator::new(RetryPolicy {
                max_retries: 2,
                initial_backoff_secs: 1,
                max_backoff_secs: 4,
            });

            let result = coord
                .run(
                    &chunks,
                    &mut store,
                    |bytes| {
                        // Fail the first call once, then succeed.
                        if !fail_clone.get() {
                            fail_clone.set(true);
                            Err(FabricError::Disconnected)
                        } else {
                            square_execute(bytes)
                        }
                    },
                    sum_reduce,
                )
                .unwrap();

            // chunk 0 (value=10): fails once, retries, succeeds => 100
            // chunk 1 (value=20): succeeds immediately => 400
            let total: u64 = postcard::from_bytes(&result.aggregate).unwrap();
            assert_eq!(total, 100 + 400);
            assert_eq!(result.succeeded, 2);
        }
    }

    #[test]
    fn teardown_clears_store() {
        let chunks = make_chunks(&[3]);
        let mut store = MemoryOutputStore::new();
        let mut coord = JobCoordinator::new(RetryPolicy::default());

        let _ = coord
            .run(&chunks, &mut store, square_execute, sum_reduce)
            .unwrap();

        assert!(store.contains(ChunkId(0)));
        coord.teardown(&mut store);
        assert!(!store.contains(ChunkId(0)));
    }

    #[test]
    fn empty_job() {
        let chunks: Vec<Box<dyn WorkChunk>> = vec![];
        let mut store = MemoryOutputStore::new();
        let mut coord = JobCoordinator::new(RetryPolicy::default());

        let result = coord
            .run(&chunks, &mut store, square_execute, |_outputs| Vec::new())
            .unwrap();

        assert_eq!(result.succeeded, 0);
        assert_eq!(result.failed, 0);
        assert_eq!(result.skipped, 0);
        assert!(result.aggregate.is_empty());
    }

    #[test]
    fn mixed_success_and_failure() {
        use std::cell::Cell;
        use std::rc::Rc;

        let call_idx = Rc::new(Cell::new(0u32));
        let call_clone = call_idx.clone();

        let chunks = make_chunks(&[2, 3, 4]);
        let mut store = MemoryOutputStore::new();
        let mut coord = JobCoordinator::new(RetryPolicy {
            max_retries: 0,
            ..RetryPolicy::default()
        });

        let result = coord
            .run(
                &chunks,
                &mut store,
                |bytes| {
                    let idx = call_clone.get();
                    call_clone.set(idx + 1);
                    if idx == 1 {
                        // Chunk 1 (value=3) fails permanently
                        Err(FabricError::Fenced)
                    } else {
                        square_execute(bytes)
                    }
                },
                sum_reduce,
            )
            .unwrap();

        assert_eq!(result.succeeded, 2);
        assert_eq!(result.failed, 1);

        // Only successful chunks contribute to aggregate.
        let total: u64 = postcard::from_bytes(&result.aggregate).unwrap();
        assert_eq!(total, 4 + 16); // 2^2 + 4^2
    }
}