intermediary

Author: juleswritescode

crates/pg_lsp/Cargo.toml

@@ -32,7 +32,7 @@ pg_base_db.workspace = true
 pg_schema_cache.workspace = true
 pg_workspace.workspace = true
 pg_diagnostics.workspace = true
-tokio = { version = "1.40.0", features = ["macros", "rt-multi-thread", "sync", "time"] }
+tokio = { version = "1.40.0", features = ["io-std", "macros", "rt-multi-thread", "sync", "time"] }
 tokio-util = "0.7.12"
 tower-lsp = "0.20.0"
 

crates/pg_lsp/src/client/client_flags.rs

@@ -6,9 +6,6 @@ use tower_lsp::lsp_types::InitializeParams;
 pub struct ClientFlags {
     /// If `true`, the server can pull configuration from the client.
     pub supports_pull_opts: bool,
-
-    /// If `true`, the client notifies the server when its configuration changes.
-    pub supports_dynamic_registration: bool,
 }
 
 impl ClientFlags {
@@ -20,17 +17,6 @@ impl ClientFlags {
             .and_then(|w| w.configuration)
             .unwrap_or(false);
 
-        let supports_dynamic_registration = params
-            .capabilities
-            .workspace
-            .as_ref()
-            .and_then(|w| w.did_change_configuration)
-            .and_then(|c| c.dynamic_registration)
-            .unwrap_or(false);
-
-        Self {
-            supports_pull_opts,
-            supports_dynamic_registration,
-        }
+        Self { supports_pull_opts }
     }
 }

crates/pg_lsp/src/db_connection.rs

@@ -2,66 +2,74 @@ use pg_schema_cache::SchemaCache;
 use sqlx::{postgres::PgListener, PgPool};
 use tokio::task::JoinHandle;
 
