grafos_sync/

barrier.rs

//! Distributed barrier for multi-party phase synchronization.
//!
//! A [`FabricBarrier`] blocks each participant at a synchronization point
//! until all `num_parties` have arrived. On the last arrival, the generation
//! counter advances and the arrival count resets, making the barrier reusable
//! across successive phases.
//!
//! # Memory layout at `base_offset`
//!
//! ```text
//! offset +0:  generation   (u64, 8 bytes)  — incremented when all parties arrive
//! offset +8:  count        (u64, 8 bytes)  — number of parties that have called wait()
//! offset +16: num_parties  (u32, 4 bytes)  — total parties expected
//! ```
//!
//! Total header: 20 bytes.

use grafos_std::error::{FabricError, Result};
use grafos_std::mem::MemLease;

const GENERATION_OFFSET: u64 = 0;
const COUNT_OFFSET: u64 = 8;
const NUM_PARTIES_OFFSET: u64 = 16;

/// Distributed barrier backed by a fabric memory lease.
///
/// All `num_parties` participants must call [`wait`](FabricBarrier::wait)
/// before any of them proceed. After all parties arrive, the generation
/// counter increments and the barrier resets for reuse across successive
/// synchronization rounds.
///
/// If a party crashes, the lease eventually expires and remaining parties
/// will observe a timeout ([`FabricError::LeaseExpired`]).
///
/// # Example
///
/// ```rust
/// # grafos_std::host::reset_mock();
/// # grafos_std::host::mock_set_fbmu_arena_size(65536);
/// use grafos_sync::FabricBarrier;
/// use grafos_std::mem::MemBuilder;
///
/// let lease = MemBuilder::new().acquire().unwrap();
/// let barrier = FabricBarrier::new(lease, 0, 1).unwrap();
///
/// // With 1 party, wait() completes immediately
/// let gen = barrier.wait(10).unwrap();
/// assert_eq!(gen, 1);
/// ```
pub struct FabricBarrier {
    lease: MemLease,
    base_offset: u64,
    num_parties: u32,
}

fn read_u64(mem: &grafos_std::mem::FabricMem, offset: u64) -> Result<u64> {
    let data = mem.read(offset, 8)?;
    if data.len() < 8 {
        return Ok(0);
    }
    let mut buf = [0u8; 8];
    buf.copy_from_slice(&data[..8]);
    Ok(u64::from_le_bytes(buf))
}

fn write_u64(mem: &grafos_std::mem::FabricMem, offset: u64, val: u64) -> Result<()> {
    mem.write(offset, &val.to_le_bytes())
}

fn write_u32(mem: &grafos_std::mem::FabricMem, offset: u64, val: u32) -> Result<()> {
    mem.write(offset, &val.to_le_bytes())
}

impl FabricBarrier {
    /// Create a new barrier for `num_parties` participants.
    ///
    /// Initializes the barrier state in leased memory at `base_offset`.
    pub fn new(lease: MemLease, base_offset: u64, num_parties: u32) -> Result<Self> {
        let mem = lease.mem();
        write_u64(mem, base_offset + GENERATION_OFFSET, 0)?;
        write_u64(mem, base_offset + COUNT_OFFSET, 0)?;
        write_u32(mem, base_offset + NUM_PARTIES_OFFSET, num_parties)?;

        Ok(FabricBarrier {
            lease,
            base_offset,
            num_parties,
        })
    }

    /// Wait at the barrier.
    ///
    /// Increments the arrival count. If this is the last party to arrive,
    /// the generation is incremented and the count is reset to zero. Otherwise,
    /// polls until the generation changes, indicating all parties have arrived.
    ///
    /// Returns the new generation number on success.
    ///
    /// `max_polls` limits how many polling iterations before returning
    /// [`FabricError::LeaseExpired`].
    pub fn wait(&self, max_polls: u32) -> Result<u64> {
        let mem = self.lease.mem();
        let base = self.base_offset;

        let gen_before = read_u64(mem, base + GENERATION_OFFSET)?;
        let count = read_u64(mem, base + COUNT_OFFSET)?;
        let new_count = count + 1;

        if new_count >= self.num_parties as u64 {
            // We are the last party — advance generation, reset count
            write_u64(mem, base + GENERATION_OFFSET, gen_before + 1)?;
            write_u64(mem, base + COUNT_OFFSET, 0)?;
            return Ok(gen_before + 1);
        }

        // Not the last — write our arrival and wait for generation to change
        write_u64(mem, base + COUNT_OFFSET, new_count)?;

        for _ in 0..max_polls {
            let gen_now = read_u64(mem, base + GENERATION_OFFSET)?;
            if gen_now != gen_before {
                return Ok(gen_now);
            }
        }

        Err(FabricError::LeaseExpired)
    }

