link/sidecar/src/rpc.rs

//! Request/reply correlation over the one shard socket.
//!
//! REST is synchronous ("give me this character"), the shard link is an async stream of lines. This
//! bridges them: a call registers a pending entry under a correlation id, sends the command, and
//! awaits a reply carrying that id. The event loop routes any incoming line whose correlation id is
//! pending back to the waiting caller; everything else flows on as a normal event.
//!
//! Three correlation fields are recognized, matching what the plugin echoes: `reqId` (queries —
//! char/roster/vendor), `code` (link.confirm → link.ok/error), and `id` (towncrier). A query's
//! `reqId` is a process-unique counter; `code`/`id` are supplied by the caller and must be unique
//! while outstanding.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use serde_json::Value;
use tokio::sync::{oneshot, Mutex};

use crate::shard::ShardHandle;

const REPLY_TIMEOUT: Duration = Duration::from_secs(10);

#[derive(Clone)]
pub struct Rpc {
    pending: Arc<Mutex<HashMap<String, oneshot::Sender<Value>>>>,
    counter: Arc<AtomicU64>,
}

#[derive(Debug)]
pub enum RpcError {
    NoShard,
    Timeout,
}

impl Rpc {
    pub fn new() -> Self {
        Self {
            pending: Arc::new(Mutex::new(HashMap::new())),
            counter: Arc::new(AtomicU64::new(1)),
        }
    }

    pub fn next_req_id(&self) -> String {
        format!("r-{}", self.counter.fetch_add(1, Ordering::Relaxed))
    }

    /// Sends `command` to the shard and awaits the reply correlated by `corr_val`. The command must
    /// already contain the correlation field (e.g. `reqId`) set to `corr_val`.
    pub async fn call(
        &self,
        shard: &ShardHandle,
        command: Value,
        corr_val: &str,
    ) -> Result<Value, RpcError> {
        let (tx, rx) = oneshot::channel();

        self.pending.lock().await.insert(corr_val.to_string(), tx);

        if !shard.send(command.to_string()).await {
            self.pending.lock().await.remove(corr_val);
            return Err(RpcError::NoShard);
        }

        match tokio::time::timeout(REPLY_TIMEOUT, rx).await {
            Ok(Ok(value)) => Ok(value),
            _ => {
                self.pending.lock().await.remove(corr_val);
                Err(RpcError::Timeout)
            }
        }
    }

    /// If this incoming value correlates to a pending call, complete it and return true (the value
    /// was a reply, not a broadcast event). Otherwise return false.
    pub async fn try_route(&self, value: &Value) -> bool {
        let candidate = ["reqId", "code", "id"]
            .iter()
            .find_map(|k| value.get(*k).and_then(|v| v.as_str()).map(str::to_string));

        let corr = match candidate {
            Some(c) => c,
            None => return false,
        };

        let sender = self.pending.lock().await.remove(&corr);
        match sender {
            Some(tx) => {
                let _ = tx.send(value.clone());
                true
            }
            None => false,
        }
    }
}