feat(gateway): let an operator bound what a machine loss costs the KV log

Author: kvinwang

dstack/gateway/gateway.toml

@@ -175,6 +175,21 @@ bootnode = ""
 data_dir = "/dstack-gateway/data"
 # Interval for periodic persistence of WaveKV data (e.g., "5s", "1m", "1h")
 persist_interval = "5m"
+# How long a KV write may sit in the page cache before the write-ahead log is
+# forced to disk.
+#
+# An fsync costs about 100us on an NVMe host and milliseconds on the virtio disk
+# a CVM gets, and it runs under the store lock, so paying one per write bounds
+# how fast this gateway can accept registrations and stalls its readers while it
+# does. Deferring it bounds what losing the *machine* costs — writes made in the
+# last window — and changes nothing about losing the process: a panic, an OOM
+# kill or a restart loses nothing either way, because the bytes are already with
+# the kernel.
+#
+# A cluster recovers the window from its peers on restart. A single-node gateway
+# does not, which is the case for setting this to "0s" — force every write
+# before it returns, as releases before WaveKV 2.1 always did.
+wal_sync_interval = "5s"
 # Enable periodic sync of instance connections to KV store
 sync_connections_enabled = true
 # Interval for syncing instance connections to KV store

dstack/gateway/src/config.rs

