grafos_rpc/

lib.rs

//! grafos-rpc -- Lease-backed RPC framework for grafOS.
//!
//! The hot path is a leased shared memory write + read instead of a network
//! round trip. Client writes a request to a leased memory region; server
//! reads from the same region and writes back a response. The underlying
//! transport is FBMU (Fabric Bootstrap Memory Unit), so co-located services
//! exchange calls without touching the network stack.
//!
//! # Quick start
//!
//! ```rust
//! use grafos_rpc::{RpcHandler, RpcClient, RpcServer};
//! use grafos_std::error::{FabricError, Result};
//!
//! // 1. Define your service handler.
//! struct Adder;
//! impl RpcHandler for Adder {
//!     fn handle(&self, method_id: u32, payload: &[u8]) -> Result<Vec<u8>> {
//!         let args: (f64, f64) = postcard::from_bytes(payload)
//!             .map_err(|_| FabricError::IoError(-200))?;
//!         postcard::to_allocvec(&(args.0 + args.1))
//!             .map_err(|_| FabricError::IoError(-201))
//!     }
//! }
//!
//! // 2. Acquire a lease (mock backend for native testing).
//! grafos_std::host::reset_mock();
//! grafos_std::host::mock_set_fbmu_arena_size(65536);
//! let lease = grafos_std::mem::MemBuilder::new().min_bytes(65536).acquire()?;
//!
//! // 3. Create client and server on the same lease.
//! let mut client = RpcClient::new(&lease);
//! let server = RpcServer::new(&lease);
//! # Ok::<(), FabricError>(())
//! ```
//!
//! # Architecture
//!
//! The shared memory arena is divided into two regions at fixed offsets:
//!
//! - **Request region** (offset 0): Written by the client, read by the server.
//! - **Response region** (offset 32768): Written by the server, read by the
//!   client.
//!
//! Each region has a status byte that drives the protocol state machine:
//!
//! | Value | Constant | Meaning |
//! |-------|----------|---------|
//! | 0 | [`EMPTY`] | Slot is idle |
//! | 1 | [`REQUEST_READY`] | Client has written a request |
//! | 2 | [`PROCESSING`] | Server is handling the request |
//! | 3 | [`RESPONSE_READY`] | Server has written the response |
//! | 4 | [`ERROR`] | Server encountered an error |
//!
//! # Wire layout
//!
//! Request region (written at offset 0):
//! ```text
//! [request_id: u64 LE (8B)] [method_id: u32 LE (4B)] [status: u8 (1B)]
//! [payload_len: u32 LE (4B)] [payload: [u8; payload_len]]
//! ```
//!
//! Response region (written at offset 32768):
//! ```text
//! [request_id: u64 LE (8B)] [status: u8 (1B)]
//! [payload_len: u32 LE (4B)] [payload: [u8; payload_len]]
//! ```
//!
//! All integers are little-endian. Payloads are serialized with
//! [postcard](https://docs.rs/postcard) (compact, `no_std`-friendly).
//! Maximum payload size is 30 KiB.

#![cfg_attr(not(feature = "std"), no_std)]

extern crate alloc;
#[cfg(test)]
extern crate self as grafos_rpc;

use alloc::vec::Vec;

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

pub mod mux;
pub mod transport;
pub use grafos_rpc_macros::grafos_rpc_service;
pub use mux::{RpcMuxClient, RpcMuxServer};
pub use transport::{
    AutoTransport, QuicTransport, RpcTransport, ServiceHandlerAdapter, SharedMemoryTransport,
};

// ---------------------------------------------------------------------------
// Status constants
// ---------------------------------------------------------------------------

/// Slot is idle — no pending request or response.
pub const EMPTY: u8 = 0;
/// Client has written a request; server should process it.
pub const REQUEST_READY: u8 = 1;
/// Server is currently processing the request.
pub const PROCESSING: u8 = 2;
/// Server has written the response; client can read it.
pub const RESPONSE_READY: u8 = 3;
/// Server encountered an error processing the request.
pub const ERROR: u8 = 4;

// ---------------------------------------------------------------------------
// Layout constants
// ---------------------------------------------------------------------------

