cosigner_client/

batch.rs

//! Coalescing of concurrent single-message signing into batched proxy calls
//! under a rate budget.
//!
//! A caller that holds one message per request — an HTTP handler, say — cannot
//! use [`sign_messages`](crate::ArchSignerT::sign_messages) on its own, because
//! it never has N messages at once. [`BatchSigner`] supplies the missing piece:
//! requests wait briefly in a bounded queue, and a dispatcher groups whatever
//! has arrived into one activity.
//!
//! ```text
//!   sign_message ─┐
//!   sign_message ─┼─► queue ─► dispatcher ─► one POST /v1/sign_batch
//!   sign_message ─┘             │  ▲
//!                               │  └── rate limiter: the wait for a permit
//!                               │      IS the batching window
//!                               └───── in-flight cap on concurrent activities
//! ```
//!
//! The window is not a fixed linger: under light load a permit is free, the
//! drain finds nothing, and the batch is one message with no added latency.
//! Under load, requests pile up while the dispatcher waits for its permit, so
//! batches grow on their own and the ceiling becomes roughly
//! `max_rps × max_batch` messages per second.
//!
//! No Turnkey limit is hard-coded here: every value is a default the consumer
//! overrides.

use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;

use arch_program::pubkey::Pubkey;
use arch_program::sanitized::ArchMessage;
use async_trait::async_trait;
use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit, Semaphore};
use tokio::time::Instant;

use crate::{
    is_retryable, ArchSigner, ArchSignerT, LocalSigner, RemoteSigner, SignError, SignResponse,
};

/// Tuning for the queued path. Every field is a starting point the consumer
/// overrides; none of them encodes a Turnkey limit.
#[derive(Clone, Copy, Debug)]
struct Tuning {
    max_rps: f64,
    max_batch: usize,
    max_in_flight: usize,
    queue_depth: usize,
    deadline: Duration,
    retries: u32,
}

impl Default for Tuning {
    fn default() -> Self {
        Self {
            // Below the sub-org ceiling so retry attempts, which also spend
            // permits, and the proxy's own bounded retry have headroom.
            max_rps: 8.0,
            max_batch: 32,
            // At ~300ms per activity this sustains far more than max_rps
            // allows, so the limiter stays the binding constraint.
            max_in_flight: 4,
            queue_depth: 512,
            deadline: Duration::from_secs(2),
            retries: 2,
        }
    }
}

/// One queued signing request and the channel its result goes back on.
struct Job {
    message: ArchMessage,
    deadline: Instant,
    reply: oneshot::Sender<Result<SignResponse, SignError>>,
}

impl Job {
    fn expired(&self) -> bool {
        Instant::now() >= self.deadline
    }
}

/// Coalesces concurrent single-message signing into batched proxy calls under a
/// configurable rate budget. Implements [`ArchSignerT`], so call sites are
/// unchanged.
///
/// A local backend is passed straight through — see
/// [`is_batching`](Self::is_batching).
///
/// # Examples
///
/// ```no_run
/// use std::time::Duration;
/// use cosigner_client::{ArchSigner, ArchSignerT, BatchSigner};
///
/// # async fn run(message: arch_program::sanitized::ArchMessage)
/// #     -> Result<(), cosigner_client::SignError> {
/// let signer = BatchSigner::spawn(ArchSigner::from_env()?)
///     .with_max_rps(8.0)
///     .with_max_batch(32)
///     .with_deadline(Duration::from_millis(2000));
///
/// // Unchanged call site; concurrent calls share one proxy round trip.
/// let response = signer.sign_message(&message).await?;
/// # let _ = response;
/// # Ok(())
/// # }
/// ```
pub struct BatchSigner {
    mode: Mode,
}

enum Mode {
    /// Local signing has no rate limit and no round trip, so queueing it would
    /// only import a concurrency cap and a queue-full failure mode into a path
    /// that has neither.
    Direct(LocalSigner),
    Queued(Queued),
}

