Enable RocksDB atomic_flush to keep the index crash-consistent during IBD

src/server/mod.rs

@@ -11,6 +11,7 @@ use crate::inc_connection_error_counter;
 use crate::server::preload::headers;
 use crate::store::memory::MemoryStore;
 use crate::store::AnyStore;
+use crate::store::Store;
 use crate::threads::blocks::blocks_infallible;
 use crate::threads::mempool::mempool_sync_infallible;
 use crate::threads::zmq::rawtx_listener_infallible;
@@ -632,6 +633,13 @@ pub async fn inner_main(
     }
     h4.await.unwrap();
 
+    // Background writers have stopped, so persist any buffered data before
+    // exiting. While the WAL is disabled during IBD, un-flushed memtables would
+    // otherwise be dropped on process exit, losing sync progress (and, without
+    // atomic_flush, risking an inconsistent on-disk state). Cheap at tip where
+    // memtables are nearly empty.
+    state.store.flush();
+
     log::info!("shutting down gracefully");
     Ok(())
 }

src/store/db.rs

@@ -108,6 +108,22 @@ impl DBStore {
         Ok(reordered)
     }
 
+    /// Flush every column-family memtable to SST files.
+    ///
+    /// With `atomic_flush` enabled at open (see `open`), this produces a
+    /// crash-consistent on-disk checkpoint across all column families even while
+    /// the WAL is disabled during IBD. Called when IBD finishes and on graceful
+    /// shutdown so that sync progress is not lost on a clean restart.
+    fn flush_all_cfs(&self) {
+        for cf_name in COLUMN_FAMILIES {
+            if let Some(cf) = self.db.cf_handle(cf_name) {
+                self.db
+                    .flush_cf(&cf)
+                    .unwrap_or_else(|e| log::error!("failed to flush CF {cf_name}: {e}"));
+            }
+        }
+    }
+
     fn create_cf_descriptors(shared_db_cache_mb: u64) -> Vec<rocksdb::ColumnFamilyDescriptor> {
         let cache_size = (shared_db_cache_mb * 1024 * 1024) as usize;
         // HyperClockCache is lock-free, reducing mutex contention under concurrent reads.
@@ -173,12 +189,12 @@ impl DBStore {
             .collect()
     }
 
-    pub fn open(
-        path: &Path,
-        shared_db_cache_mb: u64,
-        enable_statistics: bool,
-        reorg_data_keep_heights: u32,
-    ) -> Result<Self> {
+    /// Build the database-wide options used to open the store.
+    ///
+    /// Kept separate so tests can open with the exact production options (most
+    /// importantly `atomic_flush`, see below) while layering on test-only
+    /// settings such as crash simulation.
+    fn db_options(enable_statistics: bool) -> Options {
         let mut db_opts = Options::default();
 
         // Enable statistics collection for detailed metrics including bloom filter stats
@@ -195,6 +211,33 @@ impl DBStore {
         db_opts.increase_parallelism(parallelism);
         db_opts.set_max_background_jobs(parallelism);
 
+        // The WAL is disabled during IBD (see `write`) to speed up the initial
+        // sync. A single block update is one WriteBatch spanning several column
+        // families (utxos, history, block hash/height). Without the WAL each
+        // column family flushes its memtable to disk independently, so a hard
+        // crash mid-sync (OOM, SIGKILL, host loss) can persist some families
+        // ahead of others — e.g. the block-height marker advances while a UTXO
+        // that a later block spends was never flushed. On restart that surfaces
+        // as a fatal "every utxo must exist when spent" and an unrecoverable
+        // crash loop. `atomic_flush` makes multi-family flushes atomic, so the
+        // recovered on-disk state is always a consistent prefix even with the
+        // WAL off. It only changes flush grouping, not the write path, so the
+        // cost at tip (where flushes are infrequent) is negligible; the extra
+        // work falls on IBD, which is exactly when the guarantee is needed.
+        // Regression-tested by `test::atomic_flush_keeps_utxos_consistent_on_crash`.
+        db_opts.set_atomic_flush(true);
+
+        db_opts
+    }
+
+    pub fn open(
+        path: &Path,
+        shared_db_cache_mb: u64,
+        enable_statistics: bool,
+        reorg_data_keep_heights: u32,
+    ) -> Result<Self> {
+        let db_opts = Self::db_options(enable_statistics);
+
         let db = rocksdb::DB::open_cf_descriptors(
             &db_opts,
             path,
@@ -728,16 +771,14 @@ impl Store for DBStore {
 
     fn ibd_finished(&self) {
         log::info!("Initial block download finished, flushing memtables before enabling WAL...");
-        for cf_name in COLUMN_FAMILIES {
-            if let Some(cf) = self.db.cf_handle(cf_name) {
-                self.db
-                    .flush_cf(&cf)
-                    .unwrap_or_else(|e| log::error!("failed to flush CF {cf_name}: {e}"));
-            }
-        }
+        self.flush_all_cfs();
         log::info!("Memtables flushed, setting ibd to false");
         self.ibd.store(false, Ordering::Relaxed);
     }
+
+    fn flush(&self) {
+        self.flush_all_cfs();
+    }
 }
 
 fn serialize_outpoint(o: &OutPoint) -> Vec<u8> {
@@ -1042,4 +1083,271 @@ mod test {
             "ScriptHash (u64) natural ordering must match binary encoding ordering"
         );
     }
+
+    /// Copy the on-disk files of a RocksDB directory into `dst`.
+    ///
+    /// The result is the state that would survive an ungraceful crash: only data
+    /// already flushed to disk is captured, anything still in a memtable is lost.
+    /// (A RocksDB directory is flat — CURRENT, MANIFEST, OPTIONS, *.sst, etc.)
+    fn snapshot_on_disk_state(src: &std::path::Path, dst: &std::path::Path) {
+        for entry in std::fs::read_dir(src).unwrap() {
+            let entry = entry.unwrap();
+            if entry.file_type().unwrap().is_file() {
+                std::fs::copy(entry.path(), dst.join(entry.file_name())).unwrap();
+            }
+        }
+    }
+
+    /// Regression test for the "every utxo must exist when spent" corruption.
+    ///
+    /// During IBD the WAL is disabled (see `write`). A block update writes the
+    /// UTXO set and the block-height marker in a single batch, but those land in
+    /// separate column families, which RocksDB flushes independently. If the
+    /// process is killed after the height marker's memtable is flushed but
+    /// before the UTXO's is, the persisted state has a height ahead of its
+    /// UTXOs. On restart the first spend of the missing UTXO panics, and the
+    /// node crash-loops forever with no way to self-heal.
+    ///
+    /// `atomic_flush` (set in `db_options`) makes those background flushes
+    /// atomic across families, so the persisted state is always a consistent
+    /// prefix. We reproduce a hard crash by snapshotting the on-disk directory
+    /// the moment the background flush lands — before any clean-shutdown flush
+    /// could mask the bug — and reopening from the snapshot. Remove
+    /// `set_atomic_flush(true)` and this test fails: the UTXO is gone after the
+    /// simulated crash, exactly the production failure.
+    ///
+    /// Note: the flush has to be the automatic (memtable-pressure) one — a
+    /// manual single-family flush does not co-flush other families even with
+    /// atomic_flush — so we give the families a tiny write buffer and write
+    /// enough height markers to force background flushes.
+    #[test]
+    fn atomic_flush_keeps_utxos_consistent_on_crash() {
+        use crate::store::BlockMeta;
+        use std::sync::atomic::AtomicBool;
+
+        let _ = env_logger::try_init();
+
+        let live = tempfile::TempDir::new().unwrap();
+        let crashed = tempfile::TempDir::new().unwrap();
+
+        let outpoint = {
+            let mut o = OutPoint::null();
+            o.vout = 6;
+            o
+        };
+        let script_hash: u64 = 777;
+
+        {
+            // Production DB options — crucially `atomic_flush`, from `db_options`,
+            // which is the line under test. The per-family write buffers are sized
+            // tiny here so writes trigger real background flushes (the path
+            // atomic_flush governs); that only controls *when* flushes happen, not
+            // their cross-family atomicity.
+            let db_opts = DBStore::db_options(false);
+            let cf_descriptors: Vec<_> = super::COLUMN_FAMILIES
+                .iter()
+                .map(|&name| {
+                    let mut cf_opts = rocksdb::Options::default();
+                    cf_opts.set_write_buffer_size(16 * 1024);
+                    cf_opts.set_max_write_buffer_number(2);
+                    rocksdb::ColumnFamilyDescriptor::new(name, cf_opts)
+                })
+                .collect();
+            let db = DBStore {
+                db: rocksdb::DB::open_cf_descriptors(&db_opts, live.path(), cf_descriptors)
+                    .unwrap(),
+                salt: 0,
+                ibd: AtomicBool::new(true), // IBD => writes skip the WAL
+                reorg_data_keep_heights: 6,
+            };
+
+            // An earlier block creates the UTXO (utxo CF), WAL-off.
+            let created: BTreeMap<_, _> = [(outpoint, script_hash)].into_iter().collect();
+            let mut batch = rocksdb::WriteBatch::default();
+            db.insert_utxos(&mut batch, &created).unwrap();
+            db.write(batch).unwrap();
+
+            // Later blocks advance the height marker (hashes CF), WAL-off,
+            // written in many small batches. With the tiny write buffer this
+            // overflows the memtable and triggers background flushes. (A single
+            // oversized batch does not — RocksDB only re-checks the flush
+            // condition on a subsequent write — so the writes must be chunked,
+            // which also matches how blocks are indexed one at a time.) Under
+            // atomic_flush the first such flush also persists the single UTXO;
+            // without it, the UTXO is left behind in its own memtable, which
+            // never fills because it holds one entry.
+            for chunk in 0..40u32 {
+                let mut batch = rocksdb::WriteBatch::default();
+                for i in 0..500u32 {
+                    let height = chunk * 500 + i + 1;
+                    let meta = BlockMeta::new(height, BlockHash::all_zeros(), 0);
+                    db.set_hash_ts_batch(&mut batch, &meta);
+                }
+                db.write(batch).unwrap();
+            }
+
+            // Wait (bounded) for the height-marker CF to flush to disk. Under
+            // atomic_flush the UTXO CF is flushed in the same atomic step, so this
+            // is also the barrier for the UTXO.
+            let hashes_cf = db.db.cf_handle(super::HASHES_CF).unwrap();
+            let mut flushed = false;
+            for _ in 0..400 {
+                let l0 = db
+                    .db
+                    .property_int_value_cf(&hashes_cf, "rocksdb.num-files-at-level0")
+                    .unwrap()
+                    .unwrap_or(0);
+                if l0 > 0 {
+                    flushed = true;
+                    break;
+                }
+                std::thread::sleep(std::time::Duration::from_millis(25));
+            }
+            assert!(
+                flushed,
+                "test setup: height-marker CF never flushed; raise the write count \
+                 or lower the write buffer size",
+            );
+
+            // Capture the on-disk state == what survives an ungraceful crash,
+            // before the clean drop below flushes everything and masks the bug.
+            snapshot_on_disk_state(live.path(), crashed.path());
+        }
+
+        // Reopen from the crashed snapshot and verify the UTXO survived.
+        let recovered = DBStore::open(crashed.path(), 64, false, 6).unwrap();
+        let found = recovered.get_utxos(&[outpoint]).unwrap();
+        assert_eq!(
+            found,
+            vec![Some(script_hash)],
+            "a UTXO created before the last persisted height marker must survive a \
+             crash during IBD; without atomic_flush it is lost, which is the \
+             unrecoverable 'every utxo must exist when spent' corruption",
+        );
+    }
+
+    /// Benchmark the cost of `atomic_flush` on the IBD write path (WAL off).
+    ///
+    /// Replays a synthetic initial sync — many blocks, each creating UTXOs and
+    /// history entries — through the real `update` path, with `atomic_flush`
+    /// off then on, and prints wall-clock time plus flush/compaction bytes
+    /// (write amplification). Production column-family settings (default 64 MiB
+    /// write buffers) so the numbers reflect a real sync, not a worst case.
+    ///
+    /// Run: `cargo test --release --lib bench_ibd_atomic_flush -- --ignored --nocapture`
+    /// Scale with env vars `BENCH_BLOCKS` / `BENCH_CREATES_PER_BLOCK`.
+    #[test]
+    #[ignore = "benchmark; run with --release --ignored --nocapture"]
+    fn bench_ibd_atomic_flush() {
+        use crate::store::BlockMeta;
+        use std::sync::atomic::AtomicBool;
+        use std::time::Instant;
+
+        let env_u64 = |k: &str, d: u64| {
+            std::env::var(k)
+                .ok()
+                .and_then(|v| v.parse().ok())
+                .unwrap_or(d)
+        };
+        let blocks: u32 = env_u64("BENCH_BLOCKS", 3_000) as u32;
+        let creates: u32 = env_u64("BENCH_CREATES_PER_BLOCK", 2_000) as u32;
+        // 0 = production default (64 MiB) buffers. Set small (e.g. 256) to force
+        // frequent flushes and expose the worst-case atomic_flush overhead.
+        let write_buf_kib: u64 = env_u64("BENCH_WRITE_BUFFER_KIB", 0);
+        let txid = crate::be::Txid::all_zeros();
+
+        // Pull a ticker count out of the RocksDB statistics dump.
+        let ticker = |stats: &str, key: &str| -> u64 {
+            let needle = format!("{key} ");
+            stats
+                .lines()
+                .find(|l| l.trim_start().starts_with(&needle))
+                .and_then(|l| l.rsplit(':').next())
+                .and_then(|n| n.trim().parse().ok())
+                .unwrap_or(0)
+        };
+
+        let buf_desc = if write_buf_kib == 0 {
+            "64MiB buffers".to_string()
+        } else {
+            format!("{write_buf_kib}KiB buffers")
+        };
+        println!(
+            "\nIBD bench: {blocks} blocks x {creates} creates = {} UTXO writes (WAL off, {buf_desc})",
+            blocks as u64 * creates as u64,
+        );
+        println!(
+            "{:<14}{:>12}{:>14}{:>14}{:>16}",
+            "atomic_flush", "elapsed_s", "flush_MiB", "compact_MiB", "writes/s"
+        );
+
+        for atomic in [false, true] {
+            let dir = tempfile::TempDir::new().unwrap();
+            let mut db_opts = DBStore::db_options(false);
+            db_opts.set_atomic_flush(atomic);
+            db_opts.enable_statistics();
+            let cf_descriptors = if write_buf_kib == 0 {
+                DBStore::create_cf_descriptors(64)
+            } else {
+                super::COLUMN_FAMILIES
+                    .iter()
+                    .map(|&name| {
+                        let mut cf_opts = rocksdb::Options::default();
+                        cf_opts.set_write_buffer_size(write_buf_kib as usize * 1024);
+                        // history CF is written via merge; it needs the operator.
+                        if name == super::HISTORY_CF {
+                            cf_opts.set_merge_operator_associative(
+                                "concat_merge",
+                                super::concat_merge,
+                            );
+                        }
+                        rocksdb::ColumnFamilyDescriptor::new(name, cf_opts)
+                    })
+                    .collect()
+            };
+            let db = DBStore {
+                db: rocksdb::DB::open_cf_descriptors(&db_opts, dir.path(), cf_descriptors).unwrap(),
+                salt: 0,
+                ibd: AtomicBool::new(true), // IBD => WAL off, the path under test
+                reorg_data_keep_heights: 6,
+            };
+
+            let mut vout = 0u32;
+            let start = Instant::now();
+            for height in 1..=blocks {
+                let mut created = BTreeMap::new();
+                let mut history = BTreeMap::new();
+                for _ in 0..creates {
+                    let mut o = OutPoint::null();
+                    o.vout = vout;
+                    let sh = vout as u64;
+                    created.insert(o, sh);
+                    history.insert(sh, vec![TxSeen::new(txid, height, V::Vout(vout))]);
+                    vout = vout.wrapping_add(1);
+                }
+                db.update(
+                    &BlockMeta::new(height, BlockHash::all_zeros(), height),
+                    vec![],
+                    history,
+                    created,
+                )
+                .unwrap();
+            }
+            // Flush remaining memtables, as production does when IBD finishes.
+            db.flush();
+            let elapsed = start.elapsed();
+
+            let stats = db_opts.get_statistics().unwrap_or_default();
+            let mib = |bytes: u64| bytes as f64 / (1024.0 * 1024.0);
+            let total_writes = blocks as f64 * creates as f64;
+            println!(
+                "{:<14}{:>12.2}{:>14.1}{:>14.1}{:>16.0}",
+                atomic,
+                elapsed.as_secs_f64(),
+                mib(ticker(&stats, "rocksdb.flush.write.bytes")),
+                mib(ticker(&stats, "rocksdb.compact.write.bytes")),
+                total_writes / elapsed.as_secs_f64(),
+            );
+        }
+    }
 }