    /// Read the current generation counter.
    ///
    /// The generation starts at 0 and increments by one each time all
    /// `num_parties` have arrived and the barrier completes a round.
    pub fn generation(&self) -> Result<u64> {
        read_u64(self.lease.mem(), self.base_offset + GENERATION_OFFSET)
    }

    /// Returns the lease ID of the underlying memory lease for external
    /// renewal management (e.g. via [`grafos_leasekit::RenewalManager`]).
    pub fn lease_id(&self) -> u128 {
        self.lease.lease_id()
    }

    /// Returns the expiry time (unix seconds) of the underlying memory
    /// lease for external renewal management.
    pub fn expires_at_unix_secs(&self) -> u64 {
        self.lease.expires_at_unix_secs()
    }

    /// Read the current arrival count.
    ///
    /// Returns the number of parties that have called [`wait`](Self::wait) in
    /// the current generation. Resets to 0 when the last party arrives and the
    /// generation advances.
    pub fn count(&self) -> Result<u64> {
        read_u64(self.lease.mem(), self.base_offset + COUNT_OFFSET)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use grafos_std::host;
    use grafos_std::mem::MemBuilder;

    fn setup() -> MemLease {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(4096);
        MemBuilder::new().acquire().expect("acquire lease")
    }

    #[test]
    fn barrier_three_parties_sequential() {
        let lease = setup();
        let barrier = FabricBarrier::new(lease, 0, 3).expect("new barrier");

        assert_eq!(barrier.generation().unwrap(), 0);
        assert_eq!(barrier.count().unwrap(), 0);

        // Party 1 arrives — count goes to 1, waits for generation change
        // Since we're sequential, party 1 would block. Instead, we simulate
        // by directly calling wait for all three in sequence.
        // Party 1: count 0 → 1, not last, starts polling
        // But since we're single-threaded, we need to simulate this differently.
        //
        // In the sequential mock, we call wait() three times. The first two
        // increment the count, and the third sees count == num_parties and
        // advances the generation. The first two will poll and see the
        // generation has changed (because the third call already ran).
        //
        // To make this work in single-threaded tests, we simulate the
        // arrivals manually and then call wait for the last party.

        // Manually set count to 2 (simulating 2 parties already arrived)
        let mem = barrier.lease.mem();
        write_u64(mem, COUNT_OFFSET, 2).unwrap();

        // Third party calls wait — should complete immediately
        let gen = barrier.wait(10).expect("wait");
        assert_eq!(gen, 1);
        assert_eq!(barrier.generation().unwrap(), 1);
        assert_eq!(barrier.count().unwrap(), 0); // reset after all arrive
    }

    #[test]
    fn barrier_reuse_across_generations() {
        let lease = setup();
        let barrier = FabricBarrier::new(lease, 0, 2).expect("new barrier");

        // Generation 0: simulate first party arrived
        let mem = barrier.lease.mem();
        write_u64(mem, COUNT_OFFSET, 1).unwrap();

        // Second party arrives — completes generation 0
        let gen = barrier.wait(10).expect("wait gen 0");
        assert_eq!(gen, 1);
        assert_eq!(barrier.count().unwrap(), 0);

        // Generation 1: simulate first party arrived again
        write_u64(mem, barrier.base_offset + COUNT_OFFSET, 1).unwrap();

        // Second party arrives — completes generation 1
        let gen = barrier.wait(10).expect("wait gen 1");
        assert_eq!(gen, 2);
        assert_eq!(barrier.generation().unwrap(), 2);
    }

    #[test]
    fn barrier_timeout_when_not_all_arrive() {
        let lease = setup();
        let barrier = FabricBarrier::new(lease, 0, 3).expect("new barrier");

        // Only 1 party arrives (count stays at 0 before wait increments to 1)
        // wait will poll but generation never changes
        let result = barrier.wait(5);
        assert_eq!(result.unwrap_err(), FabricError::LeaseExpired);
    }
}