pending_requests/

lib.rs

//! Track in-flight requests and await their responses by key.
//!
//! A [`PendingRequests`] registry lets one task register a request under a
//! key, hand back a [`ResponseWaiter`] future, and have another task deliver
//! the matching response later. This is the classic pattern behind
//! request/response protocols multiplexed over a single connection.
//!
//! ```
//! # use std::time::Duration;
//! # use pending_requests::PendingRequests;
//! # #[tokio::main]
//! # async fn main() {
//! let requests = PendingRequests::<u64, String>::new();
//! let waiter = requests.prepare_response(1).unwrap();
//!
//! // ... elsewhere, when the response for key `1` arrives:
//! requests.handle_response(1, "pong".to_string()).unwrap();
//!
//! assert_eq!(waiter.await.unwrap(), "pong");
//! # }
//! ```

use std::{
    collections::HashMap,
    future::Future,
    hash::Hash,
    pin::Pin,
    sync::{
        Arc, Weak,
        atomic::{AtomicU64, Ordering},
    },
    task::{Context, Poll},
    time::Duration,
};

use parking_lot::Mutex;
use tokio::{
    sync::oneshot,
    time::{Sleep, sleep},
};

mod error;

pub use error::Error;
use error::Result;

/// One stored slot in the registry. The `id` distinguishes successive
/// requests that reuse the same key, so a stale waiter never removes an
/// entry that belongs to a newer request.
struct Entry<R> {
    id: u64,
    sender: oneshot::Sender<R>,
}

/// Shared state behind the registry. Cloning a [`PendingRequests`] clones the
/// `Arc` to this, so every clone observes the same pending requests, timeout,
/// and id counter.
struct Inner<K: Eq + Hash, R> {
    /// Default timeout in milliseconds, mutable at runtime via an atomic so
    /// it can be changed through a shared `&self`.
    timeout_ms: AtomicU64,
    next_id: AtomicU64,
    requests: Mutex<HashMap<K, Entry<R>>>,
}

/// A future that resolves with the response for a single request.
///
/// Resolves to:
/// - `Ok(response)` when [`PendingRequests::handle_response`] delivers a value,
/// - `Err(Error::RequestTimeout)` when the configured timeout elapses first,
/// - `Err(Error::Canceled)` when the request is canceled or the registry is dropped.
pub struct ResponseWaiter<K: Eq + Hash, R> {
    receiver: oneshot::Receiver<R>,
    sleep: Pin<Box<Sleep>>,
    inner: Weak<Inner<K, R>>,
    key: K,
    id: u64,
    /// Set once the future has resolved, so `Drop` skips cleanup.
    done: bool,
}

/// The `Sleep` is the only `!Unpin` field and it is kept behind a
/// `Pin<Box<_>>`; nothing in this struct is self-referential, so it is safe
/// to treat the waiter itself as `Unpin` regardless of `K`/`R`.
impl<K: Eq + Hash, R> Unpin for ResponseWaiter<K, R> {}

impl<K: Eq + Hash, R> ResponseWaiter<K, R> {
    /// Remove this waiter's entry from the registry, but only if it is still
    /// the same generation we registered (guards against a reused key).
    fn clear(&mut self) {
        if let Some(inner) = self.inner.upgrade() {
            let mut guard = inner.requests.lock();
            if guard.get(&self.key).is_some_and(|e| e.id == self.id) {
                guard.remove(&self.key);
            }
        }
    }
}

impl<K: Eq + Hash, R> Future for ResponseWaiter<K, R> {
    type Output = Result<R>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // All fields are `Unpin` (the `Sleep` lives behind a `Pin<Box<_>>`),
        // so we can work with a plain mutable reference.
        let this = self.get_mut();

        // A response (or a closed channel) takes priority over the timeout.
        match Pin::new(&mut this.receiver).poll(cx) {
            Poll::Ready(Ok(response)) => {
                this.done = true;
                return Poll::Ready(Ok(response));
            }
            Poll::Ready(Err(_)) => {
                // Sender was dropped without sending: canceled or registry gone.
                this.done = true;
                return Poll::Ready(Err(Error::Canceled));
            }
            Poll::Pending => {}
        }

        if this.sleep.as_mut().poll(cx).is_ready() {
            this.clear();
            this.done = true;
            return Poll::Ready(Err(Error::RequestTimeout));
        }

        Poll::Pending
    }
}

impl<K: Eq + Hash, R> Drop for ResponseWaiter<K, R> {
    fn drop(&mut self) {
        if !self.done {
            self.clear();
        }
    }
}