struct Queued {
    signer: Arc<RemoteSigner>,
    pubkey: Pubkey,
    tuning: Tuning,
    /// Started with the first queued request, so the builder setters shape the
    /// queue they configure even though they run after [`BatchSigner::spawn`].
    dispatcher: OnceLock<mpsc::Sender<Job>>,
}

impl BatchSigner {
    /// Creates the signer and the dispatcher that serves it.
    ///
    /// A local `inner` is passed through directly, with no task, channel, or
    /// limiter. For a remote `inner` the dispatcher task starts with the first
    /// queued request; apply the builder setters before then, since they
    /// configure a queue that is already in use afterwards.
    ///
    /// `inner`'s own retry budget is cleared: the dispatcher owns retry so that
    /// every attempt spends a rate-limiter permit.
    pub fn spawn(inner: ArchSigner) -> Self {
        let mode = match inner {
            ArchSigner::Local(local) => Mode::Direct(local),
            ArchSigner::Remote(remote) => {
                let pubkey = remote.pubkey();
                Mode::Queued(Queued {
                    // Retries here would bypass the rate limiter, turning a
                    // transient failure into a burst past the configured rate.
                    signer: Arc::new(remote.with_retries(0)),
                    pubkey,
                    tuning: Tuning::default(),
                    dispatcher: OnceLock::new(),
                })
            }
        };
        Self { mode }
    }

    /// Returns whether requests are coalesced, i.e. whether the backend is
    /// remote.
    ///
    /// Worth logging beside the resolved backend: under a local signer the
    /// batch metrics stay flat, which otherwise reads as a bug.
    pub fn is_batching(&self) -> bool {
        matches!(self.mode, Mode::Queued(_))
    }

    /// Returns this signer with the ceiling on activities per second.
    ///
    /// A non-positive or non-finite rate means unlimited.
    pub fn with_max_rps(self, rps: f64) -> Self {
        self.tuned(|t| t.max_rps = rps)
    }

    /// Returns this signer with the largest number of messages per activity.
    pub fn with_max_batch(self, n: usize) -> Self {
        self.tuned(|t| t.max_batch = n.max(1))
    }

    /// Returns this signer with the cap on concurrent in-flight activities.
    pub fn with_max_in_flight(self, n: usize) -> Self {
        self.tuned(|t| t.max_in_flight = n.max(1))
    }

    /// Returns this signer with the bound on queued requests. Beyond it,
    /// signing fails immediately rather than waiting.
    pub fn with_queue_depth(self, n: usize) -> Self {
        self.tuned(|t| t.queue_depth = n.max(1))
    }

    /// Returns this signer with the per-request deadline. A request still
    /// queued when it expires is dropped rather than signed.
    pub fn with_deadline(self, deadline: Duration) -> Self {
        self.tuned(|t| t.deadline = deadline)
    }

    /// Returns this signer with the dispatcher's retry budget for a whole
    /// batch. Each attempt spends a rate-limiter permit.
    pub fn with_retries(self, retries: u32) -> Self {
        self.tuned(|t| t.retries = retries)
    }

    /// Applies `set` to the queued tuning; a no-op in direct mode, where there
    /// is no budget to configure.
    fn tuned(mut self, set: impl FnOnce(&mut Tuning)) -> Self {
        if let Mode::Queued(queued) = &mut self.mode {
            set(&mut queued.tuning);
        }
        self
    }
}

impl std::fmt::Debug for BatchSigner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut out = f.debug_struct("BatchSigner");
        out.field("batching", &self.is_batching());
        if let Mode::Queued(queued) = &self.mode {
            out.field("tuning", &queued.tuning);
        }
        out.finish()
    }
}

impl Queued {
    fn sender(&self) -> &mpsc::Sender<Job> {
        self.dispatcher.get_or_init(|| {
            let (tx, rx) = mpsc::channel(self.tuning.queue_depth);
            let signer = Arc::clone(&self.signer);
            let tuning = self.tuning;
            tokio::spawn(dispatch_loop(signer, rx, tuning));
            tx
        })
    }

