Add snapshot and restore to indexer (#188)

Author: calummoore

indexer/src/adaptor.rs

@@ -1,5 +1,6 @@
 use crate::{where_query::WhereQuery, IndexerChange};
 use schema::{self, record::RecordRoot, Schema};
+use serde::{Deserialize, Serialize};
 use std::{pin::Pin, time::SystemTime};
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -21,6 +22,12 @@ pub enum Error {
     CollectionCollectionRecordNotFound { id: String },
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SnapshotValue {
+    pub key: Box<[u8]>,
+    pub value: Box<[u8]>,
+}
+
 /// The Store trait
 #[async_trait::async_trait]
 pub trait IndexerAdaptor: Send + Sync {
@@ -51,4 +58,13 @@ pub trait IndexerAdaptor: Send + Sync {
     async fn set_system_key(&self, key: &str, data: &RecordRoot) -> Result<()>;
 
     async fn get_system_key(&self, key: &str) -> Result<Option<RecordRoot>>;
+
+    async fn snapshot(
+        &self,
+        chunk_size: usize,
+    ) -> Pin<Box<dyn futures::Stream<Item = Result<Vec<SnapshotValue>>> + '_ + Send>>;
+
+    async fn restore(&self, chunk: Vec<SnapshotValue>) -> Result<()>;
+
+    async fn reset(&self) -> Result<()>;
 }

indexer/src/lib.rs

@@ -2,7 +2,7 @@
 
 // TODO: we should export schema from here, so that indexer builders
 // are using the correct schema
-use crate::adaptor::IndexerAdaptor;
+use crate::adaptor::{IndexerAdaptor, SnapshotValue};
 use crate::list_query::ListQuery;
 use crate::where_query::WhereQuery;
 use futures::stream::{FuturesUnordered, StreamExt};
@@ -75,6 +75,25 @@ impl<A: IndexerAdaptor> Indexer<A> {
         Self { adaptor }
     }
 
+    pub async fn snapshot(
+        &self,
+        chunk_size: usize,
+    ) -> Pin<Box<dyn futures::Stream<Item = Result<Vec<SnapshotValue>>> + '_ + Send>> {
+        self.adaptor
+            .snapshot(chunk_size)
+            .await
+            .map(|s| s.map_err(Error::from))
+            .boxed()
+    }
+
+    pub async fn restore(&self, chunk: Vec<SnapshotValue>) -> Result<()> {
+        Ok(self.adaptor.restore(chunk).await?)
+    }
+
+    pub async fn reset(&self) -> Result<()> {
+        Ok(self.adaptor.reset().await?)
+    }
+
     pub async fn commit(&self, height: usize, changes: Vec<IndexerChange>) -> Result<()> {
         Ok(self.adaptor.commit(height, changes).await?)
     }

indexer/src/memory.rs

@@ -1,5 +1,6 @@
-use crate::adaptor::{Error, IndexerAdaptor, Result};
+use crate::adaptor::{Error, Result, SnapshotValue};
 use crate::where_query::{WhereInequality, WhereNode, WhereQuery};
+use crate::IndexerAdaptor;
 use crate::IndexerChange;
 use schema::{
     field_path::FieldPath,
@@ -417,6 +418,21 @@ impl IndexerAdaptor for MemoryStore {
 
         Ok(state.system_data.get(key).cloned())
     }
+
+    async fn snapshot(
+        &self,
+        _: usize,
+    ) -> Pin<Box<dyn futures::Stream<Item = Result<Vec<SnapshotValue>>> + '_ + Send>> {
+        todo!()
+    }
+
+    async fn restore(&self, _: Vec<SnapshotValue>) -> Result<()> {
+        todo!()
+    }
+
+    async fn reset(&self) -> Result<()> {
+        todo!()
+    }
 }
 
 #[cfg(test)]

indexer_rocksdb/src/adaptor.rs

@@ -2,13 +2,13 @@ use crate::keys;
 use crate::result_stream::convert_stream;
 use crate::{
     key_range::{self, key_range, KeyRange},
-    proto,
+    proto, snapshot,
     store::{self, Store},
 };
 use async_recursion::async_recursion;
 use futures::{StreamExt, TryStreamExt};
 use indexer::{
-    adaptor::{self, IndexerAdaptor},
+    adaptor::{self, IndexerAdaptor, SnapshotValue},
     where_query::WhereQuery,
     IndexerChange,
 };
@@ -92,6 +92,14 @@ impl RocksDBAdaptor {
         }
     }
 