/// A registry of in-flight requests awaiting responses, keyed by `K`.
///
/// Cloning is cheap and shares all state: register a request on one clone and
/// deliver its response from another. This is the expected way to use the
/// registry across tasks (e.g. a reader task calling [`handle_response`] while
/// many sender tasks await their [`ResponseWaiter`]s).
///
/// [`handle_response`]: PendingRequests::handle_response
pub struct PendingRequests<K: Eq + Hash, R> {
    inner: Arc<Inner<K, R>>,
}

impl<K: Eq + Hash, R> Clone for PendingRequests<K, R> {
    fn clone(&self) -> Self {
        Self {
            inner: Arc::clone(&self.inner),
        }
    }
}

const DEFAULT_TIMEOUT_MS: u64 = 6000;

impl<K: Eq + Hash, R> Default for PendingRequests<K, R> {
    fn default() -> Self {
        Self {
            inner: Arc::new(Inner {
                timeout_ms: AtomicU64::new(DEFAULT_TIMEOUT_MS),
                next_id: AtomicU64::new(0),
                requests: Mutex::new(HashMap::new()),
            }),
        }
    }
}

impl<K: Eq + Hash + Clone, R> PendingRequests<K, R> {
    /// Create a registry with the default 6s timeout.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the default timeout applied to requests registered *after* this
    /// call. Takes `&self`, so it works through a shared/cloned handle.
    /// Already-registered requests keep the timeout they were created with.
    pub fn set_timeout(&self, timeout: Duration) {
        self.inner
            .timeout_ms
            .store(timeout.as_millis() as u64, Ordering::Relaxed);
    }

    /// The current default timeout.
    pub fn timeout(&self) -> Duration {
        Duration::from_millis(self.inner.timeout_ms.load(Ordering::Relaxed))
    }

    /// Number of requests currently awaiting a response.
    pub fn len(&self) -> usize {
        self.inner.requests.lock().len()
    }

    /// Whether there are no requests currently awaiting a response.
    pub fn is_empty(&self) -> bool {
        self.inner.requests.lock().is_empty()
    }

    /// Whether a request is currently pending for `key`.
    pub fn contains(&self, key: &K) -> bool {
        self.inner.requests.lock().contains_key(key)
    }

    /// Register a request under `key`, using the registry's default timeout.
    ///
    /// Returns [`Error::KeyAlreadyExists`] if a request with `key` is already
    /// pending.
    pub fn prepare_response(&self, key: K) -> Result<ResponseWaiter<K, R>> {
        self.prepare_response_with_timeout(key, self.timeout())
    }

    /// Register a request under `key` with an explicit `timeout`.
    pub fn prepare_response_with_timeout(
        &self,
        key: K,
        timeout: Duration,
    ) -> Result<ResponseWaiter<K, R>> {
        let mut guard = self.inner.requests.lock();
        if guard.contains_key(&key) {
            return Err(Error::KeyAlreadyExists);
        }

        let id = self.inner.next_id.fetch_add(1, Ordering::Relaxed);
        let (tx, rx) = oneshot::channel();
        guard.insert(key.clone(), Entry { id, sender: tx });

        Ok(ResponseWaiter {
            receiver: rx,
            sleep: Box::pin(sleep(timeout)),
            inner: Arc::downgrade(&self.inner),
            key,
            id,
            done: false,
        })
    }

    /// Deliver `response` to the request registered under `key`.
    ///
    /// Returns [`Error::KeyNotFound`] if no request is pending for `key`, or
    /// [`Error::ReceiverDropped`] if the waiter was already dropped.
    pub fn handle_response(&self, key: K, response: R) -> Result<()> {
        let entry = self
            .inner
            .requests
            .lock()
            .remove(&key)
            .ok_or(Error::KeyNotFound)?;

        entry
            .sender
            .send(response)
            .map_err(|_| Error::ReceiverDropped)
    }

    /// Cancel the request registered under `key`, causing its waiter to
    /// resolve with [`Error::Canceled`].
    ///
    /// Returns [`Error::KeyNotFound`] if no request is pending for `key`.
    pub fn cancel(&self, key: &K) -> Result<()> {
        self.inner
            .requests
            .lock()
            .remove(key)
            .map(|_| ())
            .ok_or(Error::KeyNotFound)
    }

