PiCloud/crates/orchestrator-core/src/inbox.rs

//! In-process `InboxRegistry` — the NATS-style request/reply
//! implementation for sync HTTP via the trigger outbox (design notes
//! §3).
//!
//! Workflow:
//!   1. Orchestrator allocates an `inbox_id`, calls
//!      `registry.register()` to get a oneshot receiver.
//!   2. Orchestrator writes an outbox row with `reply_to = inbox_id`.
//!   3. Dispatcher picks the row, runs the script, calls
//!      `registry.deliver(inbox_id, result)`.
//!   4. Orchestrator's `.await` on the receiver fires; it maps the
//!      `InboxResult` back into an HTTP response.
//!
//! `Delivered` means the receiver was alive when delivery hit. If the
//! orchestrator timed out and dropped the receiver before delivery,
//! `Abandoned` comes back — the dispatcher writes an
//! `abandoned_executions` row (design notes §3 #9).
//!
//! Cluster mode (v1.3+) swaps this for a Postgres `LISTEN/NOTIFY`-
//! based resolver; the `InboxResolver` trait stays the same.

use std::collections::HashMap;
use std::sync::Mutex;

use async_trait::async_trait;
use picloud_shared::{InboxDeliveryOutcome, InboxResolver, InboxResult};
use tokio::sync::oneshot;
use uuid::Uuid;

pub struct InboxRegistry {
    inner: Mutex<HashMap<Uuid, oneshot::Sender<InboxResult>>>,
}

impl InboxRegistry {
    #[must_use]
    pub fn new() -> Self {
        Self {
            inner: Mutex::new(HashMap::new()),
        }
    }

    /// Allocate a new inbox id and register the sender side. The
    /// caller awaits the returned `Receiver`; the dispatcher delivers
    /// the outcome via `deliver(id, …)`.
    #[must_use]
    pub fn register(&self) -> (Uuid, oneshot::Receiver<InboxResult>) {
        let id = Uuid::new_v4();
        let (tx, rx) = oneshot::channel();
        if let Ok(mut g) = self.inner.lock() {
            g.insert(id, tx);
        }
        (id, rx)
    }

    /// Cancel a pending inbox (orchestrator timed out and gave up).
    /// Drops the sender so any future `deliver` returns `Abandoned`.
    /// Returns `true` if the receiver was still registered.
    pub fn cancel(&self, id: Uuid) -> bool {
        self.inner
            .lock()
            .map(|mut g| g.remove(&id).is_some())
            .unwrap_or(false)
    }
}

impl Default for InboxRegistry {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl InboxResolver for InboxRegistry {
    async fn deliver(&self, inbox_id: Uuid, result: InboxResult) -> InboxDeliveryOutcome {
        let Ok(mut g) = self.inner.lock() else {
            return InboxDeliveryOutcome::Abandoned;
        };
        let Some(tx) = g.remove(&inbox_id) else {
            return InboxDeliveryOutcome::Abandoned;
        };
        // `send` returns Err iff the receiver was dropped — exactly
        // the abandoned-execution case.
        if tx.send(result).is_err() {
            InboxDeliveryOutcome::Abandoned
        } else {
            InboxDeliveryOutcome::Delivered
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use picloud_shared::ExecResponseSummary;
    use std::collections::BTreeMap;

    fn ok_result() -> InboxResult {
        InboxResult::Success(ExecResponseSummary {
            status_code: 200,
            headers: BTreeMap::new(),
            body: serde_json::json!({ "ok": true }),
        })
    }

    #[tokio::test]
    async fn register_then_deliver_resolves_receiver() {
        let reg = InboxRegistry::new();
        let (id, rx) = reg.register();
        let outcome = reg.deliver(id, ok_result()).await;
        assert_eq!(outcome, InboxDeliveryOutcome::Delivered);
        let received = rx.await.expect("receiver should fire");
        assert!(matches!(received, InboxResult::Success(_)));
    }

    #[tokio::test]
    async fn deliver_to_unknown_id_is_abandoned() {
        let reg = InboxRegistry::new();
        let outcome = reg.deliver(Uuid::new_v4(), ok_result()).await;
        assert_eq!(outcome, InboxDeliveryOutcome::Abandoned);
    }

    #[tokio::test]
    async fn dropping_receiver_then_delivering_is_abandoned() {
        let reg = InboxRegistry::new();
        let (id, rx) = reg.register();
        drop(rx);
        let outcome = reg.deliver(id, ok_result()).await;
        assert_eq!(outcome, InboxDeliveryOutcome::Abandoned);
    }

    #[tokio::test]
    async fn cancel_removes_sender() {
        let reg = InboxRegistry::new();
        let (id, _rx) = reg.register();
        assert!(reg.cancel(id));
        let outcome = reg.deliver(id, ok_result()).await;
        assert_eq!(outcome, InboxDeliveryOutcome::Abandoned);
    }
}