1use std::future::{ready, Future, Ready};
12use std::ops::ControlFlow;
13use std::pin::Pin;
14use std::task::{Context, Poll};
15
16use futures::future::Either;
17use lsp_types::notification::{self, Notification};
18use lsp_types::request::{self, Request};
19use pin_project_lite::pin_project;
20use tower_layer::Layer;
21use tower_service::Service;
22
23use crate::{
24 AnyEvent, AnyNotification, AnyRequest, Error, ErrorCode, LspService, ResponseError, Result,
25};
26
27#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
28enum State {
29 #[default]
30 Uninitialized,
31 Initializing,
32 Ready,
33 ShuttingDown,
34}
35
36#[derive(Debug, Default)]
40pub struct Lifecycle<S> {
41 service: S,
42 state: State,
43}
44
45define_getters!(impl[S] Lifecycle<S>, service: S);
46
47impl<S> Lifecycle<S> {
48 #[must_use]
50 pub fn new(service: S) -> Self {
51 Self {
52 service,
53 state: State::Uninitialized,
54 }
55 }
56}
57
58impl<S: LspService> Service<AnyRequest> for Lifecycle<S>
59where
60 S::Error: From<ResponseError>,
61{
62 type Response = S::Response;
63 type Error = S::Error;
64 type Future = ResponseFuture<S::Future>;
65
66 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
67 self.service.poll_ready(cx)
68 }
69
70 fn call(&mut self, req: AnyRequest) -> Self::Future {
71 let inner = match (self.state, &*req.method) {
72 (State::Uninitialized, request::Initialize::METHOD) => {
73 self.state = State::Initializing;
74 Either::Left(self.service.call(req))
75 }
76 (State::Uninitialized | State::Initializing, _) => {
77 Either::Right(ready(Err(ResponseError {
78 code: ErrorCode::SERVER_NOT_INITIALIZED,
79 message: "Server is not initialized yet".into(),
80 data: None,
81 }
82 .into())))
83 }
84 (_, request::Initialize::METHOD) => Either::Right(ready(Err(ResponseError {
85 code: ErrorCode::INVALID_REQUEST,
86 message: "Server is already initialized".into(),
87 data: None,
88 }
89 .into()))),
90 (State::Ready, _) => {
91 if req.method == request::Shutdown::METHOD {
92 self.state = State::ShuttingDown;
93 }
94 Either::Left(self.service.call(req))
95 }
96 (State::ShuttingDown, _) => Either::Right(ready(Err(ResponseError {
97 code: ErrorCode::INVALID_REQUEST,
98 message: "Server is shutting down".into(),
99 data: None,
100 }
101 .into()))),
102 };
103 ResponseFuture { inner }
104 }
105}
106
107impl<S: LspService> LspService for Lifecycle<S>
108where
109 S::Error: From<ResponseError>,
110{
111 fn notify(&mut self, notif: AnyNotification) -> ControlFlow<Result<()>> {
112 match &*notif.method {
113 notification::Exit::METHOD => {
114 self.service.notify(notif)?;
115 ControlFlow::Break(Ok(()))
116 }
117 notification::Initialized::METHOD => {
118 if self.state != State::Initializing {
119 return ControlFlow::Break(Err(Error::Protocol(format!(
120 "Unexpected initialized notification on state {:?}",
121 self.state
122 ))));
123 }
124 self.state = State::Ready;
125 self.service.notify(notif)?;
126 ControlFlow::Continue(())
127 }
128 _ => self.service.notify(notif),
129 }
130 }
131
132 fn emit(&mut self, event: AnyEvent) -> ControlFlow<Result<()>> {
133 self.service.emit(event)
134 }
135}
136
137pin_project! {
138 pub struct ResponseFuture<Fut: Future> {
140 #[pin]
141 inner: Either<Fut, Ready<Fut::Output>>,
142 }
143}
144
145impl<Fut: Future> Future for ResponseFuture<Fut> {
146 type Output = Fut::Output;
147
148 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
149 self.project().inner.poll(cx)
150 }
151}
152
153#[must_use]
155#[derive(Clone, Default)]
156pub struct LifecycleLayer {
157 _private: (),
158}
159
160impl<S> Layer<S> for LifecycleLayer {
161 type Service = Lifecycle<S>;
162
163 fn layer(&self, inner: S) -> Self::Service {
164 Lifecycle::new(inner)
165 }
166}