    /// Cancel every pending request, resolving all their waiters with
    /// [`Error::Canceled`]. Returns the number of requests canceled.
    ///
    /// Useful when a connection drops and all in-flight requests should fail
    /// at once.
    pub fn cancel_all(&self) -> usize {
        let mut guard = self.inner.requests.lock();
        let n = guard.len();
        guard.clear();
        n
    }
}

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

    #[tokio::test]
    async fn roundtrip_delivers_response() {
        let requests = PendingRequests::<u64, &str>::new();
        let waiter = requests.prepare_response(1).unwrap();

        requests.handle_response(1, "pong").unwrap();

        assert_eq!(waiter.await.unwrap(), "pong");
        assert!(requests.is_empty());
    }

    #[tokio::test]
    async fn duplicate_key_is_rejected() {
        let requests = PendingRequests::<u64, ()>::new();
        let _waiter = requests.prepare_response(1).unwrap();

        assert!(matches!(
            requests.prepare_response(1),
            Err(Error::KeyAlreadyExists)
        ));
    }

    #[tokio::test]
    async fn unknown_key_response_errors() {
        let requests = PendingRequests::<u64, ()>::new();
        assert!(matches!(
            requests.handle_response(42, ()),
            Err(Error::KeyNotFound)
        ));
    }

    #[tokio::test(start_paused = true)]
    async fn timeout_resolves_and_cleans_up() {
        let requests = PendingRequests::<u64, ()>::new();
        let waiter = requests
            .prepare_response_with_timeout(1, Duration::from_millis(50))
            .unwrap();

        assert!(matches!(waiter.await, Err(Error::RequestTimeout)));
        assert!(requests.is_empty(), "timed-out entry should be removed");
    }

    #[tokio::test]
    async fn cancel_resolves_waiter() {
        let requests = PendingRequests::<u64, ()>::new();
        let waiter = requests.prepare_response(1).unwrap();

        requests.cancel(&1).unwrap();

        assert!(matches!(waiter.await, Err(Error::Canceled)));
        assert!(matches!(requests.cancel(&1), Err(Error::KeyNotFound)));
    }

    #[tokio::test]
    async fn dropped_waiter_is_reported_and_cleaned() {
        let requests = PendingRequests::<u64, ()>::new();
        let waiter = requests.prepare_response(1).unwrap();

        drop(waiter);

        assert!(
            requests.is_empty(),
            "dropped waiter should remove its entry"
        );
        assert!(matches!(
            requests.handle_response(1, ()),
            Err(Error::KeyNotFound)
        ));
    }

    #[tokio::test]
    async fn dropping_registry_cancels_waiter() {
        let requests = PendingRequests::<u64, ()>::new();
        let waiter = requests.prepare_response(1).unwrap();

        drop(requests);

        assert!(matches!(waiter.await, Err(Error::Canceled)));
    }

    #[tokio::test(start_paused = true)]
    async fn reused_key_after_timeout_is_independent() {
        let requests = PendingRequests::<u64, &str>::new();
        let first = requests
            .prepare_response_with_timeout(1, Duration::from_millis(50))
            .unwrap();
        assert!(matches!(first.await, Err(Error::RequestTimeout)));

        // The same key can now be registered again and resolved normally.
        let second = requests.prepare_response(1).unwrap();
        requests.handle_response(1, "ok").unwrap();
        assert_eq!(second.await.unwrap(), "ok");
    }

    #[tokio::test]
    async fn clones_share_state() {
        let requests = PendingRequests::<u64, &str>::new();
        let other = requests.clone();

        // Register on one handle, respond from the clone.
        let waiter = requests.prepare_response(1).unwrap();
        other.handle_response(1, "pong").unwrap();

        assert_eq!(waiter.await.unwrap(), "pong");
        assert!(other.is_empty());
    }

    #[tokio::test]
    async fn contains_reflects_pending_state() {
        let requests = PendingRequests::<u64, ()>::new();
        assert!(!requests.contains(&1));

        let _waiter = requests.prepare_response(1).unwrap();
        assert!(requests.contains(&1));

        requests.cancel(&1).unwrap();
        assert!(!requests.contains(&1));
    }

    #[tokio::test]
    async fn cancel_all_fails_every_waiter() {
        let requests = PendingRequests::<u64, ()>::new();
        let a = requests.prepare_response(1).unwrap();
        let b = requests.prepare_response(2).unwrap();

        assert_eq!(requests.cancel_all(), 2);
        assert!(requests.is_empty());
        assert!(matches!(a.await, Err(Error::Canceled)));
        assert!(matches!(b.await, Err(Error::Canceled)));
    }

    #[tokio::test]
    async fn set_timeout_works_through_shared_handle() {
        let requests = PendingRequests::<u64, ()>::new();
        let handle = requests.clone();

        handle.set_timeout(Duration::from_secs(30));
        assert_eq!(requests.timeout(), Duration::from_secs(30));
    }

    #[tokio::test]
    async fn surviving_clone_keeps_waiter_alive() {
        let requests = PendingRequests::<u64, &str>::new();
        let clone = requests.clone();
        let waiter = requests.prepare_response(1).unwrap();

        // Dropping one handle must not cancel requests while a clone lives.
        drop(requests);
        clone.handle_response(1, "still here").unwrap();
        assert_eq!(waiter.await.unwrap(), "still here");
    }
}