grafos_std/
mem.rs

1//! Safe wrappers for FBMU (Fabric Bootstrap Memory Unit) host functions.
2//!
3//! This module provides byte-addressable fabric memory access through the
4//! FBMU data-plane protocol. The primary type is [`FabricMem`], which wraps
5//! the raw host function calls with a safe Rust API.
6//!
7//! # Usage
8//!
9//! ```rust
10//! use grafos_std::mem::FabricMem;
11//!
12//! # grafos_std::host::reset_mock();
13//! # grafos_std::host::mock_set_fbmu_arena_size(4096);
14//! let mem = FabricMem::hello()?;
15//! mem.write(0, b"hello fabric")?;
16//! let data = mem.read(0, 12)?;
17//! assert_eq!(&data, b"hello fabric");
18//! # Ok::<(), grafos_std::FabricError>(())
19//! ```
20//!
21//! For capacity-checked allocation, use [`MemBuilder`]:
22//!
23//! ```rust
24//! use grafos_std::mem::MemBuilder;
25//!
26//! # grafos_std::host::reset_mock();
27//! # grafos_std::host::mock_set_fbmu_arena_size(8192);
28//! let lease = MemBuilder::new().min_bytes(4096).acquire()?;
29//! lease.mem().write(0, &[42; 100])?;
30//! # Ok::<(), grafos_std::FabricError>(())
31//! ```
32
33extern crate alloc;
34use alloc::vec;
35use alloc::vec::Vec;
36
37use crate::error::{FabricError, Result};
38use crate::host;
39use crate::lease::{self, LeaseInfo, LeaseStatus, SharedLeaseState};
40
41const LEASE_TAG_MEM: u8 = 0x01;
42
43fn lease_status_from_host(code: u8) -> LeaseStatus {
44    match code {
45        host::LEASE_STATUS_ACTIVE => LeaseStatus::Active,
46        host::LEASE_STATUS_EXPIRED => LeaseStatus::Expired,
47        host::LEASE_STATUS_REVOKED => LeaseStatus::Revoked,
48        _ => LeaseStatus::Revoked,
49    }
50}
51
52fn sync_state_from_host(state: &SharedLeaseState, info: &host::FbmuLeaseInfo) {
53    lease::set_expires_at_unix_secs(state, info.expires_at_unix_secs);
54    lease::set_status(state, lease_status_from_host(info.status));
55}
56
57/// Safe handle to fabric memory via FBMU host functions.
58///
59/// Created by [`FabricMem::hello`], which performs the FBMU HELLO handshake
60/// to establish a memory data-plane session. Once created, the handle can
61/// be used for byte-addressable reads and writes within the arena.
62///
63/// # Examples
64///
65/// ```rust
66/// use grafos_std::mem::FabricMem;
67///
68/// # grafos_std::host::reset_mock();
69/// # grafos_std::host::mock_set_fbmu_arena_size(4096);
70/// let mem = FabricMem::hello()?;
71///
72/// // Write and read back
73/// mem.write(0, b"test data")?;
74/// let data = mem.read(0, 9)?;
75/// assert_eq!(&data, b"test data");
76///
77/// // Check arena capacity
78/// let size = mem.arena_size()?;
79/// assert_eq!(size, 4096);
80/// # Ok::<(), grafos_std::FabricError>(())
81/// ```
82#[derive(Debug)]
83pub struct FabricMem {
84    lease_state: Option<SharedLeaseState>,
85    host_lease_id: Option<u128>,
86}
87
88impl FabricMem {
89    /// Perform the FBMU HELLO handshake and return a memory handle.
90    ///
91    /// This establishes the memory data-plane session with the host. On
92    /// WASM targets, this calls the real `fbmu_hello` host import. On
93    /// native targets, it initializes the mock state.
94    ///
95    /// # Errors
96    ///
97    /// - [`FabricError::Disconnected`] if the host connection is unavailable.
98    /// - Other [`FabricError`] variants based on the host status code.
99    pub fn hello() -> Result<FabricMem> {
100        host::fbmu_hello()?;
101        Ok(FabricMem {
102            lease_state: None,
103            host_lease_id: None,
104        })
105    }
106
107    /// Write `data` at the given byte offset in the arena.
108    ///
109    /// # Errors
110    ///
111    /// Returns a [`FabricError`] if the host reports a write failure.
112    pub fn write(&self, offset: u64, data: &[u8]) -> Result<()> {
113        #[cfg(feature = "observe")]
114        let start = std::time::Instant::now();
115        let result = if let Some(lease_id) = self.host_lease_id {
116            if let Some(state) = &self.lease_state {
117                let info = host::fbmu_query(lease_id)?;
118                sync_state_from_host(state, &info);
119                lease::ensure_active(state)?;
120            }
121            host::fbmu_write_lease(lease_id, offset, data)
122        } else {
123            if let Some(state) = &self.lease_state {
124                lease::ensure_active(state)?;
125            }
126            host::fbmu_write(offset, data)
127        };
128        #[cfg(feature = "observe")]
129        match &result {
130            Ok(()) => {
131                let duration_us = start.elapsed().as_micros() as u64;
132                crate::observe_hooks::on_mem_write(data.len() as u64);
133                crate::observe_hooks::on_op_completed(
134                    grafos_observe::OpType::Write,
135                    duration_us,
136                    data.len() as u64,
137                );
138            }
139            Err(e) => {
140                crate::observe_hooks::on_op_error();
141                crate::observe_hooks::on_op_failed(
142                    grafos_observe::OpType::Write,
143                    &alloc::format!("{e:?}"),
144                );
145            }
146        }
147        result
148    }
149
150    /// Read `len` bytes starting at `offset` from the arena.
151    ///
152    /// Returns a `Vec<u8>` containing the data read. The returned vector
153    /// may be shorter than `len` if the host returned fewer bytes.
154    ///
155    /// # Errors
156    ///
157    /// Returns a [`FabricError`] if the host reports a read failure.
158    pub fn read(&self, offset: u64, len: u32) -> Result<Vec<u8>> {
159        #[cfg(feature = "observe")]
160        let start = std::time::Instant::now();
161        let result = if let Some(lease_id) = self.host_lease_id {
162            if let Some(state) = &self.lease_state {
163                let info = host::fbmu_query(lease_id)?;
164                sync_state_from_host(state, &info);
165                lease::ensure_active(state)?;
166            }
167            let mut buf = vec![0u8; len as usize];
168            let n = host::fbmu_read_lease(lease_id, offset, &mut buf)?;
169            buf.truncate(n);
170            Ok(buf)
171        } else {
172            if let Some(state) = &self.lease_state {
173                lease::ensure_active(state)?;
174            }
175            let mut buf = vec![0u8; len as usize];
176            let n = host::fbmu_read(offset, &mut buf)?;
177            buf.truncate(n);
178            Ok(buf)
179        };
180        #[cfg(feature = "observe")]
181        match &result {
182            Ok(buf) => {
183                let duration_us = start.elapsed().as_micros() as u64;
184                crate::observe_hooks::on_mem_read(buf.len() as u64);
185                crate::observe_hooks::on_op_completed(
186                    grafos_observe::OpType::Read,
187                    duration_us,
188                    buf.len() as u64,
189                );
190            }
191            Err(e) => {
192                crate::observe_hooks::on_op_error();
193                crate::observe_hooks::on_op_failed(
194                    grafos_observe::OpType::Read,
195                    &alloc::format!("{e:?}"),
196                );
197            }
198        }
199        result
200    }
201
202    /// Query the total arena size in bytes.
203    ///
204    /// # Errors
205    ///
206    /// Returns [`FabricError::IoError`] if the host returns a negative value.
207    pub fn arena_size(&self) -> Result<u64> {
208        if let Some(lease_id) = self.host_lease_id {
209            let info = host::fbmu_query(lease_id)?;
210            if let Some(state) = &self.lease_state {
211                sync_state_from_host(state, &info);
212            }
213            Ok(info.arena_size)
214        } else {
215            host::fbmu_get_arena_size()
216        }
217    }
218}
219
220/// A memory lease that auto-frees on drop.
221///
222/// Wraps a [`FabricMem`] handle with RAII lifecycle management. When the
223/// `MemLease` is dropped, the underlying resource will be released (once
224/// the host provides an explicit `fbmu_free()` call).
225///
226/// Created via [`MemBuilder::acquire`].
227///
228/// # Examples
229///
230/// ```rust
231/// use grafos_std::mem::MemBuilder;
232///
233/// # grafos_std::host::reset_mock();
234/// # grafos_std::host::mock_set_fbmu_arena_size(8192);
235/// let lease = MemBuilder::new().min_bytes(4096).acquire()?;
236/// lease.mem().write(0, b"leased memory")?;
237/// // lease is freed when it goes out of scope
238/// # Ok::<(), grafos_std::FabricError>(())
239/// ```
240#[derive(Debug)]
241pub struct MemLease {
242    state: SharedLeaseState,
243    mem: FabricMem,
244    _size: u64,
245}
246
247impl MemLease {
248    fn sync_from_host(&self) {
249        let Some(lease_id) = self.mem.host_lease_id else {
250            return;
251        };
252        if let Ok(info) = host::fbmu_query(lease_id) {
253            sync_state_from_host(&self.state, &info);
254        }
255    }
256
257    /// Access the underlying [`FabricMem`] handle for I/O operations.
258    pub fn mem(&self) -> &FabricMem {
259        &self.mem
260    }
261
262    /// Lease metadata snapshot (id, creation time, expiry, and status).
263    pub fn info(&self) -> LeaseInfo {
264        self.sync_from_host();
265        lease::info(&self.state)
266    }
267
268    /// Unique lease identifier.
269    pub fn lease_id(&self) -> u128 {
270        lease::lease_id(&self.state)
271    }
272
273    /// Lease creation timestamp (unix seconds).
274    pub fn created_at_unix_secs(&self) -> u64 {
275        lease::created_at_unix_secs(&self.state)
276    }
277
278    /// Lease expiry timestamp (unix seconds).
279    pub fn expires_at_unix_secs(&self) -> u64 {
280        self.sync_from_host();
281        lease::expires_at_unix_secs(&self.state)
282    }
283
284    /// Current lease status.
285    pub fn status(&self) -> LeaseStatus {
286        self.sync_from_host();
287        lease::status(&self.state)
288    }
289
290    /// Renew the lease TTL by `duration_secs`.
291    ///
292    /// Extends `expires_at` to at least `now + duration_secs`.
293    pub fn renew(&self, duration_secs: u64) -> Result<()> {
294        if let Some(lease_id) = self.mem.host_lease_id {
295            let expires_at_unix_secs = host::fbmu_renew(lease_id, duration_secs)?;
296            lease::set_expires_at_unix_secs(&self.state, expires_at_unix_secs);
297            lease::set_status(&self.state, LeaseStatus::Active);
298            return Ok(());
299        }
300        lease::renew(&self.state, duration_secs)
301    }
302
303    /// Explicitly revoke/free this lease.
304    pub fn free(&self) {
305        if let Some(lease_id) = self.mem.host_lease_id {
306            let _ = host::fbmu_free(lease_id);
307        }
308        lease::free(&self.state);
309    }
310
311    /// Create a second handle to the same underlying lease.
312    ///
313    /// The returned `MemLease` shares the same `host_lease_id` and lease
314    /// state, so reads and writes go to the same arena. This models the
315    /// real-world case where sender and receiver both hold handles to the
316    /// same remote memory region.
317    #[cfg(any(test, feature = "test-support"))]
318    pub fn dup(&self) -> MemLease {
319        MemLease {
320            state: self.state.clone(),
321            mem: FabricMem {
322                lease_state: self.mem.lease_state.clone(),
323                host_lease_id: self.mem.host_lease_id,
324            },
325            _size: self._size,
326        }
327    }
328}
329
330impl Drop for MemLease {
331    fn drop(&mut self) {
332        #[cfg(feature = "observe")]
333        crate::observe_hooks::on_lease_dropped(LEASE_TAG_MEM, lease::lease_id(&self.state));
334        lease::free(&self.state);
335    }
336}
337
338/// Builder for acquiring a fabric memory lease with capacity constraints.
339///
340/// Use the builder pattern to specify minimum capacity requirements, then
341/// call [`acquire`](MemBuilder::acquire) to perform the HELLO handshake
342/// and validate the arena size.
343///
344/// # Examples
345///
346/// ```rust
347/// use grafos_std::mem::MemBuilder;
348///
349/// # grafos_std::host::reset_mock();
350/// # grafos_std::host::mock_set_fbmu_arena_size(65536);
351/// // Require at least 8 KiB
352/// let lease = MemBuilder::new().min_bytes(8192).acquire()?;
353/// assert!(lease.mem().arena_size()? >= 8192);
354/// # Ok::<(), grafos_std::FabricError>(())
355/// ```
356pub struct MemBuilder {
357    min_bytes: u64,
358    lease_secs: u64,
359}
360
361impl MemBuilder {
362    /// Create a new builder with no minimum capacity constraint.
363    pub fn new() -> Self {
364        MemBuilder {
365            min_bytes: 0,
366            lease_secs: 300,
367        }
368    }
369
370    /// Set the minimum number of bytes required from the arena.
371    ///
372    /// If the arena reported by the host is smaller than this value,
373    /// [`acquire`](MemBuilder::acquire) will return
374    /// [`FabricError::CapacityExceeded`].
375    pub fn min_bytes(mut self, n: u64) -> Self {
376        self.min_bytes = n;
377        self
378    }
379
380    /// Set the lease TTL in seconds.
381    pub fn lease_secs(mut self, secs: u64) -> Self {
382        self.lease_secs = secs.max(1);
383        self
384    }
385
386    /// Acquire the memory lease by performing the HELLO handshake and
387    /// verifying that the arena meets the requested minimum.
388    ///
389    /// # Errors
390    ///
391    /// - [`FabricError::CapacityExceeded`] if the arena size is less than
392    ///   the configured [`min_bytes`](MemBuilder::min_bytes).
393    /// - [`FabricError::Disconnected`] if the HELLO handshake fails.
394    pub fn acquire(self) -> Result<MemLease> {
395        let result = match host::fbmu_alloc(self.min_bytes, self.lease_secs) {
396            Ok(info) => {
397                if info.arena_size < self.min_bytes {
398                    return Err(FabricError::CapacityExceeded);
399                }
400                let created_at_unix_secs = info
401                    .expires_at_unix_secs
402                    .saturating_sub(self.lease_secs.max(1));
403                let state = lease::new_shared_lease_from_parts(
404                    info.lease_id,
405                    created_at_unix_secs,
406                    info.expires_at_unix_secs,
407                    lease_status_from_host(info.status),
408                );
409                Ok(MemLease {
410                    state: state.clone(),
411                    mem: FabricMem {
412                        lease_state: Some(state),
413                        host_lease_id: Some(info.lease_id),
414                    },
415                    _size: info.arena_size,
416                })
417            }
418            Err(FabricError::Unsupported) => {
419                let state = lease::new_shared_lease(LEASE_TAG_MEM, self.lease_secs);
420                let mut mem = FabricMem::hello()?;
421                mem.lease_state = Some(state.clone());
422                let arena = mem.arena_size()?;
423                if arena < self.min_bytes {
424                    return Err(FabricError::CapacityExceeded);
425                }
426                Ok(MemLease {
427                    state,
428                    mem,
429                    _size: arena,
430                })
431            }
432            Err(e) => Err(e),
433        };
434        #[cfg(feature = "observe")]
435        if let Ok(ref lease) = result {
436            crate::observe_hooks::on_lease_acquired(
437                LEASE_TAG_MEM,
438                lease.lease_id(),
439                "local",
440                lease._size,
441            );
442        }
443        result
444    }
445
446    /// Attach to an existing memory lease by ID.
447    ///
448    /// Instead of allocating a new lease, this queries the host for the
449    /// given `lease_id` and returns a [`MemLease`] bound to it if the
450    /// lease is still active. This enables crash recovery: a replacement
451    /// tasklet can reconnect to storage that survived the previous
452    /// tasklet's death.
453    ///
454    /// # Errors
455    ///
456    /// - [`FabricError::LeaseExpired`] if the lease has expired.
457    /// - [`FabricError::Revoked`] if the lease has been revoked.
458    /// - [`FabricError::Disconnected`] if the query fails.
459    pub fn attach(lease_id: u128) -> Result<MemLease> {
460        let info = host::fbmu_query(lease_id)?;
461        let status = lease_status_from_host(info.status);
462        match status {
463            LeaseStatus::Active => {}
464            LeaseStatus::Expired => return Err(FabricError::LeaseExpired),
465            LeaseStatus::Revoked => return Err(FabricError::Revoked),
466        }
467        // We don't know the original lease duration, so estimate created_at
468        // from the current expiry. This is imprecise but sufficient for
469        // status tracking — the authoritative state comes from the host.
470        let created_at_unix_secs = 0;
471        let state = lease::new_shared_lease_from_parts(
472            info.lease_id,
473            created_at_unix_secs,
474            info.expires_at_unix_secs,
475            status,
476        );
477        let result = MemLease {
478            state: state.clone(),
479            mem: FabricMem {
480                lease_state: Some(state),
481                host_lease_id: Some(info.lease_id),
482            },
483            _size: info.arena_size,
484        };
485        #[cfg(feature = "observe")]
486        crate::observe_hooks::on_lease_acquired(
487            LEASE_TAG_MEM,
488            result.lease_id(),
489            "attach",
490            result._size,
491        );
492        Ok(result)
493    }
494}
495
496impl Default for MemBuilder {
497    fn default() -> Self {
498        Self::new()
499    }
500}
501
502#[cfg(test)]
503mod tests {
504    use super::*;
505    use crate::host;
506
507    #[test]
508    fn mem_hello_write_read_roundtrip() {
509        host::reset_mock();
510        host::mock_set_fbmu_arena_size(4096);
511
512        let mem = FabricMem::hello().expect("hello");
513        assert_eq!(mem.arena_size().unwrap(), 4096);
514
515        let data = b"fabricBIOS-grafos-std";
516        mem.write(0, data).expect("write");
517
518        let readback = mem.read(0, data.len() as u32).expect("read");
519        assert_eq!(&readback, data);
520    }
521
522    #[test]
523    fn mem_hello_error_propagates() {
524        host::reset_mock();
525        host::mock_set_fbmu_hello_error(Some(-1));
526
527        let result = FabricMem::hello();
528        assert!(result.is_err());
529        assert_eq!(result.unwrap_err(), FabricError::Disconnected);
530
531        host::mock_set_fbmu_hello_error(None);
532    }
533
534    #[test]
535    fn mem_builder_checks_capacity() {
536        host::reset_mock();
537        host::mock_set_fbmu_arena_size(1024);
538
539        // Request more than available
540        let result = MemBuilder::new().min_bytes(2048).acquire();
541        assert_eq!(result.unwrap_err(), FabricError::CapacityExceeded);
542
543        // Request exactly available
544        let lease = MemBuilder::new()
545            .min_bytes(1024)
546            .acquire()
547            .expect("acquire");
548        assert_eq!(lease.mem().arena_size().unwrap(), 1024);
549    }
550
551    #[test]
552    fn mem_builder_default_acquires() {
553        host::reset_mock();
554        host::mock_set_fbmu_arena_size(65536);
555
556        let lease = MemBuilder::new().acquire().expect("acquire");
557        let data = b"test";
558        lease.mem().write(100, data).expect("write");
559        let readback = lease.mem().read(100, 4).expect("read");
560        assert_eq!(&readback, data);
561    }
562
563    #[test]
564    fn mem_lease_expires_and_can_be_renewed() {
565        host::reset_mock();
566        host::mock_set_fbmu_arena_size(4096);
567        host::mock_set_unix_time_secs(1_000);
568
569        let lease = MemBuilder::new().lease_secs(10).acquire().expect("acquire");
570        assert_eq!(lease.status(), LeaseStatus::Active);
571        assert_eq!(lease.expires_at_unix_secs(), 1_010);
572
573        lease.renew(20).expect("renew");
574        assert_eq!(lease.expires_at_unix_secs(), 1_020);
575
576        host::mock_advance_time_secs(21);
577        assert_eq!(lease.status(), LeaseStatus::Expired);
578        assert_eq!(
579            lease.mem().write(0, b"x").unwrap_err(),
580            FabricError::LeaseExpired
581        );
582    }
583
584    #[test]
585    fn mem_lease_free_revokes_operations() {
586        host::reset_mock();
587        host::mock_set_fbmu_arena_size(4096);
588
589        let lease = MemBuilder::new().acquire().expect("acquire");
590        lease.free();
591        assert_eq!(lease.status(), LeaseStatus::Revoked);
592        assert_eq!(lease.mem().read(0, 1).unwrap_err(), FabricError::Revoked);
593    }
594
595    #[test]
596    fn mem_attach_reconnects_to_active_lease() {
597        host::reset_mock();
598        host::mock_set_fbmu_arena_size(4096);
599
600        let lease = MemBuilder::new()
601            .lease_secs(300)
602            .acquire()
603            .expect("acquire");
604        let id = lease.lease_id();
605
606        // Write data through the original lease.
607        lease.mem().write(0, b"hello").expect("write");
608
609        // Attach to the same lease by ID.
610        let attached = MemBuilder::attach(id).expect("attach");
611        assert_eq!(attached.lease_id(), id);
612        assert_eq!(attached.status(), LeaseStatus::Active);
613
614        // Read back data through the attached lease.
615        let readback = attached.mem().read(0, 5).expect("read");
616        assert_eq!(&readback, b"hello");
617    }
618
619    #[test]
620    fn mem_attach_expired_returns_error() {
621        host::reset_mock();
622        host::mock_set_fbmu_arena_size(4096);
623        host::mock_set_unix_time_secs(1_000);
624
625        let lease = MemBuilder::new().lease_secs(5).acquire().expect("acquire");
626        let id = lease.lease_id();
627
628        // Advance past expiry.
629        host::mock_advance_time_secs(10);
630
631        let result = MemBuilder::attach(id);
632        assert_eq!(result.unwrap_err(), FabricError::LeaseExpired);
633    }
634
635    #[test]
636    fn mem_attach_revoked_returns_error() {
637        host::reset_mock();
638        host::mock_set_fbmu_arena_size(4096);
639
640        let lease = MemBuilder::new().acquire().expect("acquire");
641        let id = lease.lease_id();
642        lease.free(); // revoke
643
644        let result = MemBuilder::attach(id);
645        assert_eq!(result.unwrap_err(), FabricError::Revoked);
646    }
647
648    #[test]
649    fn mem_attach_unknown_lease_returns_error() {
650        host::reset_mock();
651        host::mock_set_fbmu_arena_size(4096);
652
653        // Attach to a lease ID that was never allocated.
654        let result = MemBuilder::attach(0xDEADBEEF);
655        assert!(result.is_err());
656    }
657
658    #[test]
659    fn mem_leases_are_isolated_by_lease_id() {
660        host::reset_mock();
661        host::mock_set_fbmu_arena_size(4096);
662
663        let lease_a = MemBuilder::new().acquire().expect("acquire a");
664        let lease_b = MemBuilder::new().acquire().expect("acquire b");
665
666        lease_a.mem().write(0, b"aaa").expect("write a");
667        lease_b.mem().write(0, b"bbb").expect("write b");
668
669        let read_a = lease_a.mem().read(0, 3).expect("read a");
670        let read_b = lease_b.mem().read(0, 3).expect("read b");
671        assert_eq!(&read_a, b"aaa");
672        assert_eq!(&read_b, b"bbb");
673    }
674}
grafos_std/mem.rs

grafos_std/
mem.rs
