fix(persistence): close 8 CodeRabbit PR #136 durability findings + aof_manifest split

src/server/embedded.rs

@@ -57,6 +57,18 @@ use crate::shard::Shard;
 use crate::shard::mesh::{CHANNEL_BUFFER_SIZE, ChannelMesh};
 use crate::shard::shared_databases::ShardDatabases;
 
+/// Whether pre-loop recovery (`Shard::restore_from_persistence`) must run.
+///
+/// Recovery is required when EITHER durable persistence (`persistence_dir`) OR
+/// disk-offload is configured. The disk-offload case is the subtle one: under
+/// `--appendonly no` + disk-offload, `persistence_dir` is `None`, yet the cold
+/// tier still has on-disk state whose baseline must be restored on restart —
+/// skipping it leaves cold keys unrecoverable. Mirrors the `main.rs` gate.
+#[inline]
+fn should_run_recovery(persistence_dir: Option<&str>, disk_offload_enabled: bool) -> bool {
+    persistence_dir.is_some() || disk_offload_enabled
+}
+
 /// Run an embedded sharded Moon server until `cancel` is fired.
 ///
 /// Behaves like the production `main.rs` startup path but with cluster,
@@ -218,8 +230,13 @@ pub async fn run_embedded(
                 config.initial_keyspace_hint,
                 config.to_runtime_config(),
             );
