Recommitments refactoring (#589)

nazar-pc · web-flow · commit f49009fa87b6 · 2022-06-13T12:34:25.000+03:00
* Simplify iteration over commitments databases while avoiding some locks

* Simplify commitments metadata mutation under new data structure and move it into a separate module

* Rename `commitment_databases` module to `databases`

* Minor simplification
diff --git a/crates/subspace-farmer/src/commitments.rs b/crates/subspace-farmer/src/commitments.rs
@@ -1,10 +1,11 @@
-mod commitment_databases;
+mod databases;
+mod metadata;
 #[cfg(test)]
 mod tests;
 
 use crate::plot::{PieceOffset, Plot};
 use arc_swap::ArcSwapOption;
-use commitment_databases::{CommitmentDatabases, CreateDbEntryResult, DbEntry};
+use databases::{CommitmentDatabases, CreateDbEntryResult, DbEntry};
 use event_listener_primitives::{Bag, HandlerId};
 use parking_lot::Mutex;
 use rayon::prelude::*;
@@ -179,22 +180,8 @@ impl Commitments {
     }
 
     pub(crate) fn remove_pieces(&self, pieces: &[Piece]) -> Result<(), CommitmentError> {
-        let salts = self.inner.commitment_databases.lock().get_salts();
-
-        for salt in salts {
-            let db_entry = match self
-                .inner
-                .commitment_databases
-                .lock()
-                .get_db_entry(&salt)
-                .cloned()
-            {
-                Some(db_entry) => db_entry,
-                None => {
-                    continue;
-                }
-            };
-
+        for db_entry in self.get_db_entries() {
+            let salt = db_entry.salt();
             let db_guard = db_entry.lock();
 
             if let Some(db) = db_guard.as_ref() {
@@ -217,22 +204,8 @@ impl Commitments {
         F: Fn() -> Iter,
         Iter: Iterator<Item = (PieceOffset, &'iter [u8])>,
     {
-        let salts = self.inner.commitment_databases.lock().get_salts();
-
-        for salt in salts {
-            let db_entry = match self
-                .inner
-                .commitment_databases
-                .lock()
-                .get_db_entry(&salt)
-                .cloned()
-            {
-                Some(db_entry) => db_entry,
-                None => {
-                    continue;
-                }
-            };
-
+        for db_entry in self.get_db_entries() {
+            let salt = db_entry.salt();
             let db_guard = db_entry.lock();
 
             if let Some(db) = db_guard.as_ref() {
@@ -257,7 +230,7 @@ impl Commitments {
         range: u64,
         salt: Salt,
     ) -> Option<(Tag, PieceOffset)> {
-        let db_entry = self.get_local_db_entry(&salt)?;
+        let db_entry = self.get_db_entry(salt)?;
 
         let db_guard = db_entry.try_lock()?;
         let db = db_guard.clone()?;
@@ -282,7 +255,7 @@ impl Commitments {
         self.inner.handlers.status_change.add(callback)
     }
 
-    fn get_local_db_entry(&self, salt: &Salt) -> Option<Arc<DbEntry>> {
+    fn get_db_entry(&self, salt: Salt) -> Option<Arc<DbEntry>> {
         if let Some(current) = self.inner.current.load_full() {
             if current.salt() == salt {
                 return Some(current);
@@ -297,6 +270,14 @@ impl Commitments {
 
         None
     }
+
+    fn get_db_entries(&self) -> impl Iterator<Item = Arc<DbEntry>> {
+        self.inner
+            .current
+            .load_full()
+            .into_iter()
+            .chain(self.inner.next.load_full())
+    }
 }
 
 enum SolutionIteratorState {
diff --git a/crates/subspace-farmer/src/commitments/databases.rs b/crates/subspace-farmer/src/commitments/databases.rs
@@ -1,9 +1,8 @@
-use super::CommitmentError;
+use crate::commitments::metadata::{CommitmentMetadata, CommitmentStatus};
+use crate::commitments::CommitmentError;
 use lru::LruCache;
 use parking_lot::Mutex;
 use rocksdb::{Options, DB};
-use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
 use std::fmt;
 use std::ops::Deref;
 use std::path::PathBuf;
@@ -13,15 +12,6 @@ use tracing::error;
 
 // Cache size is just enough for last 2 salts to be stored
 const COMMITMENTS_CACHE_SIZE: usize = 2;
-const COMMITMENTS_KEY: &[u8] = b"commitments";
-
-#[derive(Debug, Copy, Clone, Eq, PartialEq, Serialize, Deserialize)]
-enum CommitmentStatus {
-    /// In-progress commitment to the part of the plot
-    InProgress,
-    /// Commitment to the whole plot and not some in-progress partial commitment
-    Created,
-}
 
 pub(super) struct CreateDbEntryResult {
     pub(super) db_entry: Arc<DbEntry>,
@@ -48,45 +38,25 @@ impl Deref for DbEntry {
 }
 
 impl DbEntry {
-    pub(super) fn salt(&self) -> &Salt {
-        &self.salt
+    pub(super) fn salt(&self) -> Salt {
+        self.salt
     }
 }
 
 #[derive(Debug)]
 pub(super) struct CommitmentDatabases {
     base_directory: PathBuf,
     databases: LruCache<Salt, Arc<DbEntry>>,
-    metadata_cache: HashMap<Salt, CommitmentStatus>,
-    metadata_db: Arc<DB>,
+    metadata: Mutex<CommitmentMetadata>,
 }
 
 impl CommitmentDatabases {
     pub(super) fn new(base_directory: PathBuf) -> Result<Self, CommitmentError> {
-        let metadata_db = DB::open_default(base_directory.join("metadata"))
-            .map_err(CommitmentError::MetadataDb)?;
-        let metadata_cache: HashMap<Salt, CommitmentStatus> = metadata_db
-            .get(COMMITMENTS_KEY)
-            .map_err(CommitmentError::MetadataDb)?
-            .map(|bytes| {
-                serde_json::from_slice::<HashMap<String, CommitmentStatus>>(&bytes)
-                    .unwrap()
-                    .into_iter()
-                    .map(|(salt, status)| (hex::decode(salt).unwrap().try_into().unwrap(), status))
-                    .collect()
-            })
-            .unwrap_or_default();
-
-        let mut commitment_databases = CommitmentDatabases {
-            base_directory: base_directory.clone(),
-            databases: LruCache::new(COMMITMENTS_CACHE_SIZE),
-            metadata_cache,
-            metadata_db: Arc::new(metadata_db),
-        };
+        let mut metadata = CommitmentMetadata::new(base_directory.join("metadata"))?;
+        let mut databases = LruCache::new(COMMITMENTS_CACHE_SIZE);
 
-        if commitment_databases
-            .metadata_cache
-            .drain_filter(|salt, status| match status {
+        metadata.mutate(|metadata| {
+            metadata.drain_filter(|salt, status| match status {
                 CommitmentStatus::InProgress => {
                     if let Err(error) =
                         std::fs::remove_dir_all(base_directory.join(hex::encode(salt)))
@@ -100,35 +70,29 @@ impl CommitmentDatabases {
                     true
                 }
                 CommitmentStatus::Created => false,
-            })
-            .next()
-            .is_some()
-        {
-            commitment_databases.persist_metadata_cache()?;
-        }
-
-        // Open databases that were fully created during previous run
-        for salt in commitment_databases.metadata_cache.keys() {
-            let db = DB::open(&Options::default(), base_directory.join(hex::encode(salt)))
-                .map_err(CommitmentError::CommitmentDb)?;
-            commitment_databases.databases.put(
-                *salt,
-                Arc::new(DbEntry {
-                    salt: *salt,
-                    db: Mutex::new(Some(Arc::new(db))),
-                }),
-            );
-        }
-
-        Ok::<_, CommitmentError>(commitment_databases)
-    }
+            });
 
-    /// Get salts for all current database entries
-    pub(super) fn get_salts(&self) -> Vec<Salt> {
-        self.databases
-            .iter()
-            .map(|(salt, _db_entry)| *salt)
-            .collect()
+            // Open databases that were fully created during previous run
+            for salt in metadata.keys() {
+                let db = DB::open(&Options::default(), base_directory.join(hex::encode(salt)))
+                    .map_err(CommitmentError::CommitmentDb)?;
+                databases.put(
+                    *salt,
+                    Arc::new(DbEntry {
+                        salt: *salt,
+                        db: Mutex::new(Some(Arc::new(db))),
+                    }),
+                );
+            }
+
+            Ok(())
+        })?;
+
+        Ok(CommitmentDatabases {
+            base_directory: base_directory.clone(),
+            databases,
+            metadata: Mutex::new(metadata),
+        })
     }
 
     /// Returns current and next `db_entry`.
@@ -176,7 +140,11 @@ impl CommitmentDatabases {
             let old_db_path = self.base_directory.join(hex::encode(old_salt));
 
             // Remove old commitments for `old_salt`
-            self.metadata_cache.remove(&old_salt);
+            self.metadata.lock().mutate(|metadata| {
+                metadata.remove(&old_salt);
+
+                Ok(())
+            })?;
 
             tokio::task::spawn_blocking(move || {
                 // Take a lock to make sure database was released by whatever user there was and we
@@ -202,35 +170,16 @@ impl CommitmentDatabases {
     }
 
     pub(super) fn mark_in_progress(&mut self, salt: Salt) -> Result<(), CommitmentError> {
-        self.update_status(salt, CommitmentStatus::InProgress)
+        self.metadata.lock().mutate(|metadata| {
+            metadata.insert(salt, CommitmentStatus::InProgress);
+            Ok(())
+        })
     }
 
     pub(super) fn mark_created(&mut self, salt: Salt) -> Result<(), CommitmentError> {
-        self.update_status(salt, CommitmentStatus::Created)
-    }
-
-    fn update_status(
-        &mut self,
-        salt: Salt,
-        status: CommitmentStatus,
-    ) -> Result<(), CommitmentError> {
-        self.metadata_cache.insert(salt, status);
-
-        self.persist_metadata_cache()
-    }
-
-    fn persist_metadata_cache(&self) -> Result<(), CommitmentError> {
-        let prepared_metadata_cache: HashMap<String, CommitmentStatus> = self
-            .metadata_cache
-            .iter()
-            .map(|(salt, status)| (hex::encode(salt), *status))
-            .collect();
-
-        self.metadata_db
-            .put(
-                COMMITMENTS_KEY,
-                &serde_json::to_vec(&prepared_metadata_cache).unwrap(),
-            )
-            .map_err(CommitmentError::MetadataDb)
+        self.metadata.lock().mutate(|metadata| {
+            metadata.insert(salt, CommitmentStatus::Created);
+            Ok(())
+        })
     }
 }
diff --git a/crates/subspace-farmer/src/commitments/metadata.rs b/crates/subspace-farmer/src/commitments/metadata.rs
@@ -0,0 +1,58 @@
+use crate::commitments::CommitmentError;
+use rocksdb::DB;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::path::PathBuf;
+use subspace_core_primitives::Salt;
+
+const COMMITMENTS_KEY: &[u8] = b"commitments";
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Serialize, Deserialize)]
+pub(super) enum CommitmentStatus {
+    /// In-progress commitment to the part of the plot
+    InProgress,
+    /// Commitment to the whole plot and not some in-progress partial commitment
+    Created,
+}
+
+#[derive(Debug)]
+pub(super) struct CommitmentMetadata {
+    db: DB,
+}
+
+impl CommitmentMetadata {
+    pub(super) fn new(path: PathBuf) -> Result<Self, CommitmentError> {
+        let db = DB::open_default(path).map_err(CommitmentError::MetadataDb)?;
+
+        Ok(Self { db })
+    }
+
+    pub(super) fn mutate<F>(&mut self, mut callback: F) -> Result<(), CommitmentError>
+    where
+        F: FnMut(&mut HashMap<Salt, CommitmentStatus>) -> Result<(), CommitmentError>,
+    {
+        let mut metadata = self
+            .db
+            .get(COMMITMENTS_KEY)
+            .map_err(CommitmentError::MetadataDb)?
+            .map(|bytes| {
+                serde_json::from_slice::<HashMap<String, CommitmentStatus>>(&bytes)
+                    .unwrap()
+                    .into_iter()
+                    .map(|(salt, status)| (hex::decode(salt).unwrap().try_into().unwrap(), status))
+                    .collect()
+            })
+            .unwrap_or_default();
+
+        callback(&mut metadata)?;
+
+        let metadata: HashMap<String, CommitmentStatus> = metadata
+            .iter()
+            .map(|(salt, status)| (hex::encode(salt), *status))
+            .collect();
+
+        self.db
+            .put(COMMITMENTS_KEY, &serde_json::to_vec(&metadata).unwrap())
+            .map_err(CommitmentError::MetadataDb)
+    }
+}
