sync_ls/server/
lsp_srv.rs

1use lsp_types::{notification::Notification as Notif, request::Request as Req, *};
2
3use super::*;
4
5type PureHandler<S, T> = fn(srv: &mut S, args: T) -> LspResult<()>;
6
7impl<S: 'static> TypedLspClient<S> {
8    /// Sends a request to the client and registers a handler handled by the
9    /// service `S`.
10    pub fn send_lsp_request<R: Req>(
11        &self,
12        params: R::Params,
13        handler: impl FnOnce(&mut S, lsp::Response) + Send + Sync + 'static,
14    ) {
15        let caster = self.caster.clone();
16        self.client
17            .send_lsp_request_::<R>(params, move |s, resp| handler(caster(s), resp))
18    }
19}
20
21impl LspClient {
22    /// Sends a request to the client and registers a handler.
23    pub fn send_lsp_request_<R: Req>(
24        &self,
25        params: R::Params,
26        handler: impl FnOnce(&mut dyn Any, lsp::Response) + Send + Sync + 'static,
27    ) {
28        let mut req_queue = self.req_queue.lock();
29        let request = req_queue.outgoing.register(
30            R::METHOD.to_owned(),
31            params,
32            Box::new(|s, resp| handler(s, resp.try_into().unwrap())),
33        );
34
35        self.sender.send_message(request.into());
36    }
37
38    /// Completes an client2server request in the request queue.
39    pub fn respond_lsp(&self, response: lsp::Response) {
40        self.respond(response.id.clone(), response.into())
41    }
42
43    /// Sends a typed notification to the client.
44    pub fn send_notification<N: Notif>(&self, params: &N::Params) {
45        self.send_notification_(lsp::Notification::new(N::METHOD.to_owned(), params));
46    }
47
48    /// Sends an untyped notification to the client.
49    pub fn send_notification_(&self, notif: lsp::Notification) {
50        self.sender.send_message(notif.into());
51    }
52}
53
54impl<Args: Initializer> LsBuilder<LspMessage, Args>
55where
56    Args::S: 'static,
57{
58    /// Registers an raw event handler.
59    pub fn with_command_(
60        mut self,
61        cmd: &'static str,
62        handler: RawHandler<Args::S, Vec<JsonValue>>,
63    ) -> Self {
64        self.command_handlers
65            .insert(cmd, Box::new(move |s, _req_id, args| handler(s, args)));
66        self
67    }
68
69    /// Registers an async command handler.
70    pub fn with_command<R: Serialize + 'static>(
71        mut self,
72        cmd: &'static str,
73        handler: AsyncHandler<Args::S, Vec<JsonValue>, R>,
74    ) -> Self {
75        self.command_handlers.insert(
76            cmd,
77            Box::new(move |s, _req_id, req| erased_response(handler(s, req))),
78        );
79        self
80    }
81
82    /// Registers an untyped notification handler.
83    pub fn with_notification_<R: Notif>(
84        mut self,
85        handler: PureHandler<Args::S, JsonValue>,
86    ) -> Self {
87        self.notif_handlers.insert(R::METHOD, Box::new(handler));
88        self
89    }
90
91    /// Registers a typed notification handler.
92    pub fn with_notification<R: Notif>(mut self, handler: PureHandler<Args::S, R::Params>) -> Self {
93        self.notif_handlers.insert(
94            R::METHOD,
95            Box::new(move |s, req| handler(s, from_json(req)?)),
96        );
97        self
98    }
99
100    /// Registers a raw request handler that handlers a kind of untyped lsp
101    /// request.
102    pub fn with_raw_request<R: Req>(mut self, handler: RawHandler<Args::S, JsonValue>) -> Self {
103        self.req_handlers
104            .insert(R::METHOD, Box::new(move |s, _req_id, req| handler(s, req)));
105        self
106    }
107
108    // todo: unsafe typed
109    /// Registers an raw request handler that handlers a kind of typed lsp
110    /// request.
111    pub fn with_request_<R: Req>(
112        mut self,
113        handler: fn(&mut Args::S, R::Params) -> ScheduleResult,
114    ) -> Self {
115        self.req_handlers.insert(
116            R::METHOD,
117            Box::new(move |s, _req_id, req| handler(s, from_json(req)?)),
118        );
119        self
120    }
121
122    /// Registers a typed request handler.
123    pub fn with_request<R: Req>(
124        mut self,
125        handler: AsyncHandler<Args::S, R::Params, R::Result>,
126    ) -> Self {
127        self.req_handlers.insert(
128            R::METHOD,
129            Box::new(move |s, _req_id, req| erased_response(handler(s, from_json(req)?))),
130        );
131        self
132    }
133}
134
135impl<Args: Initializer> LsDriver<LspMessage, Args>
136where
137    Args::S: 'static,
138{
139    /// Starts the language server on the given connection.
140    ///
141    /// If `is_replay` is true, the server will wait for all pending requests to
142    /// finish before exiting. This is useful for testing the language server.
143    ///
144    /// See [`transport::MirrorArgs`] for information about the record-replay
145    /// feature.
146    #[cfg(feature = "system")]
147    pub fn start(
148        &mut self,
149        inbox: TConnectionRx<LspMessage>,
150        is_replay: bool,
151    ) -> anyhow::Result<()> {
152        let res = self.start_(inbox);
153
154        if is_replay {
155            let client = self.client.clone();
156            let _ = std::thread::spawn(move || {
157                let since = tinymist_std::time::Instant::now();
158                let timeout = std::env::var("REPLAY_TIMEOUT")
159                    .ok()
160                    .and_then(|s| s.parse().ok())
161                    .unwrap_or(60);
162                client.handle.block_on(async {
163                    while client.has_pending_requests() {
164                        if since.elapsed().as_secs() > timeout {
165                            log::error!("replay timeout reached, {timeout}s");
166                            client.begin_panic();
167                        }
168
169                        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
170                    }
171                })
172            })
173            .join();
174        }
175
176        res
177    }
178
179    /// Starts the language server on the given connection.
180    #[cfg(feature = "system")]
181    pub fn start_(&mut self, inbox: TConnectionRx<LspMessage>) -> anyhow::Result<()> {
182        use EventOrMessage::*;
183        // todo: follow what rust analyzer does
184        // Windows scheduler implements priority boosts: if thread waits for an
185        // event (like a condvar), and event fires, priority of the thread is
186        // temporary bumped. This optimization backfires in our case: each time
187        // the `main_loop` schedules a task to run on a threadpool, the
188        // worker threads gets a higher priority, and (on a machine with
189        // fewer cores) displaces the main loop! We work around this by
190        // marking the main loop as a higher-priority thread.
191        //
192        // https://docs.microsoft.com/en-us/windows/win32/procthread/scheduling-priorities
193        // https://docs.microsoft.com/en-us/windows/win32/procthread/priority-boosts
194        // https://github.com/rust-lang/rust-analyzer/issues/2835
195        // #[cfg(windows)]
196        // unsafe {
197        //     use winapi::um::processthreadsapi::*;
198        //     let thread = GetCurrentThread();
199        //     let thread_priority_above_normal = 1;
200        //     SetThreadPriority(thread, thread_priority_above_normal);
201        // }
202
203        while let Ok(msg) = inbox.recv() {
204            const EXIT_METHOD: &str = notification::Exit::METHOD;
205            let loop_start = tinymist_std::time::now();
206            match msg {
207                Evt(event) => {
208                    let Some(event_handler) = self.events.get(&event.as_ref().type_id()) else {
209                        log::warn!("unhandled event: {:?}", event.as_ref().type_id());
210                        continue;
211                    };
212
213                    let s = match &mut self.state {
214                        State::Uninitialized(u) => ServiceState::Uninitialized(u.as_deref_mut()),
215                        State::Initializing(s) | State::Ready(s) => ServiceState::Ready(s),
216                        State::ShuttingDown => {
217                            log::warn!("server is shutting down");
218                            continue;
219                        }
220                    };
221
222                    event_handler(s, &self.client, event)?;
223                }
224                Msg(LspMessage::Request(req)) => {
225                    let client = self.client.clone();
226                    let req_id = req.id.clone();
227                    client.register_request(&req.method, &req_id, loop_start);
228                    let fut = client.schedule_tail(
229                        req_id.clone(),
230                        self.on_lsp_request(&req.method, req_id, req.params),
231                    );
232                    self.client.handle.spawn(fut);
233                }
234                Msg(LspMessage::Notification(not)) => {
235                    let is_exit = not.method == EXIT_METHOD;
236                    let track_id = self
237                        .next_not_id
238                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
239                    self.client.hook.start_notification(track_id, &not.method);
240                    let result = self.on_notification(&not.method, not.params);
241                    self.client
242                        .hook
243                        .stop_notification(track_id, &not.method, loop_start, result);
244                    if is_exit {
245                        return Ok(());
246                    }
247                }
248                Msg(LspMessage::Response(resp)) => {
249                    let s = match &mut self.state {
250                        State::Ready(s) => s,
251                        _ => {
252                            log::warn!("server is not ready yet");
253                            continue;
254                        }
255                    };
256
257                    self.client.clone().complete_lsp_request(s, resp)
258                }
259            }
260        }
261
262        log::warn!("client exited without proper shutdown sequence");
263        Ok(())
264    }
265
266    /// Handles an incoming server event.
267    #[cfg(feature = "web")]
268    pub fn on_server_event(&mut self, event_id: u32) {
269        let evt = match &self.client.sender {
270            TransportHost::Js { events, .. } => events.lock().remove(&event_id),
271            TransportHost::System(_) => {
272                panic!("cannot send server event in system transport");
273            }
274        };
275
276        if let Some(event) = evt {
277            let Some(event_handler) = self.events.get(&event.as_ref().type_id()) else {
278                log::warn!("unhandled event: {:?}", event.as_ref().type_id());
279                return;
280            };
281
282            let s = match &mut self.state {
283                State::Uninitialized(u) => ServiceState::Uninitialized(u.as_deref_mut()),
284                State::Initializing(s) | State::Ready(s) => ServiceState::Ready(s),
285                State::ShuttingDown => {
286                    log::warn!("server is shutting down");
287                    return;
288                }
289            };
290
291            let res = event_handler(s, &self.client, event);
292            if let Err(err) = res {
293                log::error!("failed to handle server event {event_id}: {err}");
294            }
295        }
296    }
297
298    /// Registers and handles a request. This should only be called once per
299    /// incoming request.
300    pub fn on_lsp_request(
301        &mut self,
302        method: &str,
303        req_id: RequestId,
304        params: JsonValue,
305    ) -> ScheduleResult {
306        match (&mut self.state, method) {
307            (State::Uninitialized(args), request::Initialize::METHOD) => {
308                // todo: what will happen if the request cannot be deserialized?
309                let params = serde_json::from_value::<Args::I>(params);
310                match params {
311                    Ok(params) => {
312                        let args = args.take().expect("already initialized");
313                        let (s, res) = args.initialize(params);
314                        self.state = State::Initializing(s);
315                        res
316                    }
317                    Err(e) => just_result(Err(invalid_request(e))),
318                }
319            }
320            (State::Uninitialized(..) | State::Initializing(..), _) => {
321                just_result(Err(not_initialized()))
322            }
323            (_, request::Initialize::METHOD) => {
324                just_result(Err(invalid_request("server is already initialized")))
325            }
326            // todo: generalize this
327            (State::Ready(..), request::ExecuteCommand::METHOD) => {
328                self.on_execute_command(req_id, params)
329            }
330            (State::Ready(s), method) => 'serve_req: {
331                let is_shutdown = method == request::Shutdown::METHOD;
332
333                let Some(handler) = self.requests.get(method) else {
334                    log::warn!("unhandled lsp request: {method}");
335                    break 'serve_req just_result(Err(method_not_found()));
336                };
337
338                let resp = handler(s, req_id, params);
339
340                if is_shutdown {
341                    self.state = State::ShuttingDown;
342                }
343
344                resp
345            }
346            (State::ShuttingDown, _) => {
347                just_result(Err(invalid_request("server is shutting down")))
348            }
349        }
350    }
351
352    /// The entry point for the `workspace/executeCommand` request.
353    fn on_execute_command(&mut self, req_id: RequestId, params: JsonValue) -> ScheduleResult {
354        let s = self.state.opt_mut().ok_or_else(not_initialized)?;
355
356        let params = from_value::<ExecuteCommandParams>(params)
357            .map_err(|e| invalid_params(e.to_string()))?;
358
359        let ExecuteCommandParams {
360            command, arguments, ..
361        } = params;
362
363        // todo: generalize this
364        if command == "tinymist.getResources" {
365            self.get_resources(req_id, arguments)
366        } else {
367            let Some(handler) = self.commands.get(command.as_str()) else {
368                log::error!("asked to execute unknown command: {command}");
369                return Err(method_not_found());
370            };
371            handler(s, req_id, arguments)
372        }
373    }
374
375    /// Handles an incoming notification.
376    pub fn on_notification(&mut self, method: &str, params: JsonValue) -> LspResult<()> {
377        let handle = |s, method: &str, params: JsonValue| {
378            let Some(handler) = self.notifications.get(method) else {
379                log::warn!("unhandled notification: {method}");
380                return Ok(());
381            };
382
383            handler(s, params)
384        };
385
386        match (&mut self.state, method) {
387            (state, notification::Initialized::METHOD) => {
388                let mut s = State::ShuttingDown;
389                std::mem::swap(state, &mut s);
390                match s {
391                    State::Initializing(s) => {
392                        *state = State::Ready(s);
393                    }
394                    _ => {
395                        std::mem::swap(state, &mut s);
396                    }
397                }
398
399                let s = match state {
400                    State::Ready(s) => s,
401                    _ => {
402                        log::warn!("server is not ready yet");
403                        return Ok(());
404                    }
405                };
406                handle(s, method, params)
407            }
408            (State::Ready(state), method) => handle(state, method, params),
409            // todo: whether it is safe to ignore notifications
410            (State::Uninitialized(..) | State::Initializing(..), method) => {
411                log::warn!("server is not ready yet, while received notification {method}");
412                Ok(())
413            }
414            (State::ShuttingDown, method) => {
415                log::warn!("server is shutting down, while received notification {method}");
416                Ok(())
417            }
418        }
419    }
420
421    /// Handles an incoming response.
422    pub fn on_lsp_response(&mut self, resp: lsp::Response) {
423        let client = self.client.clone();
424        let Some(s) = self.state_mut() else {
425            log::warn!("server is not ready yet, while received response");
426            return;
427        };
428
429        client.complete_lsp_request(s, resp)
430    }
431}