-            if let Some(ref dir) = persistence_dir {
-                shard.restore_from_persistence(dir, disk_offload_base.as_deref());
+            // Recover whenever persistence OR disk-offload is configured. Under
+            // `--appendonly no` + disk-offload, persistence_dir is None but the
+            // cold tier still needs its baseline restored on restart, else cold
+            // keys are unrecoverable (mirrors main.rs).
+            if should_run_recovery(persistence_dir.as_deref(), disk_offload_base.is_some()) {
+                let recover_dir = persistence_dir.as_deref().unwrap_or(config.dir.as_str());
+                shard.restore_from_persistence(recover_dir, disk_offload_base.as_deref());
             }
             if let Some(ref offload_base) = disk_offload_base {
                 let shard_dir = offload_base.join(format!("shard-{}", id));
@@ -355,8 +372,10 @@ pub async fn run_embedded(
     }
     let _ = snap_tx; // keep the watch sender alive for the shard's snap_rx clones
 
-    // Run the sharded listener until cancelled.
-    let per_shard_accept = cfg!(target_os = "linux");
+    // Run the sharded listener until cancelled. Use the central accept path
+    // (per_shard_accept = false) like main.rs: the per-shard tokio/Linux accept
+    // loop has a bind-order race that can hang serving (fixed in PR #136).
+    let per_shard_accept = false;
     let listener_result = server::listener::run_sharded(
         config,
         conn_txs,
@@ -463,3 +482,25 @@ fn panic_message<'a>(payload: &'a Box<dyn std::any::Any + Send + 'static>) -> &'
         "<non-string panic payload>"
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::should_run_recovery;
+
+    #[test]
+    fn recovery_runs_for_disk_offload_without_persistence() {
+        // The regression: `--appendonly no` + disk-offload means persistence_dir
+        // is None but the cold tier must still be recovered on restart.
+        assert!(
+            should_run_recovery(None, true),
+            "disk-offload-only (appendonly=no) must still run cold-tier recovery"
+        );
+        // Other combinations behave as before.
+        assert!(should_run_recovery(Some("/var/lib/moon"), false));
+        assert!(should_run_recovery(Some("/var/lib/moon"), true));
+        assert!(
+            !should_run_recovery(None, false),
+            "no persistence and no disk-offload: nothing to recover"
+        );
+    }
+}

src/shard/conn_accept.rs

@@ -301,6 +301,9 @@ pub(crate) fn spawn_migrated_tokio_connection(
     shard_id: usize,
     num_shards: usize,
     config_port: u16,
+    spill_sender: &Option<flume::Sender<crate::storage::tiered::spill_thread::SpillRequest>>,
+    spill_file_id: &Rc<std::cell::Cell<u64>>,
+    disk_offload_dir: &Option<std::path::PathBuf>,
 ) {
     use std::os::unix::io::FromRawFd;
 
@@ -365,6 +368,13 @@ pub(crate) fn spawn_migrated_tokio_connection(
 
             // Pool is built by the spawn site and threaded through here.
             let pool_for_ctx = aof_pool.as_ref().map(Arc::clone);
+            // Preserve the disk-offload spill context across migration — without
+            // this a migrated connection silently falls back to the non-spilling
+            // eviction path even when disk-offload is enabled, so cold-tier
+            // durability would depend on whether the connection had migrated.
+            let spill_tx = spill_sender.clone();
+            let spill_fid = spill_file_id.clone();
+            let do_dir = disk_offload_dir.clone();
             let conn_ctx = crate::server::conn::ConnectionContext::new(
                 sdbs,
                 shard_id,
@@ -390,9 +400,9 @@ pub(crate) fn spawn_migrated_tokio_connection(
                 all_regs,
                 all_rsm,
                 aff,
-                None,                             // spill_sender
-                Rc::new(std::cell::Cell::new(0)), // spill_file_id
-                None,                             // disk_offload_dir
+                spill_tx,  // spill_sender (preserved across migration)
+                spill_fid, // spill_file_id
+                do_dir,    // disk_offload_dir
             );
 
             // State restoration happens directly via migrated_state parameter —

src/shard/event_loop.rs

@@ -1244,6 +1244,7 @@ impl super::Shard {
                                 &cached_clock, &remote_sub_map_arc, &all_pubsub_registries,
                                 &all_remote_sub_maps, &affinity_tracker,
                                 shard_id, num_shards, config_port,
+                                &spill_sender, &spill_file_id, &disk_offload_dir,
                             );
                         }
                         #[cfg(feature = "runtime-monoio")]
@@ -1343,6 +1344,7 @@ impl super::Shard {
                                 &cached_clock, &remote_sub_map_arc, &all_pubsub_registries,
                                 &all_remote_sub_maps, &affinity_tracker,
                                 shard_id, num_shards, config_port,
+                                &spill_sender, &spill_file_id, &disk_offload_dir,
                             );
                         }
                         #[cfg(feature = "runtime-monoio")]

src/shard/persistence_tick.rs

@@ -369,7 +369,11 @@ pub(crate) fn drain_and_shutdown_spill(
         apply_spill_completions(spill_t, shard_manifest, shard_databases, shard_id);
     }
     if let Some(st) = spill_thread.take() {
-        st.shutdown();
+        // shutdown() returns any completions from the thread's final buffer
+        // flush that the drain above did not see; apply them so those cold keys
+        // are not lost (file on disk but never recorded in the manifest).
+        let leftover = st.shutdown();
+        apply_completion_vec(leftover, shard_manifest, shard_databases, shard_id);
         tracing::info!("Shard {}: spill background thread shut down", shard_id);
     }
 }
@@ -387,6 +391,18 @@ pub(crate) fn apply_spill_completions(
     shard_id: usize,
 ) {
     let completions = spill_thread.drain_completions();
+    apply_completion_vec(completions, shard_manifest, shard_databases, shard_id);
+}
+
+/// Apply a batch of spill completions: ONE manifest `add_file`+commit per file,
+/// one `cold_index` insert per KV entry within it. Shared by the live drain
+/// (`apply_spill_completions`) and the shutdown final-flush drain.
+fn apply_completion_vec(
+    completions: Vec<crate::storage::tiered::spill_thread::SpillCompletion>,
+    shard_manifest: &mut Option<crate::persistence::manifest::ShardManifest>,
+    shard_databases: &std::sync::Arc<super::shared_databases::ShardDatabases>,
+    shard_id: usize,
+) {
     if completions.is_empty() {
         return;
     }

src/storage/eviction.rs

@@ -227,16 +227,6 @@ pub fn try_evict_if_needed_with_spill_and_total(
     Ok(())
 }
 
-/// Check if eviction is needed, spilling evicted entries asynchronously via
-/// a background `SpillThread` instead of doing synchronous pwrite.
-///
-/// The async path: extracts key/value bytes, removes entry from DashTable
-/// (freeing RAM immediately), then sends a `SpillRequest` to the background
-/// thread. The pwrite is best-effort -- if the channel is full, the request
-/// is dropped (entry already removed from RAM).
-///
-/// Callers must poll `SpillThread::drain_completions()` to apply manifest
-/// and ColdIndex updates from completed spills.
 /// Seed value for a shard's spill `file_id` counter after recovery.
 ///
 /// On restart, AOF/RDB replay re-populates the hot tier and the persistence
@@ -296,6 +286,21 @@ pub fn next_spill_file_id_seed(shard_dir: Option<&Path>) -> u64 {
     }
 }
 
+/// Evict over-budget entries, spilling them to disk asynchronously via a
+/// background `SpillThread` instead of doing a synchronous pwrite on the event
+/// loop.
+///
+/// Per victim (`evict_one_async_spill`): serialize the key/value into a
+/// `SpillRequest`, **`try_send` it to the spill channel BEFORE removing the
+/// entry from the hot tier, and only remove from RAM once the send succeeds.**
+/// This ordering is fail-safe: if the channel is full or disconnected the send
+/// fails, the key stays in RAM, and the eviction loop bails out to retry on the
+/// next tick — no acknowledged data is lost. The worst case under sustained
+/// backpressure is keys remaining resident (eventually an OOM write-rejection
+/// once at budget), which is correct behaviour, not data loss.
+///
+/// Callers must poll `SpillThread::drain_completions()` to apply the manifest
+/// and `ColdIndex` updates produced by completed spills.
 pub fn try_evict_if_needed_async_spill(
     db: &mut Database,
     config: &RuntimeConfig,
@@ -357,76 +362,6 @@ pub fn try_evict_if_needed_async_spill_with_total(
     Ok(())
 }
 
-/// Evict entries to bring memory under maxmemory, returning removed
-/// (key, Entry) pairs for deferred spill OUTSIDE the write lock.
-///
-/// Inside the lock: only find_victim + db.remove (~600ns per eviction).
-/// The caller extracts value bytes from the owned Entry after releasing
-/// the lock, then sends SpillRequests to the background thread.
-pub fn try_evict_deferred(
-    db: &mut Database,
-    config: &RuntimeConfig,
-) -> Result<smallvec::SmallVec<[(Bytes, crate::storage::entry::Entry); 2]>, Frame> {
-    if config.maxmemory == 0 {
-        return Ok(smallvec::SmallVec::new());
-    }
-
-    // Per-shard budget (see `try_evict_if_needed_with_spill_and_total`).
-    let budget = config.maxmemory_per_shard();
-    let total_memory = db.estimated_memory();
-    if total_memory <= budget {
-        return Ok(smallvec::SmallVec::new());
-    }
-
-    let policy = EvictionPolicy::from_str(&config.maxmemory_policy);
-    let mut evicted = smallvec::SmallVec::new();
-    let mut current_total = total_memory;
-
-    while current_total > budget {
-        if policy == EvictionPolicy::NoEviction {
-            return Err(oom_error());
-        }
-
-        let victim = find_victim_for_policy(db, config, &policy);
-        let key = match victim {
-            Some(k) => k,
-            None => return Err(oom_error()),
-        };
-
-        let before = db.estimated_memory();
-        let key_bytes = Bytes::copy_from_slice(key.as_bytes());
-        if let Some(entry) = db.remove(key.as_bytes()) {
-            evicted.push((key_bytes, entry));
-        }
-        let after = db.estimated_memory();
-        current_total = current_total.saturating_sub(before.saturating_sub(after));
-    }
-
-    Ok(evicted)
-}
-
-/// Find a victim key using the given eviction policy.
-fn find_victim_for_policy(
-    db: &Database,
-    config: &RuntimeConfig,
-    policy: &EvictionPolicy,
-) -> Option<CompactKey> {
-    match policy {
-        EvictionPolicy::NoEviction => None,
-        EvictionPolicy::AllKeysLru => find_victim_lru(db, config.maxmemory_samples, false),
-        EvictionPolicy::AllKeysLfu => {
-            find_victim_lfu(db, config.maxmemory_samples, config.lfu_decay_time, false)
-        }
-        EvictionPolicy::AllKeysRandom => find_victim_random(db, false),
-        EvictionPolicy::VolatileLru => find_victim_lru(db, config.maxmemory_samples, true),
-        EvictionPolicy::VolatileLfu => {
-            find_victim_lfu(db, config.maxmemory_samples, config.lfu_decay_time, true)
-        }
-        EvictionPolicy::VolatileRandom => find_victim_random(db, true),
-        EvictionPolicy::VolatileTtl => find_victim_volatile_ttl(db, config.maxmemory_samples),
-    }
-}
-
 /// Evict a single key via the async spill path.
 ///
 /// Extracts the entry, removes it from DashTable (immediate RAM relief),