-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Wrap Connection to facilitate proper server shutdown
- Loading branch information
1 parent
4779dd1
commit 7d0fba9
Showing
5 changed files
with
183 additions
and
53 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,138 @@ | ||
use lsp::IoThreads; | ||
use lsp_server as lsp; | ||
use lsp_types::{notification::Notification, request::Request}; | ||
use std::sync::{Arc, Weak}; | ||
|
||
type ConnectionSender = crossbeam::channel::Sender<lsp::Message>; | ||
type ConnectionReceiver = crossbeam::channel::Receiver<lsp::Message>; | ||
|
||
pub(crate) struct ConnectionInitializer { | ||
connection: lsp::Connection, | ||
threads: lsp::IoThreads, | ||
} | ||
|
||
pub(crate) struct Connection { | ||
sender: Arc<ConnectionSender>, | ||
receiver: ConnectionReceiver, | ||
threads: lsp::IoThreads, | ||
} | ||
|
||
impl ConnectionInitializer { | ||
/// Create a new LSP server connection over stdin/stdout. | ||
pub(super) fn stdio() -> Self { | ||
let (connection, threads) = lsp::Connection::stdio(); | ||
Self { | ||
connection, | ||
threads, | ||
} | ||
} | ||
|
||
/// Starts the initialization process with the client by listening for an intialization request. | ||
/// Returns a request ID that should be passed into `initialize_finish` later, | ||
/// along with the initialization parameters that were provided. | ||
pub(super) fn initialize_start( | ||
&self, | ||
) -> crate::Result<(lsp::RequestId, lsp_types::InitializeParams)> { | ||
let (id, params) = self.connection.initialize_start()?; | ||
Ok((id, serde_json::from_value(params)?)) | ||
} | ||
|
||
/// Finishes the initialization process with the client, | ||
/// returning an initialized `Connection`. | ||
pub(super) fn initialize_finish( | ||
self, | ||
id: lsp::RequestId, | ||
server_capabilities: &lsp_types::ServerCapabilities, | ||
name: &str, | ||
version: &str, | ||
) -> crate::Result<Connection> { | ||
self.connection.initialize_finish( | ||
id, | ||
serde_json::json!({ | ||
"capabilities": server_capabilities, | ||
"serverInfo": { | ||
"name": name, | ||
"version": version | ||
} | ||
}), | ||
)?; | ||
Ok(Connection::initialized(self.connection, self.threads)) | ||
} | ||
} | ||
|
||
impl Connection { | ||
fn initialized( | ||
lsp::Connection { sender, receiver }: lsp::Connection, | ||
threads: IoThreads, | ||
) -> Self { | ||
Self { | ||
sender: Arc::new(sender), | ||
receiver, | ||
threads, | ||
} | ||
} | ||
pub(super) fn make_sender(&self) -> ClientSender { | ||
ClientSender { | ||
weak_sender: Arc::downgrade(&self.sender), | ||
} | ||
} | ||
|
||
pub(super) fn receive_messages(&self) -> crossbeam::channel::Iter<lsp::Message> { | ||
self.receiver.iter() | ||
} | ||
|
||
/// | ||
pub(super) fn handle_shutdown(&self, message: &lsp::Message) -> crate::Result<bool> { | ||
match message { | ||
lsp::Message::Request(lsp::Request { id, method, .. }) | ||
if method == lsp_types::request::Shutdown::METHOD => | ||
{ | ||
self.sender | ||
.send(lsp::Response::new_ok(id.clone(), ()).into())?; | ||
tracing::info!("Shutdown request received. Waiting for an exit notification..."); | ||
match self.receiver.recv_timeout(std::time::Duration::from_secs(30))? { | ||
lsp::Message::Notification(lsp::Notification { method, .. }) if method == lsp_types::notification::Exit::METHOD => { | ||
tracing::info!("Exit notification received. Server shutting down..."); | ||
Ok(true) | ||
}, | ||
message => anyhow::bail!("Server received unexpected message {message:?} while waiting for exit notification") | ||
} | ||
} | ||
lsp::Message::Notification(lsp::Notification { method, .. }) | ||
if method == lsp_types::notification::Exit::METHOD => | ||
{ | ||
tracing::error!("Server received an exit notification before a shutdown request was sent. Exiting..."); | ||
Ok(true) | ||
} | ||
_ => Ok(false), | ||
} | ||
} | ||
|
||
/// Joins the I/O threads that underpin this connection. | ||
pub(super) fn join_io_threads(self) -> crate::Result<()> { | ||
std::mem::drop( | ||
Arc::into_inner(self.sender) | ||
.expect("the client sender shouldn't have more than one strong reference"), | ||
); | ||
std::mem::drop(self.receiver); | ||
self.threads.join()?; | ||
Ok(()) | ||
} | ||
} | ||
|
||
/// A weak reference to an underlying sender channel for communication with the client. | ||
#[derive(Clone, Debug)] | ||
pub(crate) struct ClientSender { | ||
weak_sender: Weak<ConnectionSender>, | ||
} | ||
|
||
// note: additional wrapper functions for senders may be implemented as needed. | ||
impl ClientSender { | ||
pub(crate) fn send(&self, msg: lsp::Message) -> crate::Result<()> { | ||
let Some(sender) = self.weak_sender.upgrade() else { | ||
anyhow::bail!("The connection with the client has been closed"); | ||
}; | ||
|
||
Ok(sender.send(msg)?) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters