Add snapshot export functionality (#45)

.github/workflows/ci.yml

@@ -23,6 +23,7 @@ jobs:
         ]
     steps:
       - uses: actions/checkout@v4
+      - uses: arduino/setup-protoc@v2
       - uses: Swatinem/rust-cache@v2
       - name: Default on nightly Rust
         run: rustup default nightly

Cargo.toml

@@ -11,12 +11,14 @@ members = ["state-reconstruct-fetcher"]
 async-trait = "0.1.74"
 bincode = "1"
 blake2 = "0.10.6"
+bytes = "1.5"
 chrono = "0.4.31"
 clap = { version = "4.4.7", features = ["derive", "env"] }
 ethers = "1.0.2"
 eyre = "0.6.8"
 hex = "0.4.3"
 indexmap = { version = "2.0.2" }
+prost = "0.12"
 rocksdb = "0.21"
 serde = { version = "1.0.189", features = ["derive"] }
 serde_json = { version = "1.0.107", features = ["std"] }
@@ -27,3 +29,6 @@ tracing = "0.1.40"
 tracing-subscriber = "0.3.17"
 zksync_merkle_tree = { git = "https://github.com/matter-labs/zksync-era.git" }
 zkevm_opcode_defs = { git = "https://github.com/matter-labs/era-zkevm_opcode_defs.git" }
+
+[build-dependencies]
+prost-build = "0.12"

build.rs

@@ -0,0 +1,4 @@
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    prost_build::compile_protos(&["proto/snapshot.proto"], &["proto/"])?;
+    Ok(())
+}

proto/snapshot.proto

@@ -0,0 +1,23 @@
+syntax = "proto3";
+
+package protobuf;
+
+message SnapshotStorageLogsChunk {
+    repeated SnapshotStorageLog storage_logs = 1;
+}
+
+message SnapshotStorageLog {
+    optional bytes account_address = 1; // required; H160
+    optional bytes storage_key = 2; // required; H256
+    optional bytes storage_value = 3; // required; H256
+    optional uint32 l1_batch_number_of_initial_write = 4; // required
+    optional uint64 enumeration_index = 5; // required
+}
+
+message SnapshotFactoryDependencies {
+    repeated SnapshotFactoryDependency factory_deps = 1;
+}
+
+message SnapshotFactoryDependency {
+    optional bytes bytecode = 1; // required
+}

src/cli.rs

@@ -1,6 +1,8 @@
 use clap::{Args, Parser, Subcommand, ValueEnum};
 use state_reconstruct_fetcher::constants::ethereum;
 
+use crate::processor::snapshot;
+
 #[derive(Args)]
 pub struct L1FetcherOptions {
     /// The Ethereum JSON-RPC HTTP URL to use.
@@ -73,12 +75,22 @@ pub enum Command {
         db_path: Option<String>,
     },
 
-    /// Testing.
-    ExportSnapshot {
+    PrepareSnapshot {
         #[command(flatten)]
         l1_fetcher_options: L1FetcherOptions,
-        /// The path of the file to export the snapshot to.
-        file: Option<String>,
+        /// The path to the storage solution.
+        #[arg(short, long)]
+        db_path: Option<String>,
+    },
+    ExportSnapshot {
+        /// The path to the storage solution.
+        #[arg(short, long, default_value = snapshot::DEFAULT_DB_PATH)]
+        db_path: Option<String>,
+        /// Number of storage logs to stuff into one chunk.
+        #[arg(short, long, default_value_t = 1_000_000)]
+        chunk_size: u64,
+        /// The directory to export the snapshot files to.
+        directory: String,
     },
 }
 

src/main.rs

@@ -16,7 +16,7 @@ use std::{
 use clap::Parser;
 use cli::{Cli, Command, ReconstructSource};
 use eyre::Result;
-use processor::snapshot::SnapshotBuilder;
+use processor::snapshot::{SnapshotBuilder, SnapshotExporter};
 use state_reconstruct_fetcher::{
     constants::storage,
     l1_fetcher::{L1Fetcher, L1FetcherOptions},
@@ -152,9 +152,9 @@ async fn main() -> Result<()> {
                 println!("{result}");
             }
         }
-        Command::ExportSnapshot {
+        Command::PrepareSnapshot {
             l1_fetcher_options,
-            file,
+            db_path,
         } => {
             let fetcher_options = L1FetcherOptions {
                 http_url: l1_fetcher_options.http_url,
@@ -165,7 +165,7 @@ async fn main() -> Result<()> {
             };
 
             let fetcher = L1Fetcher::new(fetcher_options, None)?;
-            let processor = SnapshotBuilder::new(file);
+            let processor = SnapshotBuilder::new(db_path);
 
             let (tx, rx) = mpsc::channel::<CommitBlockInfoV1>(5);
             let processor_handle = tokio::spawn(async move {
@@ -175,6 +175,16 @@ async fn main() -> Result<()> {
             fetcher.run(tx).await?;
             processor_handle.await?;
         }
+        Command::ExportSnapshot {
+            db_path,
+            chunk_size,
+            directory,
+        } => {
+            let export_path = Path::new(&directory);
+            std::fs::create_dir_all(export_path)?;
+            let exporter = SnapshotExporter::new(export_path, db_path);
+            exporter.export_snapshot(chunk_size)?;
+        }
     }
 
     Ok(())

src/processor/snapshot/database.rs

@@ -11,9 +11,9 @@ use thiserror::Error;
 
 use super::types::{SnapshotFactoryDependency, SnapshotStorageLog};
 
-const STORAGE_LOGS: &str = "storage_logs";
-const INDEX_TO_KEY_MAP: &str = "index_to_key_map";
-const FACTORY_DEPS: &str = "factory_deps";
+pub const STORAGE_LOGS: &str = "storage_logs";
+pub const INDEX_TO_KEY_MAP: &str = "index_to_key_map";
+pub const FACTORY_DEPS: &str = "factory_deps";
 const METADATA: &str = "metadata";
 
 const LAST_REPEATED_KEY_INDEX: &str = "LAST_REPEATED_KEY_INDEX";
@@ -50,6 +50,21 @@ impl SnapshotDB {
         Ok(Self(db))
     }
 
+    pub fn new_read_only(db_path: PathBuf) -> Result<Self> {
+        let mut db_opts = Options::default();
+        db_opts.create_missing_column_families(true);
+        db_opts.create_if_missing(true);
+
+        let db = DB::open_cf_for_read_only(
+            &db_opts,
+            db_path,
+            vec![METADATA, STORAGE_LOGS, INDEX_TO_KEY_MAP, FACTORY_DEPS],
+            false,
+        )?;
+
+        Ok(Self(db))
+    }
+
     pub fn get_last_repeated_key_index(&self) -> Result<u64> {
         // Unwrapping column family handle here is safe because presence of
         // those CFs is ensured in construction of this DB.
@@ -90,7 +105,7 @@ impl SnapshotDB {
             .map_err(Into::into)
     }
 
-    pub fn insert_storage_log(&self, storage_log_entry: &SnapshotStorageLog) -> Result<()> {
+    pub fn insert_storage_log(&self, storage_log_entry: &mut SnapshotStorageLog) -> Result<()> {
         // Unwrapping column family handle here is safe because presence of
         // those CFs is ensured in construction of this DB.
         let index_to_key_map = self.cf_handle(INDEX_TO_KEY_MAP).unwrap();
@@ -102,6 +117,9 @@ impl SnapshotDB {
         // XXX: These should really be inside a transaction...
         let idx = self.get_last_repeated_key_index()? + 1;
 
+        // Update the enumeration index.
+        storage_log_entry.enumeration_index = idx;
+
         self.put_cf(index_to_key_map, idx.to_be_bytes(), key)?;
         self.set_last_repeated_key_index(idx)?;