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, …)`.
    ///
    /// F-Q-006: panic on poison — a poisoned mutex means a prior panic
    /// inside the lock-holding code path, so the process is unrecoverable.
    /// Silently swallowing it (the old behaviour) leaked a dead id: the
    /// sender never got inserted, the later `deliver(id, …)` saw no
    /// entry and returned Abandoned, and the caller's `await rx` blocked
    /// until the orchestrator's outer timeout. Sister registries
    /// (RouteTable, AppDomainTable) already `.expect()` here — settle
    /// on the same loud policy.
    #[must_use]
    pub fn register(&self) -> (Uuid, oneshot::Receiver<InboxResult>) {
        let id = Uuid::new_v4();
        let (tx, rx) = oneshot::channel();
        let mut g = self.inner.lock().expect("inbox registry poisoned");
        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 {
        let mut g = self.inner.lock().expect("inbox registry poisoned");
        g.remove(&id).is_some()
    }
}

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 {
        // F-Q-006 + F-Q-010: panic on poison, matching the registry's
        // register/cancel policy and the sister registries.
        let mut g = self.inner.lock().expect("inbox registry poisoned");
        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);
    }
}