write proxy: expire old write delegation streams

Right now we don't have a mechanism for expiring write delegation
streams that disconnected, but without sending a valid Disconnect
messages. That unfortunately also means that we leak libSQL connections,
which are kept in-memory indefinitely. That leaks memory as well as disk
space, because connections also keep temporary tables in place.

FIXME: we also need to keep information about expired streams for longer,
in case a client tries to send a late message. Such a message should be
rejected in order to avoid running it outside of transaction context.

Author: psarna

sqld/src/connection/libsql.rs

@@ -143,7 +143,7 @@ where
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct LibSqlConnection {
     sender: crossbeam::channel::Sender<ExecCallback>,
 }

sqld/src/connection/mod.rs

@@ -1,4 +1,4 @@
-use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::Arc;
 use tokio::time::Duration;
 
@@ -231,6 +231,14 @@ impl Drop for WaitersGuard<'_> {
     }
 }
 
+fn now_millis() -> u64 {
+    use std::time::{SystemTime, UNIX_EPOCH};
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_millis() as u64
+}
+
 #[async_trait::async_trait]
 impl<F: MakeConnection> MakeConnection for MakeThrottledConnection<F> {
     type Connection = TrackedConnection<F::Connection>;
@@ -266,14 +274,28 @@ impl<F: MakeConnection> MakeConnection for MakeThrottledConnection<F> {
         }
 
         let inner = self.connection_maker.create().await?;
-        Ok(TrackedConnection { permit, inner })
+        Ok(TrackedConnection {
+            permit,
+            inner,
+            atime: AtomicU64::new(now_millis()),
+        })
     }
 }
 
+#[derive(Debug)]
 pub struct TrackedConnection<DB> {
     inner: DB,
     #[allow(dead_code)] // just hold on to it
     permit: tokio::sync::OwnedSemaphorePermit,
+    atime: AtomicU64,
+}
+
+impl<DB> TrackedConnection<DB> {
+    pub fn idle_time(&self) -> Duration {
+        let now = now_millis();
+        let atime = self.atime.load(Ordering::Relaxed);
+        Duration::from_millis(now - atime)
+    }
 }
 
 #[async_trait::async_trait]