/// Byte offset within the arena where the request region starts.
const REQUEST_OFFSET: u64 = 0;

/// Byte offset within the arena where the response region starts.
/// Placed at 32 KiB boundary to separate from request region.
const RESPONSE_OFFSET: u64 = 32768;

/// Header size for request: request_id(8) + method_id(4) + status(1) + payload_len(4) = 17
const REQUEST_HEADER_SIZE: usize = 17;

/// Header size for response: request_id(8) + status(1) + payload_len(4) = 13
const RESPONSE_HEADER_SIZE: usize = 13;

/// Maximum payload size (limited by arena region size).
const MAX_PAYLOAD_SIZE: usize = 30 * 1024; // 30 KiB

/// Default poll iteration limit (used as timeout proxy in tests).
const DEFAULT_MAX_POLL_ITERATIONS: u64 = 1_000_000;

// ---------------------------------------------------------------------------
// Request/response encoding
// ---------------------------------------------------------------------------

fn encode_request(request_id: u64, method_id: u32, status: u8, payload: &[u8]) -> Vec<u8> {
    let mut buf = Vec::with_capacity(REQUEST_HEADER_SIZE + payload.len());
    buf.extend_from_slice(&request_id.to_le_bytes());
    buf.extend_from_slice(&method_id.to_le_bytes());
    buf.push(status);
    buf.extend_from_slice(&(payload.len() as u32).to_le_bytes());
    buf.extend_from_slice(payload);
    buf
}

fn decode_request(data: &[u8]) -> Option<(u64, u32, u8, &[u8])> {
    if data.len() < REQUEST_HEADER_SIZE {
        return None;
    }
    let request_id = u64::from_le_bytes(data[0..8].try_into().ok()?);
    let method_id = u32::from_le_bytes(data[8..12].try_into().ok()?);
    let status = data[12];
    let payload_len = u32::from_le_bytes(data[13..17].try_into().ok()?) as usize;
    if data.len() < REQUEST_HEADER_SIZE + payload_len {
        return None;
    }
    Some((request_id, method_id, status, &data[17..17 + payload_len]))
}

fn encode_response(request_id: u64, status: u8, payload: &[u8]) -> Vec<u8> {
    let mut buf = Vec::with_capacity(RESPONSE_HEADER_SIZE + payload.len());
    buf.extend_from_slice(&request_id.to_le_bytes());
    buf.push(status);
    buf.extend_from_slice(&(payload.len() as u32).to_le_bytes());
    buf.extend_from_slice(payload);
    buf
}

fn decode_response(data: &[u8]) -> Option<(u64, u8, &[u8])> {
    if data.len() < RESPONSE_HEADER_SIZE {
        return None;
    }
    let request_id = u64::from_le_bytes(data[0..8].try_into().ok()?);
    let status = data[8];
    let payload_len = u32::from_le_bytes(data[9..13].try_into().ok()?) as usize;
    if data.len() < RESPONSE_HEADER_SIZE + payload_len {
        return None;
    }
    Some((request_id, status, &data[13..13 + payload_len]))
}

// ---------------------------------------------------------------------------
// RpcHandler trait
// ---------------------------------------------------------------------------

/// Trait for handling RPC method dispatches on the server side.
///
/// Implement this for your service type to map `method_id` values to
/// handler logic. The `payload` is the postcard-serialized request arguments;
/// the return value should be the postcard-serialized response.
///
/// # Example
///
/// ```rust
/// use grafos_rpc::RpcHandler;
/// use grafos_std::error::{FabricError, Result};
/// use serde::{Serialize, Deserialize};
///
/// const METHOD_GREET: u32 = 0;
///
/// #[derive(Serialize, Deserialize)]
/// struct GreetRequest { name: String }
///
/// struct Greeter;
///
/// impl RpcHandler for Greeter {
///     fn handle(&self, method_id: u32, payload: &[u8]) -> Result<Vec<u8>> {
///         match method_id {
///             METHOD_GREET => {
///                 let req: GreetRequest = postcard::from_bytes(payload)
///                     .map_err(|_| FabricError::IoError(-200))?;
///                 let greeting = format!("hello, {}!", req.name);
///                 postcard::to_allocvec(&greeting)
///                     .map_err(|_| FabricError::IoError(-201))
///             }
///             _ => Err(FabricError::Unsupported),
///         }
///     }
/// }
/// ```
pub trait RpcHandler {
    /// Dispatch a single RPC call.
    ///
    /// - `method_id`: Identifies which method is being called (by convention,
    ///   `u32` constants starting at 0).
    /// - `payload`: Postcard-serialized request arguments.
    ///
    /// Returns postcard-serialized response bytes on success. Return
    /// `Err(FabricError::Unsupported)` for unknown method IDs.
    fn handle(&self, method_id: u32, payload: &[u8]) -> Result<Vec<u8>>;
}

