grafos_sync/

once.rs

//! One-time initialization of shared state in fabric memory.
//!
//! [`FabricOnce`] ensures a value is computed and stored exactly once in
//! leased memory, even if multiple callers race to initialize it. The
//! first writer wins; subsequent callers read the existing value.
//!
//! # Memory layout at `base_offset`
//!
//! ```text
//! offset +0:  initialized  (u8, 1 byte)   — 0 = not initialized, 1 = initialized
//! offset +1:  data_len     (u32, 4 bytes)  — byte length of stored data
//! offset +5:  data         ([u8])          — the initialized value
//! ```
//!
//! Total header: 5 bytes before the user data.

extern crate alloc;
use alloc::vec::Vec;

use grafos_std::error::Result;
use grafos_std::mem::MemLease;

const INITIALIZED_OFFSET: u64 = 0;
const DATA_LEN_OFFSET: u64 = 1;
const DATA_OFFSET: u64 = 5;

/// One-time initialization primitive backed by a fabric memory lease.
///
/// Ensures a value is computed and stored exactly once, even if multiple
/// callers race to initialize it. The first writer wins; subsequent
/// callers read the existing value.
///
/// # Example
///
/// ```rust
/// # grafos_std::host::reset_mock();
/// # grafos_std::host::mock_set_fbmu_arena_size(65536);
/// use grafos_sync::FabricOnce;
/// use grafos_std::mem::MemBuilder;
///
/// let lease = MemBuilder::new().acquire().unwrap();
/// let once = FabricOnce::new(lease, 0).unwrap();
///
/// let val = once.call_once(|| b"computed".to_vec()).unwrap();
/// assert_eq!(&val, b"computed");
///
/// // Second call returns the stored value without recomputing
/// let val2 = once.call_once(|| b"ignored".to_vec()).unwrap();
/// assert_eq!(&val2, b"computed");
/// ```
pub struct FabricOnce {
    lease: MemLease,
    base_offset: u64,
}

impl FabricOnce {
    /// Create a new `FabricOnce` at `base_offset` in the given lease.
    ///
    /// Initializes the `initialized` flag to 0 (not yet initialized).
    pub fn new(lease: MemLease, base_offset: u64) -> Result<Self> {
        let mem = lease.mem();
        mem.write(base_offset + INITIALIZED_OFFSET, &[0u8])?;
        mem.write(base_offset + DATA_LEN_OFFSET, &0u32.to_le_bytes())?;

        Ok(FabricOnce { lease, base_offset })
    }

    /// Initialize the value by calling `f`, or return the existing value
    /// if already initialized.
    ///
    /// If the `initialized` flag is already set, `f` is **not** called and
    /// the stored value is returned. Otherwise, `f` is called, the result
    /// is written to leased memory (data first, then length, then flag),
    /// and the value is returned.
    pub fn call_once<F>(&self, f: F) -> Result<Vec<u8>>
    where
        F: FnOnce() -> Vec<u8>,
    {
        let mem = self.lease.mem();
        let base = self.base_offset;

        // Check if already initialized
        let flag = mem.read(base + INITIALIZED_OFFSET, 1)?;
        if !flag.is_empty() && flag[0] != 0 {
            return self.read_stored();
        }

        // Not initialized — compute the value
        let data = f();
        let data_len = data.len() as u32;

        // Write data first, then length, then flag (ordered for crash safety)
        mem.write(base + DATA_OFFSET, &data)?;
        mem.write(base + DATA_LEN_OFFSET, &data_len.to_le_bytes())?;
        mem.write(base + INITIALIZED_OFFSET, &[1u8])?;

        // Double-check: if someone else initialized first, read their value
        let flag = mem.read(base + INITIALIZED_OFFSET, 1)?;
        if !flag.is_empty() && flag[0] == 1 {
            return self.read_stored();
        }

        Ok(data)
    }

    /// 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()
    }

    /// Check if the value has been initialized.
    ///
    /// Returns `true` if [`call_once`](Self::call_once) has completed at least
    /// once (the `initialized` flag is set in leased memory).
    pub fn is_initialized(&self) -> Result<bool> {
        let flag = self
            .lease
            .mem()
            .read(self.base_offset + INITIALIZED_OFFSET, 1)?;
        Ok(!flag.is_empty() && flag[0] != 0)
    }

    /// Read the stored value (must be initialized).
    fn read_stored(&self) -> Result<Vec<u8>> {
        let mem = self.lease.mem();
        let base = self.base_offset;

        let len_bytes = mem.read(base + DATA_LEN_OFFSET, 4)?;
        if len_bytes.len() < 4 {
            return Ok(Vec::new());
        }
        let mut buf = [0u8; 4];
        buf.copy_from_slice(&len_bytes[..4]);
        let data_len = u32::from_le_bytes(buf);

        if data_len == 0 {
            return Ok(Vec::new());
        }

        mem.read(base + DATA_OFFSET, data_len)
    }
}

#[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 once_initializes_exactly_once() {
        let lease = setup();
        let once = FabricOnce::new(lease, 0).expect("new once");

        assert!(!once.is_initialized().unwrap());

        let mut call_count = 0u32;
        let val1 = once
            .call_once(|| {
                call_count += 1;
                b"hello fabric".to_vec()
            })
            .expect("call_once 1");
        assert_eq!(call_count, 1);
        assert_eq!(&val1, b"hello fabric");
        assert!(once.is_initialized().unwrap());

        // Second call should return existing value without calling f
        let val2 = once
            .call_once(|| {
                call_count += 1;
                b"should not happen".to_vec()
            })
            .expect("call_once 2");
        assert_eq!(call_count, 1); // f was NOT called again
        assert_eq!(&val2, b"hello fabric");
    }

    #[test]
    fn once_empty_value() {
        let lease = setup();
        let once = FabricOnce::new(lease, 0).expect("new once");

        let val = once.call_once(Vec::new).expect("call_once");
        assert!(val.is_empty());
        assert!(once.is_initialized().unwrap());

        // Second call returns empty
        let val2 = once.call_once(|| b"nope".to_vec()).expect("call_once 2");
        assert!(val2.is_empty());
    }
}