fix: add config option for clearing address db on start

Signed-off-by: William Hankins <[email protected]>

Author: whankinsiv

modules/address_state/src/address_state.rs

@@ -33,6 +33,7 @@ const DEFAULT_PARAMETERS_SUBSCRIBE_TOPIC: (&str, &str) =
 
 // Configuration defaults
 const DEFAULT_ADDRESS_DB_PATH: (&str, &str) = ("db-path", "./db");
+const DEFAULT_CLEAR_ON_START: (&str, bool) = ("clear-on-start", true);
 const DEFAULT_STORE_INFO: (&str, bool) = ("store-info", false);
 const DEFAULT_STORE_TOTALS: (&str, bool) = ("store-totals", false);
 const DEFAULT_STORE_TRANSACTIONS: (&str, bool) = ("store-transactions", false);
@@ -178,6 +179,7 @@ impl AddressState {
         // Get configuration flags and query topic
         let storage_config = AddressStorageConfig {
             db_path: get_string_flag(&config, DEFAULT_ADDRESS_DB_PATH),
+            clear_on_start: get_bool_flag(&config, DEFAULT_CLEAR_ON_START),
             skip_until: None,
             store_info: get_bool_flag(&config, DEFAULT_STORE_INFO),
             store_totals: get_bool_flag(&config, DEFAULT_STORE_TOTALS),

modules/address_state/src/immutable_address_store.rs

@@ -30,9 +30,9 @@ pub struct ImmutableAddressStore {
 }
 
 impl ImmutableAddressStore {
-    pub fn new(path: impl AsRef<Path>) -> Result<Self> {
+    pub fn new(path: impl AsRef<Path>, clear_on_start: bool) -> Result<Self> {
         let path = path.as_ref();
-        if path.exists() {
+        if path.exists() && clear_on_start {
             std::fs::remove_dir_all(path)?;
         }
         let cfg = fjall::Config::new(path).max_write_buffer_size(512 * 1024 * 1024);
@@ -56,15 +56,36 @@ impl ImmutableAddressStore {
     /// for an entire epoch. Skips any partitions that have already stored the given epoch.
     /// All writes are batched and committed atomically, preventing on-disk corruption in case of failure.
     pub async fn persist_epoch(&self, epoch: u64, config: &AddressStorageConfig) -> Result<()> {
-        let persist_utxos = config.store_info
-            && !self.epoch_exists(self.utxos.clone(), ADDRESS_UTXOS_EPOCH_COUNTER, epoch).await?;
-        let persist_txs = config.store_transactions
-            && !self.epoch_exists(self.txs.clone(), ADDRESS_TXS_EPOCH_COUNTER, epoch).await?;
-        let persist_totals = config.store_totals
-            && !self.epoch_exists(self.totals.clone(), ADDRESS_TOTALS_EPOCH_COUNTER, epoch).await?;
+        // Skip if all options disabled
+        if !(config.store_info || config.store_transactions || config.store_totals) {
+            debug!("no persistence needed for epoch {epoch} (all stores disabled)");
+            return Ok(());
+        }
+
+        // Determine which partitions need persistence
+        let (persist_utxos, persist_txs, persist_totals) = if config.clear_on_start {
+            (
+                config.store_info,
+                config.store_transactions,
+                config.store_totals,
+            )
+        } else {
+            let utxos = config.store_info
+                && !self
+                    .epoch_exists(self.utxos.clone(), ADDRESS_UTXOS_EPOCH_COUNTER, epoch)
+                    .await?;
+            let txs = config.store_transactions
+                && !self.epoch_exists(self.txs.clone(), ADDRESS_TXS_EPOCH_COUNTER, epoch).await?;
+            let totals = config.store_totals
+                && !self
+                    .epoch_exists(self.totals.clone(), ADDRESS_TOTALS_EPOCH_COUNTER, epoch)
+                    .await?;
+            (utxos, txs, totals)
+        };
 
+        // Skip if all partitions have already been persisted for the epoch
         if !(persist_utxos || persist_txs || persist_totals) {
-            debug!("no persistence needed for epoch {epoch} (already persisted or disabled)");
+            debug!("no persistence needed for epoch {epoch}");
             return Ok(());
         }
 
@@ -124,22 +145,14 @@ impl ImmutableAddressStore {
         }
 
         // Metadata markers
-        if persist_utxos {
-            batch.insert(
-                &self.utxos,
-                ADDRESS_UTXOS_EPOCH_COUNTER,
-                epoch.to_le_bytes(),
-            );
-        }
-        if persist_txs {
-            batch.insert(&self.txs, ADDRESS_TXS_EPOCH_COUNTER, epoch.to_le_bytes());
-        }
-        if persist_totals {
-            batch.insert(
-                &self.totals,
-                ADDRESS_TOTALS_EPOCH_COUNTER,
-                epoch.to_le_bytes(),
-            );
+        for (enabled, part, key) in [
+            (persist_utxos, &self.utxos, ADDRESS_UTXOS_EPOCH_COUNTER),
+            (persist_txs, &self.txs, ADDRESS_TXS_EPOCH_COUNTER),
+            (persist_totals, &self.totals, ADDRESS_TOTALS_EPOCH_COUNTER),
+        ] {
+            if enabled {
+                batch.insert(part, key, epoch.to_le_bytes());
+            }
         }
 
         match batch.commit() {

modules/address_state/src/state.rs

@@ -17,6 +17,7 @@ use crate::{
 #[derive(Debug, Default, Clone)]
 pub struct AddressStorageConfig {
     pub db_path: String,
+    pub clear_on_start: bool,
     pub skip_until: Option<u64>,
 
     pub store_info: bool,
@@ -60,7 +61,7 @@ impl State {
             PathBuf::from(&config.db_path)
         };
 
-        let store = Arc::new(ImmutableAddressStore::new(&db_path)?);
+        let store = Arc::new(ImmutableAddressStore::new(&db_path, config.clear_on_start)?);
 
         let mut config = config.clone();
         config.skip_until = store.get_last_epoch_stored().await?;
@@ -252,6 +253,7 @@ mod tests {
         let dir = tempdir().unwrap();
         AddressStorageConfig {
             db_path: dir.path().to_string_lossy().into_owned(),
+            clear_on_start: true,
             skip_until: None,
             store_info: true,
             store_transactions: true,

processes/omnibus/omnibus.toml

@@ -138,6 +138,8 @@ index-by-policy = false
 clear-on-start = true
 
 [module.address-state]
+# Clear state on start up (default true)
+clear-on-start = true
 # Enables /addresses/{address}, /addresses/{address}/extended,
 # /addresses/{address}/utxos/{asset}, and /addresses/{address}/utxos endpoints
 store-info = false