// ---------------------------------------------------------------------------
// RpcClient
// ---------------------------------------------------------------------------

/// Client side of the shared-memory RPC protocol.
///
/// Writes serialized requests into the lease's request region (offset 0),
/// then polls the response region (offset 32768) until the server has
/// written back a result.
///
/// Each call assigns a monotonically increasing `request_id` so that
/// stale responses from a previous call are ignored.
///
/// # Example
///
/// ```rust
/// use grafos_rpc::RpcClient;
/// use grafos_std::mem::MemBuilder;
///
/// # grafos_std::host::reset_mock();
/// # grafos_std::host::mock_set_fbmu_arena_size(65536);
/// let lease = MemBuilder::new().min_bytes(65536).acquire()?;
/// let mut client = RpcClient::new(&lease)
///     .with_max_poll_iterations(100); // short timeout for tests
/// # Ok::<(), grafos_std::error::FabricError>(())
/// ```
pub struct RpcClient<'a> {
    lease: &'a MemLease,
    next_request_id: u64,
    max_poll_iterations: u64,
}

impl<'a> RpcClient<'a> {
    /// Create a new RPC client backed by the given memory lease.
    ///
    /// Initializes with `request_id = 1` and a default poll limit of
    /// 1,000,000 iterations. Use [`with_max_poll_iterations`](Self::with_max_poll_iterations)
    /// to adjust the timeout behavior.
    pub fn new(lease: &'a MemLease) -> Self {
        RpcClient {
            lease,
            next_request_id: 1,
            max_poll_iterations: DEFAULT_MAX_POLL_ITERATIONS,
        }
    }

    /// Set the maximum number of poll iterations before [`call`](Self::call)
    /// returns `Err(FabricError::LeaseExpired)`.
    ///
    /// A low value (e.g. 10) is useful in tests to fail fast when no server
    /// is running. The default is 1,000,000.
    pub fn with_max_poll_iterations(mut self, n: u64) -> Self {
        self.max_poll_iterations = n;
        self
    }

