Implement support for $/cancelRequest (for client requests)

[silvanshade]

Support for canceling server requests (from the client to the server) was implemented when #145 was resolved.

However, the LSP spec for $/cancelRequest allows for server request (from the server to the client) to also be canceled.

This would especially be useful with send_custom_request per #230.

I'm not sure exactly how this would be implemented but I suppose the ClientRequests struct would need to be modified to work with future::AbortHandle, similar to ServerRequests, and a corresponding cancel method added to the impl ClientRequests.

As far as actually canceling requests goes, it would probably be good to have an interface similar to how microsoft/vscode-languageserver-node does it (see here), to where you can create a "cancellation token" prior to sending the request, and then that token would become an argument to send_request (and its derivatives).

The actual cancelation call would then be token.cancel(), which would cause the response to resolve as a RequestCancelled error (after the client properly responds to $/cancelRequest).

Does that seem feasible?

[silvanshade]

Just to expand on this a bit, I think you could define CancellationToken as something like this:

use tokio::sync::watch;

#[derive(Debug, Clone)]
enum Signal {
    Init,
    Cancel,
}

struct CancellationToken {
    tx: watch::Sender<Signal>,
    pub rx: watch::Receiver<Signal>,
}

impl CancellationToken {
    pub fn new() -> Self {
        let (tx, rx) = watch::channel(Signal::Init);
        CancellationToken {
            tx,
            rx,
        }
    }

    pub fn cancel(&self) {
        self.tx.broadcast(Signal::Cancel);
    }
}

The signature for send_request would change to something like this:

async fn send_request<R>(&self, params: R::Params, token: Option<&CancellationToken>) -> Result<R::Result>
where
    R: Request;

The idea here is that you should be able to use the same cancellation token for multiple requests (by cloning token.rx as needed). I'm not entirely sure how often that use case will come up but at least I think that the node implementation supports that.

[ebkalderon]

After a bit more digging through the specification, I actually wonder whether a CancellationToken is truly the right approach. I have two reservations with the approach outlined above:

1. Passing a CancellationToken into send_request() is great, but it can't work with the tower::Service interface, should that interface ever be implemented by Client, as recorded in Implement Service for Client #226. There would be no way to pass in a CancellationToken alongside with it, unless you were to provide an implementation for both Service<Request> and Service<(Request, CancellationToken)>. This would do the trick, but I'm not sure I'm happy with it.

2. Even if we were to go with a CancellationToken API, implementing it using purely server-side async primitives doesn't seem quite right to me. The $/cancelRequest bit of the Language Server Protocol applies to clients just as much as servers, with exactly the same semantics on both (as indicated by the ➡️ ⬅️ icons in the specification). Imagine a scenario where the user cancels a mutating client request, e.g. workspace/applyEdit, using a CancellationToken, but the operation was still silently executed on the client side. I imagine this would be pretty confusing for users.

   I think CancellationToken should instead be nothing more than a struct containing the request ID in question (maybe a Client handle too), and when the .cancel() method is called, it simply fires off a $/cancelRequest notification to the client. If it makes it to the client and request cancellation is supported, the original request will return Error::request_cancelled(). If it arrives late or the feature isn't supported on the client, you still get the result back. This fact should be made explicitly clear to users in the API documentation.

[ebkalderon]

I think there are a few possible ways forward with this task that may merit further discussion.

First, I would like to reiterate my points from the comment above: we must adhere closely to the official LSP specification and signal cancellation to the client by emitting a real $/cancelRequest notification upon calling .cancel() (cannot simply cancel a pending server-side future). As such, I imagine our final CancellationToken might look something like the following:

struct CancellationTokenInner {
    request_id: Id,
    is_canceled: AtomicBool,
    client: Client,
}

#[derive(Clone)]
pub struct CancellationToken {
    inner: Arc<CancellationTokenInner>,
}

impl CancellationToken {
    pub fn is_canceled(&self) -> bool {
        self.inner.is_canceled.load(Ordering::SeqCst)
    }

    pub async fn cancel(self)  {
        if self.inner.is_canceled.swap(true, Ordering::SeqCst) {
            // Canceled already, should not send another notification.
            return;
        }

        self.inner
            .client
            .send_notification::<Cancel>(CancelParams { id: id.into() })
            .await;
    }
}

The real question is how to expose this to the user in an ergonomic way. I see two possible options here:

Option 1: CancellationToken as an argument to every request method

This requires implementing Service<(Request, CancellationToken)> on Client and updating every request method to also accept a cancel: Option<CancellationToken> argument. Usage of such an API might look like this:

Normal usage

// Send a request.
let response: P::Result = client.send_request::<MyRequest>(params, None).await?;
   
// Send a request (with cancellation)
let token = client.cancellation_token(); // Creates a cancellation token for the next request.
let future = client.send_request::<MyRequest>(params, Some(token.clone()));
// Potentially call `token.cancel()`...
let response: P::Result = future.await?;

Using tower::Service

// Send a request.
let response: Option<Response> = client.call(request).await?;
   
// Send a request (with cancellation)
let token = client.cancellation_token(); // Creates a cancellation token for the next request.
let future = client.call((request, token.clone()));
// Potentially call `token.cancel()`...
let response: Option<Response> = future.await?;

Option 2: Introduce a concrete ResponseFuture type

Alternatively, we could change every request method on Client from a simple async fn into a regular fn that returns a nameable ResponseFuture<T> instead. This future type look something like this:

#[must_use = "futures do nothing unless polled"]
pub struct ResponseFuture<T> {
    future: BoxFuture<'static, T>,
    request_id: Id,
    client: Client,
}

impl<T> ResponseFuture<T> {
    pub fn cancellable(self) -> (impl Future<Output = T> + Send + 'static, CancellationToken) {
        let token = CancellationToken::new(self.request_id, self.client);
        (self.future, token)
    }
}

impl<T> Future for ResponseFuture<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Self::Output> { ... }
}

Given this return value, the user may choose to either .await this future immediately, or else call .cancellable() to get a CancellationToken first.

Note that JSON-RPC notifications do not have a request ID associated with them and therefore cannot be canceled. As such, Client::send_notification() and all other notification methods remain async fn under this model.

Normal usage

// Send a request.
let response: MyRequest::Result = client.send_request::<MyRequest>(params).await?;
   
// Send a request (with cancellation)
let (future, token) = client.send_request::<MyRequest>(params).cancellable();
// Potentially call `token.cancel()`...
let response: MyRequest::Result = future.await?;

Using tower::Service

// Send a request.
let response: Option<Response> = client.call(request).await?;
   
// Send a request (with cancellation)
let (future, token) = client.call(request).cancellable();
// Potentially call `token.cancel()`...
let response: Option<Response> = future.await?;

At the time of writing, I'm personally leaning slightly towards the latter pattern (option 2). This approach seems less prone to user error than option 1 (the CancellationToken request ID is guaranteed to match the request being sent) and doesn't introduce breaking changes to the Client API. It also keeps Client nicely composable with tower middleware, since there's no need to handle both Service<Request> and Service<(Request, CancellationToken)> cases.

Any thoughts on either of these proposals? I'd also be happy to review alternatives, if someone has other ideas.

[davidhalter]

Just letting you know that I like you how much thought you're putting into this. I would probably also lean to the latter pattern, but that is a very uninformed opinion, so just do whatever you think fits best.