Make snapshot at shutdown opt in

Author: avinassh

libsql-server/src/config.rs

@@ -124,6 +124,7 @@ pub struct DbConfig {
     pub max_total_response_size: u64,
     pub snapshot_exec: Option<String>,
     pub checkpoint_interval: Option<Duration>,
+    pub snapshot_at_shutdown: bool,
 }
 
 impl Default for DbConfig {
@@ -139,6 +140,7 @@ impl Default for DbConfig {
             max_total_response_size: bytesize::mb(10u64),
             snapshot_exec: None,
             checkpoint_interval: None,
+            snapshot_at_shutdown: false,
         }
     }
 }

libsql-server/src/lib.rs

@@ -526,7 +526,7 @@ where
             disable_namespace: self.disable_namespaces,
         };
         let factory = PrimaryNamespaceMaker::new(conf);
-        let namespaces = NamespaceStore::new(factory, false);
+        let namespaces = NamespaceStore::new(factory, false, self.db_config.snapshot_at_shutdown);
 
         // eagerly load the default namespace when namespaces are disabled
         if self.disable_namespaces {
@@ -614,7 +614,7 @@ impl<C: Connector> Replica<C> {
             max_total_response_size: self.db_config.max_total_response_size,
         };
         let factory = ReplicaNamespaceMaker::new(conf);
-        let namespaces = NamespaceStore::new(factory, true);
+        let namespaces = NamespaceStore::new(factory, true, false);
         let replication_service = ReplicationLogProxyService::new(channel.clone(), uri.clone());
         let proxy_service = ReplicaProxyService::new(channel, uri, self.auth.clone());
 

libsql-server/src/main.rs

@@ -191,6 +191,10 @@ struct Cli {
     /// the default namespace.
     #[clap(long)]
     enable_namespaces: bool,
+
+    /// Enable snapshot at shutdown
+    #[clap(long)]
+    snapshot_at_shutdown: bool,
 }
 
 #[derive(clap::Subcommand, Debug)]
@@ -296,6 +300,7 @@ fn make_db_config(config: &Cli) -> anyhow::Result<DbConfig> {
         max_total_response_size: config.max_total_response_size.as_u64(),
         snapshot_exec: config.snapshot_exec.clone(),
         checkpoint_interval: config.checkpoint_interval_s.map(Duration::from_secs),
+        snapshot_at_shutdown: config.snapshot_at_shutdown,
     })
 }
 

libsql-server/src/namespace/mod.rs

@@ -351,16 +351,18 @@ struct NamespaceStoreInner<M: MakeNamespace> {
     make_namespace: M,
     allow_lazy_creation: bool,
     has_shutdown: AtomicBool,
+    snapshot_at_shutdown: bool,
 }
 
 impl<M: MakeNamespace> NamespaceStore<M> {
-    pub fn new(make_namespace: M, allow_lazy_creation: bool) -> Self {
+    pub fn new(make_namespace: M, allow_lazy_creation: bool, snapshot_at_shutdown: bool) -> Self {
         Self {
             inner: Arc::new(NamespaceStoreInner {
                 store: Default::default(),
                 make_namespace,
                 allow_lazy_creation,
                 has_shutdown: AtomicBool::new(false),
+                snapshot_at_shutdown,
             }),
         }
     }
@@ -598,7 +600,7 @@ impl<M: MakeNamespace> NamespaceStore<M> {
         self.inner.has_shutdown.store(true, Ordering::Relaxed);
         let mut lock = self.inner.store.write().await;
         for (name, ns) in lock.drain() {
-            ns.shutdown().await?;
+            ns.shutdown(self.inner.snapshot_at_shutdown).await?;
             trace!("shutdown namespace: `{}`", name);
         }
         Ok(())
@@ -645,9 +647,11 @@ impl<T: Database> Namespace<T> {
         Ok(())
     }
 
-    async fn shutdown(mut self) -> anyhow::Result<()> {
+    async fn shutdown(mut self, should_checkpoint: bool) -> anyhow::Result<()> {
         self.tasks.shutdown().await;
-        self.checkpoint().await?;
+        if should_checkpoint {
+            self.checkpoint().await?;
+        }
         self.db.shutdown().await?;
         Ok(())
     }

libsql-server/src/test/bottomless.rs

@@ -50,6 +50,7 @@ async fn configure_server(
             max_total_response_size: 10000000 * 4096,
             snapshot_exec: None,
             checkpoint_interval: Some(Duration::from_secs(3)),
+            snapshot_at_shutdown: false,
         },
         admin_api_config: None,
         disable_namespaces: true,