    /// Perform an RPC call.
    ///
    /// Serializes `req` with postcard, writes it to the request region with
    /// `method_id`, then polls the response region until the server responds.
    ///
    /// # Errors
    ///
    /// - `FabricError::IoError(-100)` -- serialization of `req` failed.
    /// - [`FabricError::CapacityExceeded`] -- serialized payload exceeds 30 KiB.
    /// - [`FabricError::LeaseExpired`] -- poll limit exceeded (timeout).
    /// - `FabricError::IoError(-101)` -- deserialization of response failed.
    /// - `FabricError::IoError(-102)` -- server returned [`ERROR`] status.
    /// - Other [`FabricError`] variants from the underlying memory operations.
    pub fn call<Req, Resp>(&mut self, method_id: u32, req: &Req) -> Result<Resp>
    where
        Req: serde::Serialize,
        Resp: serde::de::DeserializeOwned,
    {
        match self.lease.status() {
            LeaseStatus::Active => {}
            LeaseStatus::Revoked => return Err(FabricError::Revoked),
            _ => return Err(FabricError::LeaseExpired),
        }
        let request_id = self.next_request_id;
        self.next_request_id += 1;

        // Serialize the request payload.
        let payload = postcard::to_allocvec(req).map_err(|_| FabricError::IoError(-100))?;
        if payload.len() > MAX_PAYLOAD_SIZE {
            return Err(FabricError::CapacityExceeded);
        }

        // Write request to shared memory.
        let req_buf = encode_request(request_id, method_id, REQUEST_READY, &payload);
        self.lease.mem().write(REQUEST_OFFSET, &req_buf)?;

        // Poll response region until server writes back.
        for _ in 0..self.max_poll_iterations {
            let resp_data = self.lease.mem().read(
                RESPONSE_OFFSET,
                (RESPONSE_HEADER_SIZE + MAX_PAYLOAD_SIZE) as u32,
            )?;

            if let Some((resp_id, status, resp_payload)) = decode_response(&resp_data) {
                if resp_id == request_id {
                    match status {
                        RESPONSE_READY => {
                            // Clear the response region.
                            let clear = encode_response(0, EMPTY, &[]);
                            self.lease.mem().write(RESPONSE_OFFSET, &clear)?;

                            let result: Resp = postcard::from_bytes(resp_payload)
                                .map_err(|_| FabricError::IoError(-101))?;
                            return Ok(result);
                        }
                        ERROR => {
                            // Read the error message from payload, then clear.
                            let clear = encode_response(0, EMPTY, &[]);
                            self.lease.mem().write(RESPONSE_OFFSET, &clear)?;

                            return Err(FabricError::IoError(-102));
                        }
                        _ => {
                            // Still processing or not ready yet — continue polling.
                        }
                    }
                }
            }
        }

        Err(FabricError::LeaseExpired)
    }
}

// ---------------------------------------------------------------------------
// RpcServer
// ---------------------------------------------------------------------------

/// Server side of the shared-memory RPC protocol.
///
/// Polls the request region (offset 0) for incoming calls, dispatches them
/// via the provided [`RpcHandler`], and writes responses to the response
/// region (offset 32768).
///
/// # Example
///
/// ```rust
/// use grafos_rpc::{RpcHandler, RpcServer};
/// use grafos_std::error::{FabricError, Result};
/// use grafos_std::mem::MemBuilder;
///
/// struct Echo;
/// impl RpcHandler for Echo {
///     fn handle(&self, _method_id: u32, payload: &[u8]) -> Result<Vec<u8>> {
///         Ok(payload.to_vec())
///     }
/// }
///
/// # grafos_std::host::reset_mock();
/// # grafos_std::host::mock_set_fbmu_arena_size(65536);
/// let lease = MemBuilder::new().min_bytes(65536).acquire()?;
/// let server = RpcServer::new(&lease);
///
/// // In your event loop:
/// // let handled = server.poll_once(&Echo)?;
/// # Ok::<(), FabricError>(())
/// ```
pub struct RpcServer<'a> {
    lease: &'a MemLease,
}

impl<'a> RpcServer<'a> {
    /// Create a new RPC server backed by the given memory lease.
    pub fn new(lease: &'a MemLease) -> Self {
        RpcServer { lease }
    }

    /// Poll for and handle a single request.
    ///
    /// Reads the request region. If a request with status [`REQUEST_READY`]
    /// is present:
    ///
    /// 1. Sets status to [`PROCESSING`].
    /// 2. Calls `handler.handle(method_id, payload)`.
    /// 3. Writes [`RESPONSE_READY`] (or [`ERROR`]) to the response region.
    /// 4. Clears the request region to [`EMPTY`].
    /// 5. Returns `Ok(true)`.
    ///
    /// If no request is pending, returns `Ok(false)` without side effects.
    pub fn poll_once<H: RpcHandler>(&self, handler: &H) -> Result<bool> {
        match self.lease.status() {
            LeaseStatus::Active => {}
            LeaseStatus::Revoked => return Err(FabricError::Revoked),
            _ => return Err(FabricError::LeaseExpired),
        }

        let req_data = self.lease.mem().read(
            REQUEST_OFFSET,
            (REQUEST_HEADER_SIZE + MAX_PAYLOAD_SIZE) as u32,
        )?;

        let (request_id, method_id, status, payload) = match decode_request(&req_data) {
            Some(decoded) => decoded,
            None => return Ok(false),
        };

        if status != REQUEST_READY {
            return Ok(false);
        }

        // Mark as PROCESSING.
        let processing_buf = encode_request(request_id, method_id, PROCESSING, payload);
        self.lease.mem().write(REQUEST_OFFSET, &processing_buf)?;

        // Dispatch to handler.
        match handler.handle(method_id, payload) {
            Ok(resp_payload) => {
                let resp_buf = encode_response(request_id, RESPONSE_READY, &resp_payload);
                self.lease.mem().write(RESPONSE_OFFSET, &resp_buf)?;
            }
            Err(_) => {
                let resp_buf = encode_response(request_id, ERROR, &[]);
                self.lease.mem().write(RESPONSE_OFFSET, &resp_buf)?;
            }
        }

        // Clear request region.
        let clear = encode_request(0, 0, EMPTY, &[]);
        self.lease.mem().write(REQUEST_OFFSET, &clear)?;

        Ok(true)
    }

