lsp: Refactor duplex to avoid issues with select! + read_exact

read_exact isn't cancellation safe.

Fixes #504
pull/560/head
Blaž Hrastnik 3 years ago
parent 8714b71991
commit 385a6b5a1a

@ -5,10 +5,14 @@ use log::{error, info};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc;
use tokio::{ use tokio::{
io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter}, io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter},
process::{ChildStderr, ChildStdin, ChildStdout}, process::{ChildStderr, ChildStdin, ChildStdout},
sync::mpsc::{unbounded_channel, Sender, UnboundedReceiver, UnboundedSender}, sync::{
mpsc::{unbounded_channel, Sender, UnboundedReceiver, UnboundedSender},
Mutex,
},
}; };
#[derive(Debug)] #[derive(Debug)]
@ -35,14 +39,7 @@ enum ServerMessage {
#[derive(Debug)] #[derive(Debug)]
pub struct Transport { pub struct Transport {
id: usize, id: usize,
client_tx: UnboundedSender<(usize, jsonrpc::Call)>, pending_requests: Mutex<HashMap<jsonrpc::Id, Sender<Result<Value>>>>,
client_rx: UnboundedReceiver<Payload>,
pending_requests: HashMap<jsonrpc::Id, Sender<Result<Value>>>,
server_stdin: BufWriter<ChildStdin>,
server_stdout: BufReader<ChildStdout>,
server_stderr: BufReader<ChildStderr>,
} }
impl Transport { impl Transport {
@ -60,15 +57,14 @@ impl Transport {
let transport = Self { let transport = Self {
id, id,
server_stdout, pending_requests: Mutex::new(HashMap::default()),
server_stdin,
server_stderr,
client_tx,
client_rx,
pending_requests: HashMap::default(),
}; };
tokio::spawn(transport.duplex()); let transport = Arc::new(transport);
tokio::spawn(Self::recv(transport.clone(), server_stdout, client_tx));
tokio::spawn(Self::err(transport.clone(), server_stderr));
tokio::spawn(Self::send(transport, server_stdin, client_rx));
(rx, tx) (rx, tx)
} }
@ -109,12 +105,12 @@ impl Transport {
//TODO: reuse vector //TODO: reuse vector
let mut content = vec![0; content_length]; let mut content = vec![0; content_length];
reader.read_exact(&mut content).await?; reader.read_exact(&mut content).await?;
let msg = String::from_utf8(content).context("invalid utf8 from server")?; let msg = std::str::from_utf8(&content).context("invalid utf8 from server")?;
info!("<- {}", msg); info!("<- {}", msg);
// try parsing as output (server response) or call (server request) // try parsing as output (server response) or call (server request)
let output: serde_json::Result<ServerMessage> = serde_json::from_str(&msg); let output: serde_json::Result<ServerMessage> = serde_json::from_str(msg);
Ok(output?) Ok(output?)
} }
@ -130,40 +126,55 @@ impl Transport {
Ok(()) Ok(())
} }
async fn send_payload_to_server(&mut self, payload: Payload) -> Result<()> { async fn send_payload_to_server(
&self,
server_stdin: &mut BufWriter<ChildStdin>,
payload: Payload,
) -> Result<()> {
//TODO: reuse string //TODO: reuse string
let json = match payload { let json = match payload {
Payload::Request { chan, value } => { Payload::Request { chan, value } => {
self.pending_requests.insert(value.id.clone(), chan); self.pending_requests
.lock()
.await
.insert(value.id.clone(), chan);
serde_json::to_string(&value)? serde_json::to_string(&value)?
} }
Payload::Notification(value) => serde_json::to_string(&value)?, Payload::Notification(value) => serde_json::to_string(&value)?,
Payload::Response(error) => serde_json::to_string(&error)?, Payload::Response(error) => serde_json::to_string(&error)?,
}; };
self.send_string_to_server(json).await self.send_string_to_server(server_stdin, json).await
} }
async fn send_string_to_server(&mut self, request: String) -> Result<()> { async fn send_string_to_server(
&self,
server_stdin: &mut BufWriter<ChildStdin>,
request: String,
) -> Result<()> {
info!("-> {}", request); info!("-> {}", request);
// send the headers // send the headers
self.server_stdin server_stdin
.write_all(format!("Content-Length: {}\r\n\r\n", request.len()).as_bytes()) .write_all(format!("Content-Length: {}\r\n\r\n", request.len()).as_bytes())
.await?; .await?;
// send the body // send the body
self.server_stdin.write_all(request.as_bytes()).await?; server_stdin.write_all(request.as_bytes()).await?;
self.server_stdin.flush().await?; server_stdin.flush().await?;
Ok(()) Ok(())
} }
async fn process_server_message(&mut self, msg: ServerMessage) -> Result<()> { async fn process_server_message(
&self,
client_tx: &UnboundedSender<(usize, jsonrpc::Call)>,
msg: ServerMessage,
) -> Result<()> {
match msg { match msg {
ServerMessage::Output(output) => self.process_request_response(output).await?, ServerMessage::Output(output) => self.process_request_response(output).await?,
ServerMessage::Call(call) => { ServerMessage::Call(call) => {
self.client_tx client_tx
.send((self.id, call)) .send((self.id, call))
.context("failed to send a message to server")?; .context("failed to send a message to server")?;
// let notification = Notification::parse(&method, params); // let notification = Notification::parse(&method, params);
@ -172,7 +183,7 @@ impl Transport {
Ok(()) Ok(())
} }
async fn process_request_response(&mut self, output: jsonrpc::Output) -> Result<()> { async fn process_request_response(&self, output: jsonrpc::Output) -> Result<()> {
let (id, result) = match output { let (id, result) = match output {
jsonrpc::Output::Success(jsonrpc::Success { id, result, .. }) => { jsonrpc::Output::Success(jsonrpc::Success { id, result, .. }) => {
info!("<- {}", result); info!("<- {}", result);
@ -186,6 +197,8 @@ impl Transport {
let tx = self let tx = self
.pending_requests .pending_requests
.lock()
.await
.remove(&id) .remove(&id)
.expect("pending_request with id not found!"); .expect("pending_request with id not found!");
@ -200,34 +213,51 @@ impl Transport {
Ok(()) Ok(())
} }
async fn duplex(mut self) { async fn recv(
transport: Arc<Self>,
mut server_stdout: BufReader<ChildStdout>,
client_tx: UnboundedSender<(usize, jsonrpc::Call)>,
) {
let mut recv_buffer = String::new(); let mut recv_buffer = String::new();
let mut err_buffer = String::new();
loop { loop {
tokio::select! { match Self::recv_server_message(&mut server_stdout, &mut recv_buffer).await {
// client -> server Ok(msg) => {
msg = self.client_rx.recv() => { transport
match msg { .process_server_message(&client_tx, msg)
Some(msg) => { .await
self.send_payload_to_server(msg).await.unwrap() .unwrap();
}, }
None => break Err(err) => {
error!("err: <- {:?}", err);
break;
}
} }
} }
// server -> client
msg = Self::recv_server_message(&mut self.server_stdout, &mut recv_buffer) => {
match msg {
Ok(msg) => {
self.process_server_message(msg).await.unwrap();
} }
Err(_) => {
error!("err: <- {:?}", msg); async fn err(_transport: Arc<Self>, mut server_stderr: BufReader<ChildStderr>) {
let mut recv_buffer = String::new();
loop {
match Self::recv_server_error(&mut server_stderr, &mut recv_buffer).await {
Ok(_) => {}
Err(err) => {
error!("err: <- {:?}", err);
break; break;
},
} }
} }
_msg = Self::recv_server_error(&mut self.server_stderr, &mut err_buffer) => {}
} }
} }
async fn send(
transport: Arc<Self>,
mut server_stdin: BufWriter<ChildStdin>,
mut client_rx: UnboundedReceiver<Payload>,
) {
while let Some(msg) = client_rx.recv().await {
transport
.send_payload_to_server(&mut server_stdin, msg)
.await
.unwrap()
}
} }
} }

Loading…
Cancel
Save