-#[derive(Debug)]
 pub(crate) struct DbConnection {
     pool: PgPool,
     connection_string: String,
-    schema_update_handle: Option<JoinHandle<()>>,
+    schema_update_handle: JoinHandle<()>,
+    close_tx: tokio::sync::oneshot::Sender<()>,
 }
 
 impl DbConnection {
-    pub(crate) async fn new(connection_string: String) -> Result<Self, sqlx::Error> {
-        let pool = PgPool::connect(&connection_string).await?;
-        Ok(Self {
-            pool,
-            connection_string: connection_string,
-            schema_update_handle: None,
-        })
-    }
-
-    pub(crate) fn connected_to(&self, connection_string: &str) -> bool {
-        connection_string == self.connection_string
-    }
-
-    pub(crate) async fn close(self) {
-        if self.schema_update_handle.is_some() {
-            self.schema_update_handle.unwrap().abort();
-        }
-        self.pool.close().await;
-    }
-
-    pub(crate) async fn listen_for_schema_updates<F>(
-        &mut self,
+    pub(crate) async fn new<F>(
+        connection_string: String,
         on_schema_update: F,
-    ) -> anyhow::Result<()>
+    ) -> Result<Self, sqlx::Error>
     where
         F: Fn(SchemaCache) -> () + Send + 'static,
     {
-        let mut listener = PgListener::connect_with(&self.pool).await?;
+        let pool = PgPool::connect(&connection_string).await?;
+
+        let mut listener = PgListener::connect_with(&pool).await?;
         listener.listen_all(["postgres_lsp", "pgrst"]).await?;
 
-        let pool = self.pool.clone();
+        let (close_tx, close_rx) = tokio::sync::oneshot::channel::<()>();
+
+        let cloned_pool = pool.clone();
+
+        let schema_update_handle: JoinHandle<()> = tokio::spawn(async move {
+            let mut moved_rx = close_rx;
 
-        let handle: JoinHandle<()> = tokio::spawn(async move {
             loop {
-                match listener.recv().await {
-                    Ok(not) => {
-                        if not.payload().to_string() == "reload schema" {
-                            let schema_cache = SchemaCache::load(&pool).await;
-                            on_schema_update(schema_cache);
-                        };
+                tokio::select! {
+                    res = listener.recv() => {
+                        match res {
+                            Ok(not) => {
+                                if not.payload().to_string() == "reload schema" {
+                                    let schema_cache = SchemaCache::load(&cloned_pool).await;
+                                    on_schema_update(schema_cache);
+                                };
+                            }
+                            Err(why) => {
+                                eprintln!("Error receiving notification: {:?}", why);
+                                break;
+                            }
+                        }
                     }
-                    Err(why) => {
-                        eprintln!("Error receiving notification: {:?}", why);
-                        break;
+
+                    _ = &mut moved_rx => {
+                        return;
                     }
                 }
             }
         });
 
-        self.schema_update_handle = Some(handle);
+        Ok(Self {
+            pool,
+            connection_string: connection_string,
+            schema_update_handle,
+            close_tx,
+        })
+    }
+
+    pub(crate) fn connected_to(&self, connection_string: &str) -> bool {
+        connection_string == self.connection_string
+    }
+
+    pub(crate) async fn close(self) {
+        let _ = self.close_tx.send(());
+        let _ = self.schema_update_handle.await;
 
-        Ok(())
+        self.pool.close().await;
     }
 
     pub(crate) fn get_pool(&self) -> PgPool {

crates/pg_lsp/src/debouncer.rs

@@ -1,21 +1,33 @@
-use std::{future::Future, pin::Pin};
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{atomic::AtomicBool, Arc},
+};
 
 type AsyncBlock = Pin<Box<dyn Future<Output = ()> + 'static + Send>>;
 
 pub(crate) struct SimpleTokioDebouncer {
     handle: tokio::task::JoinHandle<()>,
     tx: tokio::sync::mpsc::Sender<AsyncBlock>,
+    shutdown_flag: Arc<AtomicBool>,
 }
 
 impl SimpleTokioDebouncer {
     pub fn new(timeout: std::time::Duration) -> Self {
         let (tx, mut rx) = tokio::sync::mpsc::channel(100);
 
+        let shutdown_flag = Arc::new(AtomicBool::new(false));
+        let shutdown_flag_clone = shutdown_flag.clone();
+
         let handle = tokio::spawn(async move {
             let mut maybe_args: Option<AsyncBlock> = None;
             let mut instant = tokio::time::Instant::now() + timeout;
 
             loop {
+                if shutdown_flag_clone.load(std::sync::atomic::Ordering::Relaxed) {
+                    break;
+                }
+
                 tokio::select! {
                     // If the timeout is reached, execute and reset the last received action
                     _ = tokio::time::sleep_until(instant) => {
@@ -42,14 +54,30 @@ impl SimpleTokioDebouncer {
             }
         });
 
-        Self { handle, tx }
+        Self {
+            handle,
+            tx,
+            shutdown_flag,
+        }
     }
 
     pub async fn debounce(&self, block: AsyncBlock) {
+        if self
+            .shutdown_flag
+            .load(std::sync::atomic::Ordering::Relaxed)
+        {
+            println!(
+                "Trying to debounce tasks, but the Debouncer is in the process of shutting down."
+            );
+            return;
+        }
+
         self.tx.send(block).await.unwrap();
     }
 
-    pub async fn shutdown(self) {
-        let _ = self.handle.await.unwrap();
+    pub async fn shutdown(&self) {
+        self.shutdown_flag
+            .store(true, std::sync::atomic::Ordering::Relaxed);
+        let _ = self.handle.abort(); //  we don't care about any errors during shutdown
     }
 }

crates/pg_lsp/src/main.rs

@@ -1,9 +1,12 @@
-use pg_lsp::server::Server;
-use tower_lsp::LspService;
+use pg_lsp::server::LspServer;
+use tower_lsp::{LspService, Server};
 
 #[tokio::main]
-async fn main() -> anyhow::Result<()> {
-    let (server, client_socket) = LspService::build(|client| Server::new(client)).finish();
+async fn main() {
+    let stdin = tokio::io::stdin();
+    let stdout = tokio::io::stdout();
 
-    Ok(())
+    let (service, socket) = LspService::new(|client| LspServer::new(client));
+
+    Server::new(stdin, stdout, socket).serve(service).await;
 }

crates/pg_lsp/src/server.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use notification::ShowMessage;
 use pg_commands::CommandType;
 use tokio::sync::RwLock;
@@ -13,18 +15,18 @@ use crate::utils::file_path;
 use crate::utils::normalize_uri;
 use crate::utils::to_proto;
 
-pub struct Server {
-    client: Client,
-    session: Session,
+pub struct LspServer {
+    client: Arc<Client>,
+    session: Arc<Session>,
     client_capabilities: RwLock<Option<ClientFlags>>,
     debouncer: SimpleTokioDebouncer,
 }
 
-impl Server {
+impl LspServer {
     pub fn new(client: Client) -> Self {
         Self {
-            client,
-            session: Session::new(),
+            client: Arc::new(client),
+            session: Arc::new(Session::new()),
             client_capabilities: RwLock::new(None),
             debouncer: SimpleTokioDebouncer::new(std::time::Duration::from_millis(500)),
         }
@@ -120,8 +122,8 @@ impl Server {
     }
 
     async fn publish_diagnostics_debounced(&self, mut uri: Url) {
-        let session = self.session.clone();
-        let client = self.client.clone();
+        let client = Arc::clone(&self.client);
+        let session = Arc::clone(&self.session);
 
         self.debouncer
             .debounce(Box::pin(async move {
@@ -157,7 +159,7 @@ impl Server {
 }
 
 #[tower_lsp::async_trait]
-impl LanguageServer for Server {
+impl LanguageServer for LspServer {
     async fn initialize(&self, params: InitializeParams) -> jsonrpc::Result<InitializeResult> {
         let flags = ClientFlags::from_initialize_request_params(&params);
         self.client_capabilities.blocking_write().replace(flags);
@@ -199,11 +201,13 @@ impl LanguageServer for Server {
     }
 
     async fn shutdown(&self) -> jsonrpc::Result<()> {
-        // TODO: Shutdown stuff.
+        self.session.shutdown().await;
+        self.debouncer.shutdown().await;
 
         self.client
             .log_message(MessageType::INFO, "Postgres LSP terminated.")
             .await;
+
         Ok(())
     }
 
@@ -247,7 +251,10 @@ impl LanguageServer for Server {
                     self.client
                         .show_message(
                             MessageType::ERROR,
-                            format!("Used Client Options from params but failed to set them: {}", err),
+                            format!(
+                                "Used Client Options from params but failed to set them: {}",
+                                err
+                            ),
                         )
                         .await
                 }