    /// Serve requests in a loop for up to `max_iterations` polls.
    ///
    /// Returns the number of requests actually handled (i.e. the number of
    /// times [`poll_once`](Self::poll_once) returned `Ok(true)`).
    ///
    /// This is a convenience for testing where you know a fixed number of
    /// calls will arrive. In production, call [`poll_once`](Self::poll_once)
    /// from your own event loop.
    pub fn serve_n<H: RpcHandler>(&self, handler: &H, max_iterations: u64) -> Result<u64> {
        let mut handled = 0u64;
        for _ in 0..max_iterations {
            if self.poll_once(handler)? {
                handled += 1;
            }
        }
        Ok(handled)
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use grafos_std::host;
    use grafos_std::mem::MemBuilder;
    use serde::{Deserialize, Serialize};

    // -- Calculator service (manual, no proc macro) --

    /// Request arguments for Calculator methods.
    #[derive(Serialize, Deserialize)]
    struct CalcArgs {
        a: f64,
        b: f64,
    }

    const METHOD_ADD: u32 = 0;
    const METHOD_MULTIPLY: u32 = 1;

    /// Calculator service implementation.
    struct Calculator;

    impl RpcHandler for Calculator {
        fn handle(&self, method_id: u32, payload: &[u8]) -> Result<Vec<u8>> {
            match method_id {
                METHOD_ADD => {
                    let args: CalcArgs =
                        postcard::from_bytes(payload).map_err(|_| FabricError::IoError(-200))?;
                    let result = args.a + args.b;
                    postcard::to_allocvec(&result).map_err(|_| FabricError::IoError(-201))
                }
                METHOD_MULTIPLY => {
                    let args: CalcArgs =
                        postcard::from_bytes(payload).map_err(|_| FabricError::IoError(-200))?;
                    let result = args.a * args.b;
                    postcard::to_allocvec(&result).map_err(|_| FabricError::IoError(-201))
                }
                _ => Err(FabricError::Unsupported),
            }
        }
    }

    /// Helper: run a single call through client → server → client in one thread.
    ///
    /// This simulates the protocol by interleaving client write, server poll,
    /// and client read, since both share the same mock memory.
    fn call_roundtrip<Req, Resp>(
        lease: &MemLease,
        handler: &impl RpcHandler,
        method_id: u32,
        req: &Req,
    ) -> Result<Resp>
    where
        Req: serde::Serialize,
        Resp: serde::de::DeserializeOwned,
    {
        // Step 1: Client writes request.
        // Use a fixed request_id (1) since call_roundtrip is stateless.
        let request_id: u64 = 1;

        let payload = postcard::to_allocvec(req).map_err(|_| FabricError::IoError(-100))?;
        let req_buf = encode_request(request_id, method_id, REQUEST_READY, &payload);
        lease.mem().write(REQUEST_OFFSET, &req_buf)?;

        // Step 2: Server processes.
        let server = RpcServer::new(lease);
        let handled = server.poll_once(handler)?;
        assert!(handled);

        // Step 3: Client reads response.
        let resp_data = lease.mem().read(
            RESPONSE_OFFSET,
            (RESPONSE_HEADER_SIZE + MAX_PAYLOAD_SIZE) as u32,
        )?;
        let (resp_id, status, resp_payload) = decode_response(&resp_data).unwrap();
        assert_eq!(resp_id, request_id);
        assert_eq!(status, RESPONSE_READY);

        let result: Resp =
            postcard::from_bytes(resp_payload).map_err(|_| FabricError::IoError(-101))?;

        // Clear response.
        let clear = encode_response(0, EMPTY, &[]);
        lease.mem().write(RESPONSE_OFFSET, &clear)?;

        Ok(result)
    }

    #[test]
    fn calculator_add_roundtrip() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);

