lspkit_sidecar/
correlator.rs

1//! Request-ID correlation.
2//!
3//! Out-of-process backends respond asynchronously and out of order. The
4//! correlator hands out monotonic request IDs and maps responses back to the
5//! waiting future.
6
7use std::collections::HashMap;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::sync::Mutex;
10
11use tokio::sync::oneshot;
12
13/// Monotonic request identifier.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
15pub struct RequestId(u64);
16
17impl RequestId {
18    /// Raw value.
19    #[must_use]
20    pub const fn get(self) -> u64 {
21        self.0
22    }
23}
24
25/// Errors from the correlator.
26#[non_exhaustive]
27#[derive(Debug, thiserror::Error)]
28pub enum CorrelationError {
29    /// A response arrived for an unknown request ID.
30    #[error("unknown request id: {0}")]
31    UnknownId(u64),
32    /// The pending future was already resolved.
33    #[error("request {0} already completed")]
34    AlreadyResolved(u64),
35}
36
37/// Pairs outbound requests with inbound responses by ID.
38#[derive(Default)]
39pub struct Correlator<R> {
40    next_id: AtomicU64,
41    pending: Mutex<HashMap<RequestId, oneshot::Sender<R>>>,
42}
43
44impl<R> Correlator<R> {
45    /// New correlator with no pending requests.
46    #[must_use]
47    pub fn new() -> Self {
48        Self {
49            next_id: AtomicU64::new(1),
50            pending: Mutex::new(HashMap::new()),
51        }
52    }
53
54    /// Number of pending requests.
55    #[must_use]
56    pub fn pending_count(&self) -> usize {
57        self.pending.lock().map_or(0, |g| g.len())
58    }
59}
60
61impl<R: Send + 'static> Correlator<R> {
62    /// Allocate a new request ID and a receiver that will resolve when the
63    /// matching response arrives.
64    pub fn register(&self) -> (RequestId, oneshot::Receiver<R>) {
65        let id = RequestId(self.next_id.fetch_add(1, Ordering::SeqCst));
66        let (tx, rx) = oneshot::channel();
67        if let Ok(mut guard) = self.pending.lock() {
68            guard.insert(id, tx);
69        }
70        (id, rx)
71    }
72
73    /// Resolve a pending request with `response`.
74    ///
75    /// # Errors
76    /// Returns [`CorrelationError::UnknownId`] when no pending request
77    /// matches, or [`CorrelationError::AlreadyResolved`] if the receiver was
78    /// dropped.
79    pub fn resolve(&self, id: RequestId, response: R) -> Result<(), CorrelationError> {
80        let sender = self
81            .pending
82            .lock()
83            .ok()
84            .and_then(|mut guard| guard.remove(&id))
85            .ok_or(CorrelationError::UnknownId(id.get()))?;
86        sender
87            .send(response)
88            .map_err(|_| CorrelationError::AlreadyResolved(id.get()))
89    }
90}
91
92impl<R> std::fmt::Debug for Correlator<R> {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        f.debug_struct("Correlator")
95            .field("pending", &self.pending_count())
96            .finish()
97    }
98}
lspkit_sidecar/correlator.rs

lspkit_sidecar/
correlator.rs
