Merge pull request #18972 from osiewicz/drop-outgoing-messages-on-background-thread

lsp-server: Drop outgoing messages on background thread

Author: Veykril

lib/lsp-server/src/msg.rs

@@ -181,15 +181,15 @@ impl Message {
 
         Ok(Some(msg))
     }
-    pub fn write(self, w: &mut impl Write) -> io::Result<()> {
+    pub fn write(&self, w: &mut impl Write) -> io::Result<()> {
         self._write(w)
     }
-    fn _write(self, w: &mut dyn Write) -> io::Result<()> {
+    fn _write(&self, w: &mut dyn Write) -> io::Result<()> {
         #[derive(Serialize)]
-        struct JsonRpc {
+        struct JsonRpc<'a> {
             jsonrpc: &'static str,
             #[serde(flatten)]
-            msg: Message,
+            msg: &'a Message,
         }
         let text = serde_json::to_string(&JsonRpc { jsonrpc: "2.0", msg: self })?;
         write_msg_text(w, &text)

lib/lsp-server/src/socket.rs

@@ -15,8 +15,11 @@ pub(crate) fn socket_transport(
     stream: TcpStream,
 ) -> (Sender<Message>, Receiver<Message>, IoThreads) {
     let (reader_receiver, reader) = make_reader(stream.try_clone().unwrap());
-    let (writer_sender, writer) = make_write(stream);
-    let io_threads = make_io_threads(reader, writer);
+    let (writer_sender, writer, messages_to_drop) = make_write(stream);
+    let dropper = std::thread::spawn(move || {
+        messages_to_drop.into_iter().for_each(drop);
+    });
+    let io_threads = make_io_threads(reader, writer, dropper);
     (writer_sender, reader_receiver, io_threads)
 }
 
@@ -36,11 +39,21 @@ fn make_reader(stream: TcpStream) -> (Receiver<Message>, thread::JoinHandle<io::
     (reader_receiver, reader)
 }
 
-fn make_write(mut stream: TcpStream) -> (Sender<Message>, thread::JoinHandle<io::Result<()>>) {
+fn make_write(
+    mut stream: TcpStream,
+) -> (Sender<Message>, thread::JoinHandle<io::Result<()>>, Receiver<Message>) {
     let (writer_sender, writer_receiver) = bounded::<Message>(0);
+    let (drop_sender, drop_receiver) = bounded::<Message>(0);
     let writer = thread::spawn(move || {
-        writer_receiver.into_iter().try_for_each(|it| it.write(&mut stream)).unwrap();
+        writer_receiver
+            .into_iter()
+            .try_for_each(|it| {
+                let result = it.write(&mut stream);
+                let _ = drop_sender.send(it);
+                result
+            })
+            .unwrap();
         Ok(())
     });
-    (writer_sender, writer)
+    (writer_sender, writer, drop_receiver)
 }

lib/lsp-server/src/stdio.rs

@@ -11,15 +11,24 @@ use crate::Message;
 
 /// Creates an LSP connection via stdio.
 pub(crate) fn stdio_transport() -> (Sender<Message>, Receiver<Message>, IoThreads) {
+    let (drop_sender, drop_receiver) = bounded::<Message>(0);
     let (writer_sender, writer_receiver) = bounded::<Message>(0);
     let writer = thread::Builder::new()
         .name("LspServerWriter".to_owned())
         .spawn(move || {
             let stdout = stdout();
             let mut stdout = stdout.lock();
-            writer_receiver.into_iter().try_for_each(|it| it.write(&mut stdout))
+            writer_receiver.into_iter().try_for_each(|it| {
+                let result = it.write(&mut stdout);
+                let _ = drop_sender.send(it);
+                result
+            })
         })
         .unwrap();
+    let dropper = thread::Builder::new()
+        .name("LspMessageDropper".to_owned())
+        .spawn(move || drop_receiver.into_iter().for_each(drop))
+        .unwrap();
     let (reader_sender, reader_receiver) = bounded::<Message>(0);
     let reader = thread::Builder::new()
         .name("LspServerReader".to_owned())
@@ -41,21 +50,23 @@ pub(crate) fn stdio_transport() -> (Sender<Message>, Receiver<Message>, IoThread
             Ok(())
         })
         .unwrap();
-    let threads = IoThreads { reader, writer };
+    let threads = IoThreads { reader, writer, dropper };
     (writer_sender, reader_receiver, threads)
 }
 
 // Creates an IoThreads
 pub(crate) fn make_io_threads(
     reader: thread::JoinHandle<io::Result<()>>,
     writer: thread::JoinHandle<io::Result<()>>,
+    dropper: thread::JoinHandle<()>,
 ) -> IoThreads {
-    IoThreads { reader, writer }
+    IoThreads { reader, writer, dropper }
 }
 
 pub struct IoThreads {
     reader: thread::JoinHandle<io::Result<()>>,
     writer: thread::JoinHandle<io::Result<()>>,
+    dropper: thread::JoinHandle<()>,
 }
 
 impl IoThreads {
@@ -64,6 +75,12 @@ impl IoThreads {
             Ok(r) => r?,
             Err(err) => std::panic::panic_any(err),
         }
+        match self.dropper.join() {
+            Ok(_) => (),
+            Err(err) => {
+                std::panic::panic_any(err);
+            }
+        }
         match self.writer.join() {
             Ok(r) => r,
             Err(err) => {