    /// Queues `message`, returning the channel its result arrives on.
    ///
    /// Kept separate from awaiting so a caller can enqueue a whole slice before
    /// blocking on any of it, letting the dispatcher batch the lot.
    ///
    /// # Errors
    /// [`SignError::Signing`] when the queue is full or the dispatcher has
    /// stopped. A fast rejection beats unbounded latency when the caller's work
    /// expires.
    fn enqueue(
        &self,
        message: &ArchMessage,
    ) -> Result<oneshot::Receiver<Result<SignResponse, SignError>>, SignError> {
        let (reply, wait) = oneshot::channel();
        let job = Job {
            message: message.clone(),
            deadline: Instant::now() + self.tuning.deadline,
            reply,
        };
        self.sender().try_send(job).map_err(|e| match e {
            mpsc::error::TrySendError::Full(_) => SignError::Signing("batch queue full".into()),
            mpsc::error::TrySendError::Closed(_) => {
                SignError::Signing("batch dispatcher stopped".into())
            }
        })?;
        Ok(wait)
    }

    /// Awaits one queued result under the deadline.
    ///
    /// The timeout is enforced here as well as at drain time, so a dispatcher
    /// that died cannot leave a caller waiting forever.
    async fn reply(
        &self,
        wait: oneshot::Receiver<Result<SignResponse, SignError>>,
    ) -> Result<SignResponse, SignError> {
        match tokio::time::timeout(self.tuning.deadline, wait).await {
            Ok(Ok(result)) => result,
            Ok(Err(_)) => Err(SignError::Signing("batch dispatcher stopped".into())),
            Err(_) => Err(SignError::Signing("batch deadline exceeded".into())),
        }
    }
}

#[async_trait]
impl ArchSignerT for BatchSigner {
    fn pubkey(&self) -> Pubkey {
        match &self.mode {
            Mode::Direct(local) => local.pubkey(),
            Mode::Queued(queued) => queued.pubkey,
        }
    }

    async fn sign_message(&self, message: &ArchMessage) -> Result<SignResponse, SignError> {
        match &self.mode {
            Mode::Direct(local) => local.sign_message(message).await,
            Mode::Queued(queued) => {
                let wait = queued.enqueue(message)?;
                queued.reply(wait).await
            }
        }
    }

    async fn sign_messages(
        &self,
        messages: &[ArchMessage],
    ) -> Result<Vec<Result<SignResponse, SignError>>, SignError> {
        let queued = match &self.mode {
            Mode::Direct(local) => return local.sign_messages(messages).await,
            Mode::Queued(queued) => queued,
        };

        // Every message is queued as its own job rather than forwarded as one
        // batch, so the rate budget stays authoritative. All of them are
        // enqueued before anything is awaited, so the dispatcher can still
        // group them.
        let waits: Vec<_> = messages.iter().map(|m| queued.enqueue(m)).collect();
        let mut results = Vec::with_capacity(waits.len());
        for wait in waits {
            results.push(match wait {
                Ok(wait) => queued.reply(wait).await,
                Err(err) => Err(err),
            });
        }
        Ok(results)
    }
}

/// Groups queued jobs into activities under the rate and concurrency budgets.
///
/// Blocking on a permit before draining is what creates the batching window:
/// whatever arrives during the wait rides along in the next activity.
async fn dispatch_loop(signer: Arc<RemoteSigner>, mut rx: mpsc::Receiver<Job>, tuning: Tuning) {
    let in_flight = Arc::new(Semaphore::new(tuning.max_in_flight));
    let limiter = Arc::new(RateLimiter::new(tuning.max_rps));

    while let Some(first) = rx.recv().await {
        let Ok(slot) = Arc::clone(&in_flight).acquire_owned().await else {
            break;
        };
        limiter.acquire().await;

        let mut batch = vec![first];
        while batch.len() < tuning.max_batch {
            match rx.try_recv() {
                Ok(job) => batch.push(job),
                Err(_) => break,
            }
        }

        // A request whose deadline passed while it waited is worthless; drop it
        // rather than spend a signature on it.
        let (batch, expired): (Vec<Job>, Vec<Job>) =
            batch.into_iter().partition(|job| !job.expired());
        for job in expired {
            let _ = job
                .reply
                .send(Err(SignError::Signing("batch deadline exceeded".into())));
        }
        if batch.is_empty() {
            continue;
        }

        // Dispatch off the loop so the network round trip never blocks the next
        // batch from forming.
        tokio::spawn(dispatch(
            Arc::clone(&signer),
            Arc::clone(&limiter),
            batch,
            tuning.retries,
            slot,
        ));
    }
}

/// Signs one batch and answers every job in it.
///
/// `_slot` is held for the whole activity, so dropping it on return is what
/// frees the in-flight budget.
async fn dispatch(
    signer: Arc<RemoteSigner>,
    limiter: Arc<RateLimiter>,
    batch: Vec<Job>,
    retries: u32,
    _slot: OwnedSemaphorePermit,
) {
    let messages: Vec<ArchMessage> = batch.iter().map(|job| job.message.clone()).collect();

    let mut attempt = 0;
    let outcome = loop {
        match signer.sign_messages(&messages).await {
            Ok(results) => break Ok(results),
            Err(err) if attempt < retries && is_retryable(&err) => {
                attempt += 1;
                // A retry spends a permit like any other attempt, so a
                // transient failure cannot burst past the configured rate.
                limiter.acquire().await;
            }
            Err(err) => break Err(err),
        }
    };

    match outcome {
        Ok(results) => {
            for (job, result) in batch.into_iter().zip(results) {
                let _ = job.reply.send(result);
            }
        }
        Err(err) => {
            for job in batch {
                let _ = job.reply.send(Err(err.clone()));
            }
        }
    }
}

/// Admits at most `max_rps` acquisitions per second by spacing them evenly.
///
/// Even spacing rather than a refilling bucket: a burst is exactly what the
/// per-sub-organization rate limit rejects, so there is nothing to gain by
/// allowing one.
struct RateLimiter {
    interval: Duration,
    /// Earliest instant the next acquisition may proceed.
    next: Mutex<Option<Instant>>,
}

impl RateLimiter {
    fn new(max_rps: f64) -> Self {
        // A non-positive or non-finite rate would otherwise mean "never admit",
        // stalling the dispatcher; treat it as unlimited.
        let interval = if max_rps.is_finite() && max_rps > 0.0 {
            Duration::from_secs_f64(1.0 / max_rps)
        } else {
            Duration::ZERO
        };
        Self {
            interval,
            next: Mutex::new(None),
        }
    }

    async fn acquire(&self) {
        if self.interval.is_zero() {
            return;
        }
        let at = {
            let mut next = self
                .next
                .lock()
                .expect("no code panics while holding the rate-limiter lock");
            let now = Instant::now();
            let at = next.map_or(now, |scheduled| scheduled.max(now));
            *next = Some(at + self.interval);
            at
        };
        tokio::time::sleep_until(at).await;
    }
}

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

    #[test]
    fn tuning_defaults_leave_headroom_under_the_sub_org_ceiling() {
        let tuning = Tuning::default();
        assert!(tuning.max_rps < 10.0, "retries also spend permits");
        assert!(tuning.max_batch > 1);
        assert!(tuning.queue_depth > tuning.max_batch);
    }

    #[test]
    fn zero_and_negative_rates_are_unlimited_rather_than_stalled() {
        for rps in [0.0, -1.0, f64::NAN, f64::INFINITY] {
            assert!(
                RateLimiter::new(rps).interval.is_zero(),
                "rps {rps} must not stall the dispatcher"
            );
        }
    }

    #[test]
    fn rate_is_the_reciprocal_of_the_interval() {
        assert_eq!(RateLimiter::new(8.0).interval, Duration::from_millis(125));
        assert_eq!(RateLimiter::new(2.0).interval, Duration::from_millis(500));
    }

    #[tokio::test]
    async fn spacing_is_cumulative_across_acquisitions() {
        let limiter = RateLimiter::new(50.0);
        let started = Instant::now();
        for _ in 0..4 {
            limiter.acquire().await;
        }
        // First is immediate, then three 20ms gaps.
        assert!(
            started.elapsed() >= Duration::from_millis(55),
            "elapsed {:?}",
            started.elapsed()
        );
    }
}