libsql-server: use LRU cache to store active namespaces

bottomless/src/replicator.rs

@@ -376,6 +376,33 @@ impl Replicator {
         })
     }
 
+    /// Checks if there exists any backup of given database
+    pub async fn has_backup_of(db_name: impl AsRef<str>, options: &Options) -> Result<bool> {
+        let prefix = match &options.db_id {
+            Some(db_id) => format!("{db_id}-"),
+            None => format!("ns-:{}-", db_name.as_ref()),
+        };
+        let config = options.client_config().await?;
+        let client = Client::from_conf(config);
+        let bucket = options.bucket_name.clone();
+
+        match client.head_bucket().bucket(&bucket).send().await {
+            Ok(_) => tracing::trace!("Bucket {bucket} exists and is accessible"),
+            Err(e) => {
+                tracing::trace!("Bucket checking error: {e}");
+                return Err(e.into());
+            }
+        }
+
+        let mut last_frame = 0;
+        let list_objects = client.list_objects().bucket(&bucket).prefix(&prefix);
+        let response = list_objects.send().await?;
+        let _ = Self::try_get_last_frame_no(response, &mut last_frame);
+        tracing::trace!("Last frame of {prefix}: {last_frame}");
+
+        Ok(last_frame > 0)
+    }
+
     pub async fn shutdown_gracefully(&mut self) -> Result<()> {
         tracing::info!("bottomless replicator: shutting down...");
         // 1. wait for all committed WAL frames to be committed locally

libsql-server/Cargo.toml

@@ -41,6 +41,7 @@ metrics = "0.21.1"
 metrics-util = "0.15"
 metrics-exporter-prometheus = "0.12.2"
 mimalloc = { version = "0.1.36", default-features = false }
+moka = { version = "0.12.1", features = ["future"] }
 nix = { version = "0.26.2", features = ["fs"] }
 once_cell = "1.17.0"
 parking_lot = "0.12.1"

libsql-server/src/error.rs

@@ -96,6 +96,9 @@ pub enum Error {
     NamespaceStoreShutdown,
     #[error("Unable to update metastore: {0}")]
     MetaStoreUpdateFailure(Box<dyn std::error::Error + Send + Sync>),
+    // This is for errors returned by moka
+    #[error(transparent)]
+    Ref(#[from] std::sync::Arc<Self>),
 }
 
 trait ResponseError: std::error::Error {
@@ -109,6 +112,12 @@ trait ResponseError: std::error::Error {
 impl ResponseError for Error {}
 
 impl IntoResponse for Error {
+    fn into_response(self) -> axum::response::Response {
+        (&self).into_response()
+    }
+}
+
+impl IntoResponse for &Error {
     fn into_response(self) -> axum::response::Response {
         use Error::*;
 
@@ -156,6 +165,7 @@ impl IntoResponse for Error {
             UrlParseError(_) => self.format_err(StatusCode::BAD_REQUEST),
             NamespaceStoreShutdown => self.format_err(StatusCode::SERVICE_UNAVAILABLE),
             MetaStoreUpdateFailure(_) => self.format_err(StatusCode::INTERNAL_SERVER_ERROR),
+            Ref(this) => this.as_ref().into_response(),
         }
     }
 }
@@ -230,7 +240,7 @@ pub enum LoadDumpError {
 
 impl ResponseError for LoadDumpError {}
 
-impl IntoResponse for LoadDumpError {
+impl IntoResponse for &LoadDumpError {
     fn into_response(self) -> axum::response::Response {
         use LoadDumpError::*;
 
@@ -250,7 +260,7 @@ impl IntoResponse for LoadDumpError {
 
 impl ResponseError for ForkError {}
 
-impl IntoResponse for ForkError {
+impl IntoResponse for &ForkError {
     fn into_response(self) -> axum::response::Response {
         match self {
             ForkError::Internal(_)

libsql-server/src/lib.rs

@@ -100,6 +100,7 @@ pub struct Server<C = HttpConnector, A = AddrIncoming, D = HttpsConnector<HttpCo
     pub heartbeat_config: Option<HeartbeatConfig>,
     pub disable_namespaces: bool,
     pub shutdown: Arc<Notify>,
+    pub max_active_namespaces: usize,
 }
 
 impl<C, A, D> Default for Server<C, A, D> {
@@ -117,6 +118,7 @@ impl<C, A, D> Default for Server<C, A, D> {
             heartbeat_config: Default::default(),
             disable_namespaces: true,
             shutdown: Default::default(),
+            max_active_namespaces: 100,
         }
     }
 }
@@ -384,6 +386,7 @@ where
                     db_config: self.db_config.clone(),
                     base_path: self.path.clone(),
                     auth: auth.clone(),
+                    max_active_namespaces: self.max_active_namespaces,
                 };
                 let (namespaces, proxy_service, replication_service) = replica.configure().await?;
                 self.rpc_client_config = None;
@@ -422,6 +425,7 @@ where
                     extensions,
                     base_path: self.path.clone(),
                     disable_namespaces: self.disable_namespaces,
+                    max_active_namespaces: self.max_active_namespaces,
                     join_set: &mut join_set,
                     auth: auth.clone(),
                 };
@@ -487,6 +491,7 @@ struct Primary<'a, A> {
     extensions: Arc<[PathBuf]>,
     base_path: Arc<Path>,
     disable_namespaces: bool,
+    max_active_namespaces: usize,
     auth: Arc<Auth>,
     join_set: &'a mut JoinSet<anyhow::Result<()>>,
 }
@@ -520,12 +525,12 @@ where
         let meta_store_path = conf.base_path.join("metastore");
 
         let factory = PrimaryNamespaceMaker::new(conf);
-
         let namespaces = NamespaceStore::new(
             factory,
             false,
             self.db_config.snapshot_at_shutdown,
             meta_store_path,
+            self.max_active_namespaces,
         )
         .await?;
 
@@ -602,6 +607,7 @@ struct Replica<C> {
     db_config: DbConfig,
     base_path: Arc<Path>,
     auth: Arc<Auth>,
+    max_active_namespaces: usize,
 }
 
 impl<C: Connector> Replica<C> {
@@ -627,7 +633,14 @@ impl<C: Connector> Replica<C> {
         let meta_store_path = conf.base_path.join("metastore");
 
         let factory = ReplicaNamespaceMaker::new(conf);
-        let namespaces = NamespaceStore::new(factory, true, false, meta_store_path).await?;
+        let namespaces = NamespaceStore::new(
+            factory,
+            true,
+            false,
+            meta_store_path,
+            self.max_active_namespaces,
+        )
+        .await?;
         let replication_service = ReplicationLogProxyService::new(channel.clone(), uri.clone());
         let proxy_service = ReplicaProxyService::new(channel, uri, self.auth.clone());
 

libsql-server/src/main.rs

@@ -195,6 +195,10 @@ struct Cli {
     /// Enable snapshot at shutdown
     #[clap(long)]
     snapshot_at_shutdown: bool,
+
+    /// Max active namespaces kept in-memory
+    #[clap(long, env = "SQLD_MAX_ACTIVE_NAMESPACES", default_value = "100")]
+    max_active_namespaces: usize,
 }
 
 #[derive(clap::Subcommand, Debug)]
@@ -506,6 +510,7 @@ async fn build_server(config: &Cli) -> anyhow::Result<Server> {
         disable_default_namespace: config.disable_default_namespace,
         disable_namespaces: !config.enable_namespaces,
         shutdown,
+        max_active_namespaces: config.max_active_namespaces,
     })
 }
 

libsql-server/src/namespace/fork.rs

@@ -109,7 +109,6 @@ impl ForkTask<'_> {
                 self.dest_namespace.clone(),
                 RestoreOption::Latest,
                 self.bottomless_db_id,
-                true,
                 // Forking works only on primary and
                 // PrimaryNamespaceMaker::create ignores
                 // reset_cb param

libsql-server/src/namespace/meta_store.rs

@@ -177,6 +177,13 @@ impl MetaStore {
             inner: HandleState::External(change_tx, rx),
         }
     }
+
+    // TODO: we need to either make sure that the metastore is restored
+    // before we start accepting connections or we need to contact bottomless
+    // here to check if a namespace exists. Preferably the former.
+    pub fn exists(&self, namespace: &NamespaceName) -> bool {
+        self.inner.lock().configs.contains_key(namespace)
+    }
 }
 
 impl MetaStoreHandle {