        let lease = MemBuilder::new().min_bytes(65536).acquire().unwrap();
        let calc = Calculator;

        let result: f64 =
            call_roundtrip(&lease, &calc, METHOD_ADD, &CalcArgs { a: 1.5, b: 2.5 }).unwrap();
        assert!((result - 4.0).abs() < f64::EPSILON);
    }

    #[test]
    fn calculator_multiply_roundtrip() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);

        let lease = MemBuilder::new().min_bytes(65536).acquire().unwrap();
        let calc = Calculator;

        let result: f64 =
            call_roundtrip(&lease, &calc, METHOD_MULTIPLY, &CalcArgs { a: 3.0, b: 7.0 }).unwrap();
        assert!((result - 21.0).abs() < f64::EPSILON);
    }

    #[test]
    fn unknown_method_returns_error() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);

        let lease = MemBuilder::new().min_bytes(65536).acquire().unwrap();
        let calc = Calculator;

        // Write request with unknown method.
        let payload = postcard::to_allocvec(&CalcArgs { a: 1.0, b: 1.0 }).unwrap();
        let req_buf = encode_request(1, 99, REQUEST_READY, &payload);
        lease.mem().write(REQUEST_OFFSET, &req_buf).unwrap();

        let server = RpcServer::new(&lease);
        server.poll_once(&calc).unwrap();

        // Read response — should be ERROR.
        let resp_data = lease
            .mem()
            .read(
                RESPONSE_OFFSET,
                (RESPONSE_HEADER_SIZE + MAX_PAYLOAD_SIZE) as u32,
            )
            .unwrap();
        let (resp_id, status, _) = decode_response(&resp_data).unwrap();
        assert_eq!(resp_id, 1);
        assert_eq!(status, ERROR);
    }

    #[test]
    fn timeout_when_no_server() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);

        let lease = MemBuilder::new().min_bytes(65536).acquire().unwrap();

        // Client sends request but no server ever processes it.
        let mut client = RpcClient::new(&lease).with_max_poll_iterations(10);
        let result: Result<f64> = client.call(METHOD_ADD, &CalcArgs { a: 1.0, b: 2.0 });

        assert_eq!(result.unwrap_err(), FabricError::LeaseExpired);
    }

    #[test]
    fn server_poll_returns_false_when_empty() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);

        let lease = MemBuilder::new().min_bytes(65536).acquire().unwrap();
        let calc = Calculator;

        // Clear memory — write EMPTY status to request region.
        let clear = encode_request(0, 0, EMPTY, &[]);
        lease.mem().write(REQUEST_OFFSET, &clear).unwrap();

        let server = RpcServer::new(&lease);
        assert!(!server.poll_once(&calc).unwrap());
    }

    #[test]
    fn multiple_sequential_calls() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);

        let lease = MemBuilder::new().min_bytes(65536).acquire().unwrap();
        let calc = Calculator;

        // Call 1: add
        let r1: f64 =
            call_roundtrip(&lease, &calc, METHOD_ADD, &CalcArgs { a: 10.0, b: 20.0 }).unwrap();
        assert!((r1 - 30.0).abs() < f64::EPSILON);

        // Call 2: multiply
        let r2: f64 =
            call_roundtrip(&lease, &calc, METHOD_MULTIPLY, &CalcArgs { a: 5.0, b: 6.0 }).unwrap();
        assert!((r2 - 30.0).abs() < f64::EPSILON);

        // Call 3: add again
        let r3: f64 =
            call_roundtrip(&lease, &calc, METHOD_ADD, &CalcArgs { a: -1.0, b: 1.0 }).unwrap();
        assert!(r3.abs() < f64::EPSILON);
    }

    #[test]
    fn encode_decode_request_roundtrip() {
        let payload = b"hello";
        let encoded = encode_request(42, 7, REQUEST_READY, payload);
        let (rid, mid, status, p) = decode_request(&encoded).unwrap();
        assert_eq!(rid, 42);
        assert_eq!(mid, 7);
        assert_eq!(status, REQUEST_READY);
        assert_eq!(p, payload);
    }

    #[test]
    fn encode_decode_response_roundtrip() {
        let payload = b"world";
        let encoded = encode_response(99, RESPONSE_READY, payload);
        let (rid, status, p) = decode_response(&encoded).unwrap();
        assert_eq!(rid, 99);
        assert_eq!(status, RESPONSE_READY);
        assert_eq!(p, payload);
    }

    #[test]
    fn decode_request_rejects_truncated() {
        assert!(decode_request(&[0; 10]).is_none());
        assert!(decode_request(&[]).is_none());
    }

    #[test]
    fn decode_response_rejects_truncated() {
        assert!(decode_response(&[0; 5]).is_none());
        assert!(decode_response(&[]).is_none());
    }

    // -- Proc macro generated service test --

    #[grafos_rpc_service]
    pub trait Greeter {
        fn greet(&self, name: String) -> String;
        fn add(&self, a: f64, b: f64) -> f64;
    }

    struct GreeterImpl;

    impl GreeterServer for GreeterImpl {
        fn greet(&self, name: String) -> String {
            alloc::format!("hello, {}!", name)
        }

        fn add(&self, a: f64, b: f64) -> f64 {
            a + b
        }
    }

    #[test]
    fn proc_macro_generates_server_dispatch() {
        let svc = GreeterImpl;

        // method_id 0 = greet
        let payload = postcard::to_allocvec(&("world".to_string(),)).unwrap();
        let resp = svc.dispatch(0, &payload).unwrap();
        let greeting: String = postcard::from_bytes(&resp).unwrap();
        assert_eq!(greeting, "hello, world!");

        // method_id 1 = add
        let payload = postcard::to_allocvec(&(1.5f64, 2.5f64)).unwrap();
        let resp = svc.dispatch(1, &payload).unwrap();
        let sum: f64 = postcard::from_bytes(&resp).unwrap();
        assert!((sum - 4.0).abs() < f64::EPSILON);
    }

    #[test]
    fn proc_macro_unknown_method_returns_unsupported() {
        let svc = GreeterImpl;
        let result = svc.dispatch(99, &[]);
        assert_eq!(result.unwrap_err(), FabricError::Unsupported);
    }

    #[test]
    fn rpc_client_call_on_revoked_lease() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);

        let lease = MemBuilder::new().min_bytes(65536).acquire().unwrap();
        let mut client = RpcClient::new(&lease).with_max_poll_iterations(10);

        // Free the lease to revoke it.
        lease.free();

        let result: Result<f64> = client.call(METHOD_ADD, &CalcArgs { a: 1.0, b: 2.0 });
        assert_eq!(result.unwrap_err(), FabricError::Revoked);
    }

    #[test]
    fn rpc_server_poll_on_revoked_lease() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);

        let lease = MemBuilder::new().min_bytes(65536).acquire().unwrap();
        let server = RpcServer::new(&lease);
        let calc = Calculator;

        // Free the lease to revoke it.
        lease.free();

        let result = server.poll_once(&calc);
        assert_eq!(result.unwrap_err(), FabricError::Revoked);
    }

    #[test]
    fn proc_macro_client_roundtrip_via_quic_transport() {
        // Use a QuicTransport that dispatches to the GreeterImpl server.
        let svc = GreeterImpl;
        let transport = crate::QuicTransport::new(alloc::boxed::Box::new(
            move |method_id, _req_id, payload| svc.dispatch(method_id, payload),
        ));

        let client = GreeterClient::new(alloc::boxed::Box::new(transport));

        let greeting = client.greet("fabric".to_string()).unwrap();
        assert_eq!(greeting, "hello, fabric!");

        let sum = client.add(10.0, 20.0).unwrap();
        assert!((sum - 30.0).abs() < f64::EPSILON);
    }
}