Make sure scanner database is accessed using the correct types (#8112)

* impl TryFrom<zcash_primitives::BlockHeight> for Height

* Add type-safe read and write database methods

* Only allow typed access to the scanner DB

* Update docs

* Implement a common method as a trait

* Fix imports

* Tidy state imports

* Activate tracing logging macros in the whole scanner crate

* Fix dead code warnings

zebra-chain/src/block/height.rs

@@ -112,6 +112,15 @@ impl From<Height> for BlockHeight {
     }
 }
 
+impl TryFrom<BlockHeight> for Height {
+    type Error = &'static str;
+
+    /// Checks that the `height` is within the valid [`Height`] range.
+    fn try_from(height: BlockHeight) -> Result<Self, Self::Error> {
+        Self::try_from(u32::from(height))
+    }
+}
+
 /// A difference between two [`Height`]s, possibly negative.
 ///
 /// This can represent the difference between any height values,

zebra-scan/src/lib.rs

@@ -4,6 +4,9 @@
 #![doc(html_logo_url = "https://zfnd.org/wp-content/uploads/2022/03/zebra-icon.png")]
 #![doc(html_root_url = "https://docs.rs/zebra_scan")]
 
+#[macro_use]
+extern crate tracing;
+
 pub mod config;
 pub mod init;
 pub mod scan;

zebra-scan/src/scan.rs

@@ -9,7 +9,6 @@ use std::{
 use color_eyre::{eyre::eyre, Report};
 use itertools::Itertools;
 use tower::{buffer::Buffer, util::BoxService, Service, ServiceExt};
-use tracing::info;
 
 use zcash_client_backend::{
     data_api::ScannedBlock,

zebra-scan/src/storage.rs

@@ -6,9 +6,7 @@ use zebra_chain::{
     block::Height,
     parameters::{Network, NetworkUpgrade},
 };
-use zebra_state::{
-    SaplingScannedDatabaseEntry, SaplingScannedDatabaseIndex, TransactionIndex, TransactionLocation,
-};
+use zebra_state::TransactionIndex;
 
 use crate::config::Config;
 
@@ -17,11 +15,9 @@ pub mod db;
 // Public types and APIs
 pub use db::{SaplingScannedResult, SaplingScanningKey};
 
-use self::db::ScannerWriteBatch;
-
 /// We insert an empty results entry to the database every this interval for each stored key,
 /// so we can track progress.
-const INSERT_CONTROL_INTERVAL: u32 = 1_000;
+pub const INSERT_CONTROL_INTERVAL: u32 = 1_000;
 
 /// Store key info and results of the scan.
 ///
@@ -84,11 +80,7 @@ impl Storage {
 
         // It's ok to write some keys and not others during shutdown, so each key can get its own
         // batch. (They will be re-written on startup anyway.)
-        let mut batch = ScannerWriteBatch::default();
-
-        batch.insert_sapling_key(self, sapling_key, birthday);
-
-        self.write_batch(batch);
+        self.insert_sapling_key(sapling_key, birthday);
     }
 
     /// Returns all the keys and their last scanned heights.
@@ -104,6 +96,9 @@ impl Storage {
     /// Add the sapling results for `height` to the storage. The results can be any map of
     /// [`TransactionIndex`] to [`SaplingScannedResult`].
     ///
+    /// All the results for the same height must be written at the same time, to avoid partial
+    /// writes during shutdown.
+    ///
     /// Also adds empty progress tracking entries every `INSERT_CONTROL_INTERVAL` blocks if needed.
     ///
     /// # Performance / Hangs
@@ -116,37 +111,7 @@ impl Storage {
         height: Height,
         sapling_results: BTreeMap<TransactionIndex, SaplingScannedResult>,
     ) {
-        // We skip heights that have one or more results, so the results for each height must be
-        // in a single batch.
-        let mut batch = ScannerWriteBatch::default();
-
-        // Every `INSERT_CONTROL_INTERVAL` we add a new entry to the scanner database for each key
-        // so we can track progress made in the last interval even if no transaction was yet found.
-        let needs_control_entry =
-            height.0 % INSERT_CONTROL_INTERVAL == 0 && sapling_results.is_empty();
-
-        // Add scanner progress tracking entry for key.
-        // Defensive programming: add the tracking entry first, so that we don't accidentally
-        // overwrite real results with it. (This is currently prevented by the empty check.)
-        if needs_control_entry {
-            batch.insert_sapling_height(self, sapling_key, height);
-        }
-
-        for (index, sapling_result) in sapling_results {
-            let index = SaplingScannedDatabaseIndex {
-                sapling_key: sapling_key.clone(),
-                tx_loc: TransactionLocation::from_parts(height, index),
-            };
-
-            let entry = SaplingScannedDatabaseEntry {
-                index,
-                value: Some(sapling_result),
-            };
-
-            batch.insert_sapling_result(self, entry);
-        }
-
-        self.write_batch(batch);
+        self.insert_sapling_results(sapling_key, height, sapling_results)
     }
 
     /// Returns all the results for a sapling key, for every scanned block height.

zebra-scan/src/storage/db.rs

@@ -5,7 +5,6 @@ use std::path::Path;
 use semver::Version;
 
 use zebra_chain::parameters::Network;
-use zebra_state::{DiskWriteBatch, ReadDisk};
 
 use crate::Config;
 
@@ -86,15 +85,15 @@ impl Storage {
         // Report where we are for each key in the database.
         let keys = new_storage.sapling_keys_last_heights();
         for (key_num, (_key, height)) in keys.iter().enumerate() {
-            tracing::info!(
+            info!(
                 "Last scanned height for key number {} is {}, resuming at {}",
                 key_num,
                 height.as_usize(),
                 height.next().expect("height is not maximum").as_usize(),
             );
         }
 
-        tracing::info!("loaded Zebra scanner cache");
+        info!("loaded Zebra scanner cache");
 
         new_storage
     }
@@ -134,28 +133,6 @@ impl Storage {
     /// Returns true if the database is empty.
     pub fn is_empty(&self) -> bool {
         // Any column family that is populated at (or near) startup can be used here.
-        self.db.zs_is_empty(&self.sapling_tx_ids_cf())
-    }
-}
-
-// General writing
-
-/// Wrapper type for scanner database writes.
-#[must_use = "batches must be written to the database"]
-#[derive(Default)]
-pub struct ScannerWriteBatch(pub DiskWriteBatch);
-
-// Redirect method calls to DiskWriteBatch for convenience.
-impl std::ops::Deref for ScannerWriteBatch {
-    type Target = DiskWriteBatch;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-impl std::ops::DerefMut for ScannerWriteBatch {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
+        self.sapling_tx_ids_cf().zs_is_empty()
     }
 }