+    pub fn snapshot(&self, chunk_size: usize) -> snapshot::SnapshotIterator {
+        self.store.snapshot(chunk_size)
+    }
+
+    pub fn restore(&self, data: snapshot::SnapshotChunk) -> Result<()> {
+        Ok(self.store.restore(data)?)
+    }
+
     pub async fn _get(&self, collection_id: &str, record_id: &str) -> Result<Option<RecordRoot>> {
         let key = keys::Key::new_data(collection_id.to_string(), record_id.to_string())?;
 
@@ -613,6 +621,28 @@ impl IndexerAdaptor for RocksDBAdaptor {
     async fn get_system_key(&self, key: &str) -> adaptor::Result<Option<RecordRoot>> {
         Ok(self._get_system_record(key).await?)
     }
+
+    async fn snapshot(
+        &self,
+        chunk_size: usize,
+    ) -> Pin<Box<dyn futures::Stream<Item = adaptor::Result<Vec<SnapshotValue>>> + '_ + Send>> {
+        let res = futures::stream::iter(self.store.snapshot(chunk_size));
+        let stream = Box::pin(res.map(|s| {
+            s.map_err(store::StoreError::from)
+                .map_err(Error::from)
+                .map_err(adaptor::Error::from)
+        }));
+
+        stream.boxed()
+    }
+
+    async fn restore(&self, chunk: Vec<SnapshotValue>) -> adaptor::Result<()> {
+        Ok(self.store.restore(chunk).map_err(Error::from)?)
+    }
+
+    async fn reset(&self) -> adaptor::Result<()> {
+        Ok(self.store.reset().map_err(Error::from)?)
+    }
 }
 
 impl From<Error> for adaptor::Error {

indexer_rocksdb/src/snapshot.rs

@@ -1,5 +1,5 @@
+use indexer::adaptor::SnapshotValue;
 use rocksdb::{IteratorMode, DB};
-use serde::{Deserialize, Serialize};
 
 pub type Result<T> = std::result::Result<T, Error>;
 
@@ -12,16 +12,9 @@ pub enum Error {
     BincodeError(#[from] bincode::Error),
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct SnapshotValue {
-    pub(crate) key: Box<[u8]>,
-    pub(crate) value: Box<[u8]>,
-}
-
 pub type SnapshotChunk = Vec<SnapshotValue>;
 
 pub struct SnapshotIterator<'a> {
-    // db: &'a DB,
     chunk_size: usize,
     iter: rocksdb::DBIteratorWithThreadMode<'a, rocksdb::DB>,
 }

polybase/src/db.rs

@@ -3,9 +3,11 @@ use crate::mempool::Mempool;
 use crate::txn::{self, CallTxn};
 use futures_util::{future, StreamExt};
 use gateway::Gateway;
-use indexer::{adaptor::IndexerAdaptor, IndexerChange};
+use indexer::{
+    adaptor::{IndexerAdaptor, SnapshotValue},
+    IndexerChange,
+};
 use indexer::{auth_user::AuthUser, list_query::ListQuery, Indexer};
-use indexer_rocksdb::snapshot::{SnapshotChunk, SnapshotIterator};
 use parking_lot::Mutex;
 use schema::{
     self, methods,
@@ -15,9 +17,9 @@ use schema::{
     },
     Schema,
 };
-use serde::{Deserialize, Serialize};
 use solid::proposal::{self};
 use std::cmp::min;
+use std::pin::Pin;
 use std::time::{Duration, SystemTime};
 use tokio::sync::mpsc;
 use tokio::sync::Mutex as AsyncMutex;
@@ -108,11 +110,6 @@ pub enum UserError {
     UnauthorizedCall,
 }
 
-#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
-pub struct DbSnapshot {
-    pub index: Vec<u8>,
-}
-
 pub enum DbWaitResult<T> {
     Updated(T),
     NotModified,
@@ -652,22 +649,25 @@ impl<A: IndexerAdaptor> Db<A> {
     }
 
     /// Reset all data in the database
-    pub fn reset(&self) -> Result<()> {
-        todo!()
-        //Ok(self.indexer.reset()?)
+    pub async fn reset(&self) -> Result<()> {
+        Ok(self.indexer.reset().await?)
     }
 
     /// Create a snapshot iterator, that can be used to iterate over the
     /// entire database in chunks
-    pub fn snapshot_iter(&self, _chunk_size: usize) -> SnapshotIterator {
-        todo!()
-        //self.indexer.snapshot(chunk_size)
+    pub async fn snapshot_iter(
+        &self,
+        chunk_size: usize,
+    ) -> Pin<Box<dyn futures::Stream<Item = Result<Vec<SnapshotValue>>> + '_ + Send>> {
+        self.indexer
+            .snapshot(chunk_size)
+            .await
+            .map(|s| s.map_err(Error::from))
+            .boxed()
     }
 
-    pub fn restore_chunk(&self, _chunk: SnapshotChunk) -> Result<()> {
-        todo!()
-        //self.indexer.restore(chunk)?;
-        //Ok(())
+    pub async fn restore_chunk(&self, chunk: Vec<SnapshotValue>) -> Result<()> {
+        Ok(self.indexer.restore(chunk).await?)
     }
 
     #[tracing::instrument(skip(self))]

polybase/src/main.rs

@@ -290,7 +290,7 @@ async fn main() -> Result<()> {
         let mut last_commit = Instant::now();
 
         while !shutdown.load(Ordering::Relaxed) {
-            let network = Arc::clone(&network);
+            let network: Arc<Network> = Arc::clone(&network);
 
             tokio::select! {
                 // Db only produces CallTxn events, that should be propogated
@@ -378,7 +378,7 @@ async fn main() -> Result<()> {
 
                             // Reset the database
                             #[allow(clippy::expect_used)]
-                            db.reset().expect("Failed to reset database");
+                            db.reset().await.expect("Failed to reset database");
 
                             info!("Db reset ready for snapshot, sending accept");
 
@@ -401,8 +401,8 @@ async fn main() -> Result<()> {
                             // and this snapshot may take a while to complete
                             tokio::spawn(async move {
                                 // 100MB chunks
-                                let snapshot_iter = db.snapshot_iter(config.snapshot_chunk_size);
-                                for chunk in snapshot_iter {
+                                let mut snapshot_iter = db.snapshot_iter(config.snapshot_chunk_size).await;
+                                while let Some(chunk) = snapshot_iter.next().await {
                                     let peer_id = from_peer_id.clone();
                                     match chunk {
                                         Ok(chunk) => {
@@ -450,7 +450,7 @@ async fn main() -> Result<()> {
                             if let Some(chunk) = chunk {
                                 // We should panic if we are unable to restore
                                 #[allow(clippy::unwrap_used)]
-                                db.restore_chunk(chunk).unwrap();
+                                db.restore_chunk(chunk).await.unwrap();
                             } else {
                                 // We are finished, reset solid with the new proposal state from the snapshot
                                 #[allow(clippy::unwrap_used)]

polybase/src/network/events.rs

@@ -1,4 +1,4 @@
-use indexer_rocksdb::snapshot::SnapshotChunk;
+use indexer::adaptor::SnapshotValue;
 use serde::{Deserialize, Serialize};
 use solid::proposal::ProposalAccept;
 use solid::proposal::ProposalManifest;
@@ -35,7 +35,7 @@ pub enum NetworkEvent {
     /// A chunk of data for us to load into the database.
     SnapshotChunk {
         id: usize,
-        chunk: Option<SnapshotChunk>,
+        chunk: Option<Vec<SnapshotValue>>,
     },
 
     /// A transaction sent to another peer, which we should add to our Mempool