on master help sign in
lsp: Refactor duplex to avoid issues with select! + read_exact

read_exact isn't cancellation safe. Fixes #504

commit385a6b5a1adddfc26e917982641530e1a7c7aa81 authorBlaž Hrastnik <blaz@mxxn.io> dateSat 7 Aug 2021, 15:02 (5 years ago) parent8714b71
1 file changed+81−51
helix-lsp/src/transport.rs+81−51
helix-lsp/src/transport.rs+81−51
@@ -5,10 +5,14 @@ use log::{error, info};
55use serde::{Deserialize, Serialize};
66use serde_json::Value;
77use std::collections::HashMap;
8+use std::sync::Arc;
89use tokio::{
910 io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter},
1011 process::{ChildStderr, ChildStdin, ChildStdout},
11− sync::mpsc::{unbounded_channel, Sender, UnboundedReceiver, UnboundedSender},
12+ sync::{
13+ mpsc::{unbounded_channel, Sender, UnboundedReceiver, UnboundedSender},
14+ Mutex,
15+ },
1216};
1317
1418#[derive(Debug)]
@@ -35,14 +39,7 @@ enum ServerMessage {
3539#[derive(Debug)]
3640pub struct Transport {
3741 id: usize,
38− client_tx: UnboundedSender<(usize, jsonrpc::Call)>,
39− client_rx: UnboundedReceiver<Payload>,
40−
41− pending_requests: HashMap<jsonrpc::Id, Sender<Result<Value>>>,
42−
43− server_stdin: BufWriter<ChildStdin>,
44− server_stdout: BufReader<ChildStdout>,
45− server_stderr: BufReader<ChildStderr>,
42+ pending_requests: Mutex<HashMap<jsonrpc::Id, Sender<Result<Value>>>>,
4643}
4744
4845impl Transport {
@@ -60,15 +57,14 @@ impl Transport {
6057
6158 let transport = Self {
6259 id,
63− server_stdout,
64− server_stdin,
65− server_stderr,
66− client_tx,
67− client_rx,
68− pending_requests: HashMap::default(),
60+ pending_requests: Mutex::new(HashMap::default()),
6961 };
7062
71− tokio::spawn(transport.duplex());
63+ let transport = Arc::new(transport);
64+
65+ tokio::spawn(Self::recv(transport.clone(), server_stdout, client_tx));
66+ tokio::spawn(Self::err(transport.clone(), server_stderr));
67+ tokio::spawn(Self::send(transport, server_stdin, client_rx));
7268
7369 (rx, tx)
7470 }
@@ -109,12 +105,12 @@ impl Transport {
109105 //TODO: reuse vector
110106 let mut content = vec![0; content_length];
111107 reader.read_exact(&mut content).await?;
112− let msg = String::from_utf8(content).context("invalid utf8 from server")?;
108+ let msg = std::str::from_utf8(&content).context("invalid utf8 from server")?;
113109
114110 info!("<- {}", msg);
115111
116112 // try parsing as output (server response) or call (server request)
117− let output: serde_json::Result<ServerMessage> = serde_json::from_str(&msg);
113+ let output: serde_json::Result<ServerMessage> = serde_json::from_str(msg);
118114
119115 Ok(output?)
120116 }
@@ -130,40 +126,55 @@ impl Transport {
130126 Ok(())
131127 }
132128
133− async fn send_payload_to_server(&mut self, payload: Payload) -> Result<()> {
129+ async fn send_payload_to_server(
130+ &self,
131+ server_stdin: &mut BufWriter<ChildStdin>,
132+ payload: Payload,
133+ ) -> Result<()> {
134134 //TODO: reuse string
135135 let json = match payload {
136136 Payload::Request { chan, value } => {
137− self.pending_requests.insert(value.id.clone(), chan);
137+ self.pending_requests
138+ .lock()
139+ .await
140+ .insert(value.id.clone(), chan);
138141 serde_json::to_string(&value)?
139142 }
140143 Payload::Notification(value) => serde_json::to_string(&value)?,
141144 Payload::Response(error) => serde_json::to_string(&error)?,
142145 };
143− self.send_string_to_server(json).await
146+ self.send_string_to_server(server_stdin, json).await
144147 }
145148
146− async fn send_string_to_server(&mut self, request: String) -> Result<()> {
149+ async fn send_string_to_server(
150+ &self,
151+ server_stdin: &mut BufWriter<ChildStdin>,
152+ request: String,
153+ ) -> Result<()> {
147154 info!("-> {}", request);
148155
149156 // send the headers
150− self.server_stdin
157+ server_stdin
151158 .write_all(format!("Content-Length: {}\r\n\r\n", request.len()).as_bytes())
152159 .await?;
153160
154161 // send the body
155− self.server_stdin.write_all(request.as_bytes()).await?;
162+ server_stdin.write_all(request.as_bytes()).await?;
156163
157− self.server_stdin.flush().await?;
164+ server_stdin.flush().await?;
158165
159166 Ok(())
160167 }
161168
162− async fn process_server_message(&mut self, msg: ServerMessage) -> Result<()> {
169+ async fn process_server_message(
170+ &self,
171+ client_tx: &UnboundedSender<(usize, jsonrpc::Call)>,
172+ msg: ServerMessage,
173+ ) -> Result<()> {
163174 match msg {
164175 ServerMessage::Output(output) => self.process_request_response(output).await?,
165176 ServerMessage::Call(call) => {
166− self.client_tx
177+ client_tx
167178 .send((self.id, call))
168179 .context("failed to send a message to server")?;
169180 // let notification = Notification::parse(&method, params);
@@ -172,7 +183,7 @@ impl Transport {
172183 Ok(())
173184 }
174185
175− async fn process_request_response(&mut self, output: jsonrpc::Output) -> Result<()> {
186+ async fn process_request_response(&self, output: jsonrpc::Output) -> Result<()> {
176187 let (id, result) = match output {
177188 jsonrpc::Output::Success(jsonrpc::Success { id, result, .. }) => {
178189 info!("<- {}", result);
@@ -186,6 +197,8 @@ impl Transport {
186197
187198 let tx = self
188199 .pending_requests
200+ .lock()
201+ .await
189202 .remove(&id)
190203 .expect("pending_request with id not found!");
191204
@@ -200,34 +213,51 @@ impl Transport {
200213 Ok(())
201214 }
202215
203− async fn duplex(mut self) {
216+ async fn recv(
217+ transport: Arc<Self>,
218+ mut server_stdout: BufReader<ChildStdout>,
219+ client_tx: UnboundedSender<(usize, jsonrpc::Call)>,
220+ ) {
204221 let mut recv_buffer = String::new();
205− let mut err_buffer = String::new();
206222 loop {
207− tokio::select! {
208− // client -> server
209− msg = self.client_rx.recv() => {
210− match msg {
211− Some(msg) => {
212− self.send_payload_to_server(msg).await.unwrap()
213− },
214− None => break
215− }
223+ match Self::recv_server_message(&mut server_stdout, &mut recv_buffer).await {
224+ Ok(msg) => {
225+ transport
226+ .process_server_message(&client_tx, msg)
227+ .await
228+ .unwrap();
216229 }
217− // server -> client
218− msg = Self::recv_server_message(&mut self.server_stdout, &mut recv_buffer) => {
219− match msg {
220− Ok(msg) => {
221− self.process_server_message(msg).await.unwrap();
222− }
223− Err(_) => {
224− error!("err: <- {:?}", msg);
225− break;
226− },
227− }
230+ Err(err) => {
231+ error!("err: <- {:?}", err);
232+ break;
228233 }
229− _msg = Self::recv_server_error(&mut self.server_stderr, &mut err_buffer) => {}
230234 }
231235 }
232236 }
237+
238+ async fn err(_transport: Arc<Self>, mut server_stderr: BufReader<ChildStderr>) {
239+ let mut recv_buffer = String::new();
240+ loop {
241+ match Self::recv_server_error(&mut server_stderr, &mut recv_buffer).await {
242+ Ok(_) => {}
243+ Err(err) => {
244+ error!("err: <- {:?}", err);
245+ break;
246+ }
247+ }
248+ }
249+ }
250+
251+ async fn send(
252+ transport: Arc<Self>,
253+ mut server_stdin: BufWriter<ChildStdin>,
254+ mut client_rx: UnboundedReceiver<Payload>,
255+ ) {
256+ while let Some(msg) = client_rx.recv().await {
257+ transport
258+ .send_payload_to_server(&mut server_stdin, msg)
259+ .await
260+ .unwrap()
261+ }
262+ }
233263}