agent-mesh-bus 0.6.3

//! Request-reply correlation.
//!
//! Each outbound [`crate::inbox::BusMessage::Request`] carries a
//! [`CorrelationId`] generated by the requester. The matching
//! [`crate::inbox::BusMessage::Reply`] echoes that id back so the bus
//! can route the response to the right waiter. [`ReplyWaiter`] holds
//! the in-flight oneshot senders, keyed by correlation.

use rand::RngCore;
use std::collections::HashMap;
use std::sync::Mutex;
use tokio::sync::oneshot;

/// A 16-byte correlation identifier.
///
/// 128 bits is enough collision space that the requester can draw a
/// fresh one at random without coordination, even at very high request
/// rates. The bytes are echoed verbatim in the reply.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CorrelationId(pub [u8; 16]);

impl CorrelationId {
    /// Draw a fresh random correlation id from the OS RNG.
    #[must_use]
    pub fn new_random() -> Self {
        let mut id = [0u8; 16];
        rand::thread_rng().fill_bytes(&mut id);
        Self(id)
    }

    /// Full 32-character hex representation. For log lines and the
    /// `CorrelationId` round-trip in `BusMessage`.
    #[must_use]
    pub fn hex(&self) -> String {
        hex::encode(self.0)
    }
}

/// Holds oneshot senders for in-flight requests, keyed by correlation.
///
/// `register` returns the receiver half of a oneshot; the bus's inbox
/// later calls `deliver` when the matching reply arrives. If the
/// request times out or the bus shuts down before the reply, the
/// waiter is dropped (or explicitly cancelled) and any late reply is
/// silently discarded.
#[derive(Default)]
pub struct ReplyWaiter {
    waiters: Mutex<HashMap<CorrelationId, oneshot::Sender<Vec<u8>>>>,
}

impl ReplyWaiter {
    /// Build a fresh waiter with no in-flight requests.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a fresh correlation id and return the receiver the
    /// caller will await for the eventual reply.
    pub fn register(&self, id: CorrelationId) -> oneshot::Receiver<Vec<u8>> {
        let (tx, rx) = oneshot::channel();
        self.waiters
            .lock()
            .expect("reply waiter poisoned")
            .insert(id, tx);
        rx
    }

    /// Deliver a reply payload to whoever is waiting on `id`.
    ///
    /// Returns `true` on success, `false` if there is no waiter for
    /// this id (already cancelled, timed out, or never registered).
    pub fn deliver(&self, id: CorrelationId, payload: Vec<u8>) -> bool {
        if let Some(tx) = self
            .waiters
            .lock()
            .expect("reply waiter poisoned")
            .remove(&id)
        {
            tx.send(payload).is_ok()
        } else {
            false
        }
    }

    /// Drop the waiter for `id` without delivering anything.
    ///
    /// Called by the bus when a request times out — keeps the waiters
    /// map from accumulating dead entries.
    pub fn cancel(&self, id: &CorrelationId) {
        self.waiters
            .lock()
            .expect("reply waiter poisoned")
            .remove(id);
    }

    /// Number of in-flight waiters. For tests + diagnostics.
    #[must_use]
    pub fn pending(&self) -> usize {
        self.waiters.lock().expect("reply waiter poisoned").len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn correlation_ids_are_distinct() {
        let a = CorrelationId::new_random();
        let b = CorrelationId::new_random();
        assert_ne!(a, b, "two random ids must not collide");
        assert_eq!(a.hex().len(), 32);
    }

    #[test]
    fn correlation_id_hex_is_deterministic() {
        let id = CorrelationId([0xab; 16]);
        assert_eq!(id.hex(), "abababababababababababababababab");
    }

    #[tokio::test]
    async fn register_returns_receiver_completed_on_deliver() {
        let w = ReplyWaiter::new();
        let id = CorrelationId::new_random();
        let rx = w.register(id);
        assert_eq!(w.pending(), 1);
        assert!(w.deliver(id, b"reply".to_vec()));
        assert_eq!(rx.await.unwrap(), b"reply");
        assert_eq!(w.pending(), 0, "deliver removes the waiter");
    }

    #[test]
    fn deliver_unknown_id_returns_false() {
        let w = ReplyWaiter::new();
        assert!(!w.deliver(CorrelationId::new_random(), b"orphan".to_vec()));
    }

    #[tokio::test]
    async fn cancel_removes_waiter() {
        let w = ReplyWaiter::new();
        let id = CorrelationId::new_random();
        let rx = w.register(id);
        assert_eq!(w.pending(), 1);
        w.cancel(&id);
        assert_eq!(w.pending(), 0);
        // Receiver gets closed → recv resolves to Err.
        assert!(rx.await.is_err());
        // Late delivery is silently dropped.
        assert!(!w.deliver(id, b"late".to_vec()));
    }

    #[tokio::test]
    async fn deliver_twice_only_first_wins() {
        let w = ReplyWaiter::new();
        let id = CorrelationId::new_random();
        let rx = w.register(id);
        assert!(w.deliver(id, b"first".to_vec()));
        // Second delivery → no waiter left, returns false.
        assert!(!w.deliver(id, b"second".to_vec()));
        assert_eq!(rx.await.unwrap(), b"first");
    }
}