From b1bbb13978c6d4a337322e4a88b716a0b9b26064 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Wed, 25 Nov 2020 15:04:18 +1000
Subject: [PATCH] Make debug_stop_at_height and ephemeral work together (#1339)

* Make debug_stop_at_height and ephemeral work together

* if `debug_stop_at_height` and `ephemeral` are set, delete the database
  files after reaching the stop height
* drop or flush the database before `debug_stop_at_height` exits Zebra
---
 zebra-state/src/service/finalized_state.rs | 37 ++++++++++++++++++----
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/zebra-state/src/service/finalized_state.rs b/zebra-state/src/service/finalized_state.rs
index 73298244..ba93f9c0 100644
--- a/zebra-state/src/service/finalized_state.rs
+++ b/zebra-state/src/service/finalized_state.rs
@@ -65,7 +65,9 @@ impl FinalizedState {
                     "state is already at the configured height"
                 );
 
-                // There's no need to sync before exit, because the trees have just been opened
+                // RocksDB can do a cleanup when column families are opened.
+                // So we want to drop it before we exit.
+                std::mem::drop(new_state);
                 std::process::exit(0);
             }
         }
@@ -253,7 +255,11 @@ impl FinalizedState {
 
         if result.is_ok() && self.is_at_stop_height(height) {
             tracing::info!(?height, ?hash, "stopping at configured height");
-
+            // We'd like to drop the database here, because that closes the
+            // column families and the database. But Rust's ownership rules
+            // make that difficult, so we just flush instead.
+            self.db.flush().expect("flush is successful");
+            self.delete_ephemeral();
             std::process::exit(0);
         }
 
@@ -327,6 +333,27 @@ impl FinalizedState {
                 block.transactions[index as usize].clone()
             })
     }
+
+    /// If the database is `ephemeral`, delete it.
+    fn delete_ephemeral(&self) {
+        if self.ephemeral {
+            let path = self.db.path();
+            tracing::debug!("removing temporary database files {:?}", path);
+            // We'd like to use `rocksdb::Env::mem_env` for ephemeral databases,
+            // but the Zcash blockchain might not fit in memory. So we just
+            // delete the database files instead.
+            //
+            // We'd like to call `DB::destroy` here, but calling destroy on a
+            // live DB is undefined behaviour:
+            // https://github.com/facebook/rocksdb/wiki/RocksDB-FAQ#basic-readwrite
+            //
+            // So we assume that all the database files are under `path`, and
+            // delete them using standard filesystem APIs. Deleting open files
+            // might cause errors on non-Unix platforms, so we ignore the result.
+            // (The OS will delete them eventually anyway.)
+            let _res = std::fs::remove_dir_all(path);
+        }
+    }
 }
 
 // Drop isn't guaranteed to run, such as when we panic, or if someone stored
@@ -335,11 +362,7 @@ impl FinalizedState {
 // up automatically eventually.
 impl Drop for FinalizedState {
     fn drop(&mut self) {
-        if self.ephemeral {
-            let path = self.db.path();
-            tracing::debug!("removing temporary database files {:?}", path);
-            let _res = std::fs::remove_dir_all(path);
-        }
+        self.delete_ephemeral()
     }
 }