@@ -543,6 +543,10 @@ pub struct SyncConfig {
     /// Interval for periodic WAL persistence (default: 10s)
     #[serde(with = "serde_duration")]
     pub persist_interval: Duration,
+    /// How long a KV write may sit in the page cache before the write-ahead log
+    /// is forced to disk. Zero forces every write before it returns.
+    #[serde(with = "serde_duration")]
+    pub wal_sync_interval: Duration,
     /// Enable periodic sync of instance connections to KV store
     pub sync_connections_enabled: bool,
     /// Interval for syncing instance connections to KV store

dstack/gateway/src/distributed_certbot.rs

@@ -807,15 +807,16 @@ mod tests {
 
     fn test_certbot(data_dir: &std::path::Path) -> DistributedCertBot {
         let kv_store =
-            Arc::new(KvStore::new(1, vec![], data_dir).expect("failed to create kv store"));
+            Arc::new(KvStore::new(1, vec![], data_dir, None).expect("failed to create kv store"));
         DistributedCertBot::new(kv_store, Arc::new(CertResolver::new()), None)
     }
 
     #[test]
     fn lock_writes_wake_the_persistent_push_path() {
         let data_dir = tempfile::tempdir().expect("failed to create temp dir");
-        let kv_store =
-            Arc::new(KvStore::new(1, vec![], data_dir.path()).expect("failed to create kv store"));
+        let kv_store = Arc::new(
+            KvStore::new(1, vec![], data_dir.path(), None).expect("failed to create kv store"),
+        );
         let notifier = Arc::new(CountingNotifier::default());
         let certbot = DistributedCertBot::new(
             kv_store,

dstack/gateway/src/kv/mod.rs

@@ -688,13 +688,28 @@ impl KvStore {
     /// operator, which for a single-node deployment holding the only copy of
     /// the ACME account and DNS credentials is the difference between a restart
     /// and a rebuild.
+    ///
+    /// `wal_sync_interval` is how long a write may sit in the page cache before
+    /// the log is forced to disk; `None` forces every write before it returns.
+    /// It applies only to the persistent store — the ephemeral one keeps no log
+    /// and never touches the disk.
     pub fn new(
         my_node_id: NodeId,
         peer_ids: Vec<NodeId>,
         data_dir: impl AsRef<Path>,
+        wal_sync_interval: Option<Duration>,
     ) -> Result<Self> {
         let data_dir = data_dir.as_ref();
-        let persistent = match Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir) {
+        let node_config = wavekv::NodeConfig {
+            wal_sync_interval,
+            ..Default::default()
+        };
+        let persistent = match Node::with_persistence_and_config(
+            my_node_id,
+            peer_ids.clone(),
+            data_dir,
+            node_config.clone(),
+        ) {
             Ok(node) => node,
             Err(err) if is_storage_failure(&err) => {
                 return Err(err).with_context(|| {
@@ -721,8 +736,13 @@ impl KvStore {
                     data_dir.display(),
                     quarantined.display(),
                 );
-                Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir)
-                    .context("failed to create persistent wavekv node on a fresh data dir")?
+                Node::with_persistence_and_config(
+                    my_node_id,
+                    peer_ids.clone(),
+                    data_dir,
+                    node_config,
+                )
+                .context("failed to create persistent wavekv node on a fresh data dir")?
             }
         };
 
@@ -1033,6 +1053,15 @@ impl KvStore {
         self.persistent.persist_if_dirty()
     }
 
+    /// Force the write-ahead log to disk if the configured window has elapsed.
+    ///
+    /// Returns whether an fsync happened. A no-op when no window is configured
+    /// — every write was already forced — or when nothing has been written
+    /// since the last one, so an idle gateway costs a lock acquisition.
+    pub fn sync_wal_if_due(&self) -> Result<bool> {
+        self.persistent.sync_wal_if_due()
+    }
+
     // ==================== Peer Management ====================
 
     pub fn add_peer(&self, peer_id: NodeId) -> Result<()> {
@@ -1625,7 +1654,7 @@ mod acme_credentials_tests {
     use super::*;
 
     fn test_kv(data_dir: &std::path::Path) -> KvStore {
-        KvStore::new(1, vec![], data_dir).expect("failed to create kv store")
+        KvStore::new(1, vec![], data_dir, None).expect("failed to create kv store")
     }
 
     #[test]
@@ -1772,7 +1801,7 @@ mod value_encoding_tests {
     #[test]
     fn legacy_positional_records_are_still_readable() {
         let dir = tempfile::tempdir().expect("failed to create temp dir");
-        let kv = KvStore::new(1, vec![], dir.path()).expect("failed to create kv store");
+        let kv = KvStore::new(1, vec![], dir.path(), None).expect("failed to create kv store");
 
         let legacy = rmp_serde::encode::to_vec(&CertRenewLock {
             started_at: 1_700_000_000,
@@ -1836,6 +1865,51 @@ mod value_encoding_tests {
         }
     }
 
+    /// The configured window has to reach the store, not just the config file.
+    ///
+    /// An fsync per write runs under the store lock, so it bounds how fast this
+    /// gateway accepts registrations; the window is what buys that back, and it
+    /// buys nothing if the interval stops at `SyncConfig`.
+    #[test]
+    fn a_configured_window_holds_writes_out_of_the_disk_until_it_elapses() {
+        let dir = tempfile::tempdir().expect("failed to create temp dir");
+        let window = Duration::from_millis(20);
+        let kv = KvStore::new(1, vec![], dir.path(), Some(window)).expect("kv store");
+
+        kv.set_node_status(1, NodeStatus::Up).expect("write");
+        assert_eq!(
+            kv.persistent().read().wal_sync_count(),
+            0,
+            "a write inside the window must not reach the disk"
+        );
+        assert!(
+            !kv.sync_wal_if_due().expect("sync check"),
+            "nothing is due before the window elapses"
+        );
+
+        std::thread::sleep(window * 3);
+        assert!(kv.sync_wal_if_due().expect("sync"), "the window elapsed");
+        assert_eq!(kv.persistent().read().wal_sync_count(), 1);
+    }
+
+    /// Without a window every write is on the disk before it returns, which is
+    /// what every release before this one did and what a single-node gateway
+    /// holding the only copy of its ACME account still wants.
+    #[test]
+    fn without_a_window_every_write_reaches_the_disk_before_it_returns() {
+        let dir = tempfile::tempdir().expect("failed to create temp dir");
+        let kv = KvStore::new(1, vec![], dir.path(), None).expect("kv store");
+
+        kv.set_node_status(1, NodeStatus::Up).expect("write");
+        kv.set_node_status(2, NodeStatus::Down).expect("write");
+
+        assert_eq!(kv.persistent().read().wal_sync_count(), 2);
+        assert!(
+            !kv.sync_wal_if_due().expect("sync check"),
+            "with no window there is never anything owing"
+        );
+    }
+
     /// An instance record as some later release will declare it.
     #[derive(Debug, Serialize, Deserialize)]
     struct FutureInstanceData {
@@ -1854,7 +1928,7 @@ mod value_encoding_tests {
     #[test]
     fn a_record_keeps_the_fields_its_writer_does_not_know() {
         let dir = tempfile::tempdir().expect("failed to create temp dir");
-        let kv = KvStore::new(1, vec![], dir.path()).expect("failed to create kv store");
+        let kv = KvStore::new(1, vec![], dir.path(), None).expect("failed to create kv store");
 
         let written_by_a_newer_node = encode(&FutureInstanceData {
             app_id: "app".to_string(),
@@ -1913,7 +1987,7 @@ mod value_encoding_tests {
     #[test]
     fn a_snapshot_does_not_inherit_the_previous_writes_fields() {
         let dir = tempfile::tempdir().expect("failed to create temp dir");
-        let kv = KvStore::new(1, vec![], dir.path()).expect("failed to create kv store");
+        let kv = KvStore::new(1, vec![], dir.path(), None).expect("failed to create kv store");
 
         #[derive(Debug, Serialize, Deserialize)]
         struct FutureCertData {
@@ -1973,7 +2047,7 @@ mod sync_wire_tests {
     use wavekv::sync::SyncEnvelope;
 
     fn store(dir: &std::path::Path, id: NodeId, peers: Vec<NodeId>) -> KvStore {
-        KvStore::new(id, peers, dir).expect("failed to create kv store")
+        KvStore::new(id, peers, dir, None).expect("failed to create kv store")
     }
 
     #[test]
@@ -2048,7 +2122,8 @@ mod wavekv_v1_migration_tests {
                 .expect("write trailing v1 WAL entry");
         }
 
-        let upgraded = KvStore::new(1, Vec::new(), dir.path()).expect("open v1 data after upgrade");
+        let upgraded =
+            KvStore::new(1, Vec::new(), dir.path(), None).expect("open v1 data after upgrade");
         assert_eq!(
             upgraded
                 .persistent()
@@ -2078,7 +2153,8 @@ mod wavekv_v1_migration_tests {
         upgraded.persist_if_dirty().expect("persist upgraded data");
         drop(upgraded);
 
-        let restarted = KvStore::new(1, Vec::new(), dir.path()).expect("restart upgraded store");
+        let restarted =
+            KvStore::new(1, Vec::new(), dir.path(), None).expect("restart upgraded store");
         for (key, expected) in [(key, value), (wal_key, wal_value), (new_key, new_value)] {
             assert_eq!(
                 restarted
@@ -2131,7 +2207,7 @@ mod corruption_tests {
     use super::*;
 
     fn test_kv(data_dir: &std::path::Path) -> KvStore {
-        KvStore::new(1, vec![], data_dir).expect("failed to create kv store")
+        KvStore::new(1, vec![], data_dir, None).expect("failed to create kv store")
     }
 
     fn put_raw(kv: &KvStore, key: &str, value: &[u8]) {
@@ -2245,7 +2321,7 @@ mod corruption_tests {
         let data_dir = dir.path().join("kv");
         std::fs::write(&data_dir, b"not a directory").expect("failed to create blocker");
 
-        let Err(err) = KvStore::new(1, vec![], &data_dir) else {
+        let Err(err) = KvStore::new(1, vec![], &data_dir, None) else {
             panic!("startup must fail when the storage is unusable");
         };
         assert!(
@@ -2285,7 +2361,8 @@ mod corruption_tests {
         // replicated, so it must not keep the gateway from booting.
         std::fs::write(data_dir.join("node_1.wal"), b"garbage").expect("failed to corrupt wal");
 
-        let kv = KvStore::new(1, vec![], &data_dir).expect("startup must survive a corrupt wal");
+        let kv =
+            KvStore::new(1, vec![], &data_dir, None).expect("startup must survive a corrupt wal");
         let loaded = kv.load_all_instances();
         assert!(loaded.decoded.is_empty());
         assert!(loaded.undecodable.is_empty());

dstack/gateway/src/main_service.rs

@@ -263,8 +263,13 @@ impl ProxyInner {
 
         // Initialize WaveKV store without peers (peers will be added dynamically from bootnode)
         let kv_store = Arc::new(
-            KvStore::new(config.sync.node_id, vec![], &config.sync.data_dir)
-                .context("failed to initialize WaveKV store")?,
+            KvStore::new(
+                config.sync.node_id,
+                vec![],
+                &config.sync.data_dir,
+                (!config.sync.wal_sync_interval.is_zero()).then_some(config.sync.wal_sync_interval),
+            )
+            .context("failed to initialize WaveKV store")?,
         );
         info!(
             "WaveKV store initialized: node_id={}, sync_enabled={}",
@@ -1035,6 +1040,27 @@ fn start_wavekv_watch_task(proxy: Proxy) -> Result<()> {
         info!("WaveKV: periodic persistence enabled (interval: {persist_interval:?})");
     }
 
+    // Force the write-ahead log on the window the operator configured. Nothing
+    // else forces it on a schedule — a snapshot does, but that is minutes apart
+    // — so without this the window would be a hope rather than a bound. Ticking
+    // at the window itself is enough to make it one: the store measures from
+    // its last fsync, so a write is forced at the first tick that finds the
+    // window elapsed, never more than one window after it landed.
+    let wal_sync_interval = proxy.config.sync.wal_sync_interval;
+    if !wal_sync_interval.is_zero() {
+        let kv_store_for_wal = kv_store.clone();
+        tokio::spawn(async move {
+            let mut ticker = tokio::time::interval(wal_sync_interval);
+            loop {
+                ticker.tick().await;
+                if let Err(err) = kv_store_for_wal.sync_wal_if_due() {
+                    error!("WaveKV: forcing the write-ahead log failed: {err:?}");
+                }
+            }
+        });
+        info!("WaveKV: deferred WAL sync enabled (window: {wal_sync_interval:?})");
+    }
+
     // Start periodic connection sync task
     if proxy.config.sync.sync_connections_enabled {
         let sync_interval = proxy.config.sync.sync_connections_interval;