lsp_max/service/
pending.rs

1//! Types for tracking cancelable client-to-server JSON-RPC requests.
2
3use std::fmt::{self, Debug, Formatter};
4use std::future::Future;
5use std::sync::Arc;
6
7use dashmap::{mapref::entry::Entry, DashMap};
8use futures::future::{self, Either};
9use tracing::{debug, info};
10
11use super::ExitedError;
12use crate::jsonrpc::{Error, Id, Response};
13
14/// A hashmap containing pending server requests, keyed by request ID.
15#[derive(Default)]
16pub struct Pending(Arc<DashMap<Id, future::AbortHandle>>);
17
18impl Pending {
19    /// Creates a new pending server requests map.
20    pub fn new() -> Self {
21        Pending(Arc::new(DashMap::new()))
22    }
23
24    /// Executes the given async request handler, keyed by the given request ID.
25    ///
26    /// If a cancel request is issued before the future is finished resolving, this will resolve to
27    /// a "canceled" error response, and the pending request handler future will be dropped.
28    pub fn execute<F>(
29        &self,
30        id: Id,
31        fut: F,
32    ) -> impl Future<Output = Result<Option<Response>, ExitedError>> + Send + 'static
33    where
34        F: Future<Output = Result<Option<Response>, ExitedError>> + Send + 'static,
35    {
36        if let Entry::Vacant(entry) = self.0.entry(id.clone()) {
37            let (handler_fut, abort_handle) = future::abortable(fut);
38            entry.insert(abort_handle);
39
40            let requests = self.0.clone();
41            Either::Left(async move {
42                let abort_result = handler_fut.await;
43                requests.remove(&id); // Remove abort handle now to avoid double cancellation.
44
45                if let Ok(handler_result) = abort_result {
46                    handler_result
47                } else {
48                    Ok(Some(Response::from_error(id, Error::request_cancelled())))
49                }
50            })
51        } else {
52            Either::Right(async { Ok(Some(Response::from_error(id, Error::invalid_request()))) })
53        }
54    }
55
56    /// Attempts to cancel the running request handler corresponding to this ID.
57    ///
58    /// This will force the future to resolve to a "canceled" error response. If the future has
59    /// already completed, this method call will do nothing.
60    pub fn cancel(&self, id: &Id) {
61        if let Some((_, handle)) = self.0.remove(id) {
62            handle.abort();
63            info!("successfully cancelled request with ID: {}", id);
64        } else {
65            debug!(
66                "client asked to cancel request {}, but no such pending request exists, ignoring",
67                id
68            );
69        }
70    }
71
72    /// Cancels all pending request handlers, if any.
73    pub fn cancel_all(&self) {
74        self.0.retain(|_, handle| {
75            handle.abort();
76            false
77        });
78    }
79}
80
81impl Debug for Pending {
82    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
83        f.debug_set()
84            .entries(self.0.iter().map(|entry| entry.key().clone()))
85            .finish()
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use serde_json::json;
92
93    use super::*;
94
95    #[tokio::test(flavor = "current_thread")]
96    async fn executes_server_request() {
97        let pending = Pending::new();
98
99        let id = Id::Number(1);
100        let id2 = id.clone();
101        let response = pending
102            .execute(id.clone(), async {
103                Ok(Some(Response::from_ok(id2, json!({}))))
104            })
105            .await;
106
107        assert_eq!(response, Ok(Some(Response::from_ok(id, json!({})))));
108    }
109
110    #[tokio::test(flavor = "current_thread")]
111    async fn cancels_server_request() {
112        let pending = Pending::new();
113
114        let id = Id::Number(1);
115        let handler_fut = tokio::spawn(pending.execute(id.clone(), future::pending()));
116
117        pending.cancel(&id);
118
119        let res = handler_fut.await.expect("task panicked");
120        assert_eq!(
121            res,
122            Ok(Some(Response::from_error(id, Error::request_cancelled())))
123        );
124    }
125}
lsp_max/service/pending.rs

lsp_max/service/
pending.rs