@@ -286,6 +308,7 @@ impl<DB: Connection> Connection for TrackedConnection<DB> {
         builder: B,
         replication_index: Option<FrameNo>,
     ) -> crate::Result<(B, State)> {
+        self.atime.store(now_millis(), Ordering::Relaxed);
         self.inner
             .execute_program(pgm, auth, builder, replication_index)
             .await
@@ -298,6 +321,7 @@ impl<DB: Connection> Connection for TrackedConnection<DB> {
         auth: Authenticated,
         replication_index: Option<FrameNo>,
     ) -> crate::Result<DescribeResult> {
+        self.atime.store(now_millis(), Ordering::Relaxed);
         self.inner.describe(sql, auth, replication_index).await
     }
 
@@ -308,6 +332,7 @@ impl<DB: Connection> Connection for TrackedConnection<DB> {
 
     #[inline]
     async fn checkpoint(&self) -> Result<()> {
+        self.atime.store(now_millis(), Ordering::Relaxed);
         self.inner.checkpoint().await
     }
 }

sqld/src/lib.rs

@@ -480,7 +480,20 @@ where
         }
 
         if let Some(config) = self.rpc_config.take() {
+            let proxy_service =
+                ProxyService::new(namespaces.clone(), None, self.disable_namespaces);
+            // Garbage collect proxy clients every 30 seconds
+            self.join_set.spawn({
+                let clients = proxy_service.clients();
+                async move {
+                    loop {
+                        tokio::time::sleep(Duration::from_secs(30)).await;
+                        rpc::proxy::garbage_collect(&mut *clients.write().await).await;
+                    }
+                }
+            });
             self.join_set.spawn(run_rpc_server(
+                proxy_service,
                 config.acceptor,
                 config.tls_config,
                 self.idle_shutdown_kicker.clone(),
@@ -498,7 +511,16 @@ where
 
         let proxy_service =
             ProxyService::new(namespaces.clone(), Some(self.auth), self.disable_namespaces);
-
+        // Garbage collect proxy clients every 30 seconds
+        self.join_set.spawn({
+            let clients = proxy_service.clients();
+            async move {
+                loop {
+                    tokio::time::sleep(Duration::from_secs(30)).await;
+                    rpc::proxy::garbage_collect(&mut *clients.write().await).await;
+                }
+            }
+        });
         Ok((namespaces, proxy_service, logger_service))
     }
 }

sqld/src/rpc/mod.rs

@@ -24,15 +24,14 @@ pub mod replication_log_proxy;
 pub const NAMESPACE_DOESNT_EXIST: &str = "NAMESPACE_DOESNT_EXIST";
 pub(crate) const NAMESPACE_METADATA_KEY: &str = "x-namespace-bin";
 
-#[allow(clippy::too_many_arguments)]
 pub async fn run_rpc_server<A: crate::net::Accept>(
+    proxy_service: ProxyService,
     acceptor: A,
     maybe_tls: Option<TlsConfig>,
     idle_shutdown_layer: Option<IdleShutdownKicker>,
     namespaces: NamespaceStore<PrimaryNamespaceMaker>,
     disable_namespaces: bool,
 ) -> anyhow::Result<()> {
-    let proxy_service = ProxyService::new(namespaces.clone(), None, disable_namespaces);
     let logger_service = ReplicationLogService::new(
         namespaces.clone(),
         idle_shutdown_layer.clone(),

sqld/src/rpc/proxy.rs

@@ -265,7 +265,7 @@ pub mod rpc {
 }
 
 pub struct ProxyService {
-    clients: RwLock<HashMap<Uuid, Arc<TrackedConnection<LibSqlConnection>>>>,
+    clients: Arc<RwLock<HashMap<Uuid, Arc<TrackedConnection<LibSqlConnection>>>>>,
     namespaces: NamespaceStore<PrimaryNamespaceMaker>,
     auth: Option<Arc<Auth>>,
     disable_namespaces: bool,
@@ -277,13 +277,19 @@ impl ProxyService {
         auth: Option<Arc<Auth>>,
         disable_namespaces: bool,
     ) -> Self {
+        let clients: Arc<RwLock<HashMap<Uuid, Arc<TrackedConnection<LibSqlConnection>>>>> =
+            Default::default();
         Self {
-            clients: Default::default(),
+            clients,
             namespaces,
             auth,
             disable_namespaces,
         }
     }
+
+    pub fn clients(&self) -> Arc<RwLock<HashMap<Uuid, Arc<TrackedConnection<LibSqlConnection>>>>> {
+        self.clients.clone()
+    }
 }
 
 #[derive(Debug, Default)]
@@ -441,6 +447,28 @@ impl QueryResultBuilder for ExecuteResultBuilder {
     }
 }
 
+// Disconnects all clients that have been idle for more than 30 seconds.
+// FIXME: we should also keep a list of recently disconnected clients,
+// and if one should arrive with a late message, it should be rejected
+// with an error. A similar mechanism is already implemented in hrana-over-http.
+pub async fn garbage_collect(
+    clients: &mut HashMap<Uuid, Arc<TrackedConnection<LibSqlConnection>>>,
+) {
+    let mut to_remove = Vec::new();
+    let limit = std::time::Duration::from_secs(30);
+    for (client_id, db) in clients.iter() {
+        if db.idle_time() > limit {
+            to_remove.push(*client_id);
+        }
+    }
+
+    tracing::trace!("garbage collecting clients: {to_remove:?}");
+    for client_id in to_remove {
+        clients.remove(&client_id);
+    }
+    tracing::trace!("remaining client handles: {:?}", clients);
+}
+
 #[tonic::async_trait]
 impl Proxy for ProxyService {
     async fn execute(