feat(state-keeper): implement asynchronous RocksDB cache

checks-config/era.dic

@@ -119,6 +119,7 @@ stdout
 GCS
 websocket
 struct
+struct's
 localhost
 TOML
 config

core/lib/state/src/rocksdb/mod.rs

@@ -599,7 +599,7 @@ impl RocksdbStorage {
     /// # Panics
     ///
     /// Panics on RocksDB errors.
-    async fn l1_batch_number(&self) -> Option<L1BatchNumber> {
+    pub async fn l1_batch_number(&self) -> Option<L1BatchNumber> {
         let cf = StateKeeperColumnFamily::State;
         let db = self.db.clone();
         let number_bytes =

core/lib/zksync_core/src/state_keeper/batch_executor/main_executor.rs

@@ -11,16 +11,20 @@ use multivm::{
     MultiVMTracer, VmInstance,
 };
 use once_cell::sync::OnceCell;
-use tokio::sync::{mpsc, watch};
+use tokio::{
+    runtime::Handle,
+    sync::{mpsc, watch, RwLock},
+};
 use zksync_dal::ConnectionPool;
-use zksync_state::{RocksdbStorage, StorageView, WriteStorage};
+use zksync_state::{ReadStorage, StorageView, WriteStorage};
 use zksync_types::{vm_trace::Call, Transaction, U256};
 use zksync_utils::bytecode::CompressedBytecodeInfo;
 
 use super::{BatchExecutor, BatchExecutorHandle, Command, TxExecutionResult};
 use crate::{
     metrics::{InteractionType, TxStage, APP_METRICS},
     state_keeper::{
+        cached_storage::CachedStorage,
         metrics::{TxExecutionStage, EXECUTOR_METRICS, KEEPER_METRICS},
         types::ExecutionMetricsForCriteria,
     },
@@ -30,13 +34,11 @@ use crate::{
 /// Creates a "real" batch executor which maintains the VM (as opposed to the test builder which doesn't use the VM).
 #[derive(Debug, Clone)]
 pub struct MainBatchExecutor {
-    state_keeper_db_path: String,
-    pool: ConnectionPool,
     save_call_traces: bool,
     max_allowed_tx_gas_limit: U256,
     upload_witness_inputs_to_gcs: bool,
-    enum_index_migration_chunk_size: usize,
     optional_bytecode_compression: bool,
+    cached_storage: Arc<RwLock<CachedStorage>>,
 }
 
 impl MainBatchExecutor {
@@ -50,13 +52,15 @@ impl MainBatchExecutor {
         optional_bytecode_compression: bool,
     ) -> Self {
         Self {
-            state_keeper_db_path,
-            pool,
             save_call_traces,
             max_allowed_tx_gas_limit,
             upload_witness_inputs_to_gcs,
-            enum_index_migration_chunk_size,
             optional_bytecode_compression,
+            cached_storage: Arc::new(RwLock::new(CachedStorage::new(
+                pool,
+                state_keeper_db_path,
+                enum_index_migration_chunk_size,
+            ))),
         }
     }
 }
@@ -69,20 +73,6 @@ impl BatchExecutor for MainBatchExecutor {
         system_env: SystemEnv,
         stop_receiver: &watch::Receiver<bool>,
     ) -> Option<BatchExecutorHandle> {
-        let mut secondary_storage = RocksdbStorage::builder(self.state_keeper_db_path.as_ref())
-            .await
-            .expect("Failed initializing state keeper storage");
-        secondary_storage.enable_enum_index_migration(self.enum_index_migration_chunk_size);
-        let mut conn = self
-            .pool
-            .access_storage_tagged("state_keeper")
-            .await
-            .unwrap();
-        let secondary_storage = secondary_storage
-            .synchronize(&mut conn, stop_receiver)
-            .await
-            .expect("Failed synchronizing secondary state keeper storage")?;
-
         // Since we process `BatchExecutor` commands one-by-one (the next command is never enqueued
         // until a previous command is processed), capacity 1 is enough for the commands channel.
         let (commands_sender, commands_receiver) = mpsc::channel(1);
@@ -94,13 +84,22 @@ impl BatchExecutor for MainBatchExecutor {
         };
         let upload_witness_inputs_to_gcs = self.upload_witness_inputs_to_gcs;
 
+        let cached_storage = self.cached_storage.clone();
+        let stop_receiver = stop_receiver.clone();
         let handle = tokio::task::spawn_blocking(move || {
-            executor.run(
-                secondary_storage,
-                l1_batch_params,
-                system_env,
-                upload_witness_inputs_to_gcs,
-            )
+            let rt_handle = Handle::current();
+            let mut cached_storage = rt_handle.block_on(cached_storage.write());
+            if let Some(storage) = rt_handle
+                .block_on(cached_storage.access_storage(rt_handle.clone(), stop_receiver))
+                .unwrap()
+            {
+                executor.run(
+                    storage,
+                    l1_batch_params,
+                    system_env,
+                    upload_witness_inputs_to_gcs,
+                );
+            };
         });
         Some(BatchExecutorHandle {
             handle,
@@ -124,9 +123,9 @@ struct CommandReceiver {
 }
 
 impl CommandReceiver {
-    pub(super) fn run(
+    pub(super) fn run<S: ReadStorage>(
         mut self,
-        secondary_storage: RocksdbStorage,
+        secondary_storage: S,
         l1_batch_params: L1BatchEnv,
         system_env: SystemEnv,
         upload_witness_inputs_to_gcs: bool,

core/lib/zksync_core/src/state_keeper/cached_storage.rs

@@ -0,0 +1,253 @@
+use std::fmt::Debug;
+
+use anyhow::Context;
+use tokio::{runtime::Handle, sync::watch, task::JoinHandle};
+use zksync_dal::{ConnectionPool, StorageProcessor};
+use zksync_state::{PostgresStorage, ReadStorage, RocksdbStorage};
+use zksync_types::L1BatchNumber;
+
+type BoxReadStorage<'a> = Box<dyn ReadStorage + Send + 'a>;
+
+/// Encapsulates a storage that can produce a short-lived [`ReadStorage`] implementation backed by
+/// either Postgres or RocksDB (if it's caught up). Maintains internal state of how behind RocksDB
+/// is compared to Postgres and actively tries to catch it up in the background.
+///
+/// This struct's main design purpose is to be able to produce [`ReadStorage`] implementation with
+/// as little blocking operations as possible to ensure liveliness.
+#[derive(Debug)]
+pub struct CachedStorage {
+    pool: ConnectionPool,
+    state_keeper_db_path: String,
+    enum_index_migration_chunk_size: usize,
+    state: CachedStorageState,
+}
+
+#[derive(Debug)]
+enum CachedStorageState {
+    /// Cached storage has not been initialized yet, the state of RocksDB to Postgres is undefined
+    NotInitialized,
+    /// RocksDB is trying to catch up to Postgres asynchronously (the task is represented by the
+    /// handle contained inside)
+    CatchingUp {
+        /// Handle owning permission to join on an asynchronous task to catch up RocksDB to
+        /// Postgres.
+        ///
+        /// # `await` value:
+        ///
+        /// - Returns `Ok(Some(rocksdb))` if the process succeeded and `rocksdb` is caught up
+        /// - Returns `Ok(None)` if the process is interrupted
+        /// - Returns `Err(err)` for any propagated Postgres/RocksDB error
+        rocksdb_sync_handle: JoinHandle<anyhow::Result<Option<RocksdbStorage>>>,
+    },
+    /// RocksDB has finished catching up in the near past but is not guaranteed to be in complete
+    /// sync with Postgres. That being said, except for extreme circumstances the gap should be
+    /// small (see [`CATCH_UP_L1_BATCHES_TOLERANCE`]).
+    CaughtUp {
+        /// Last L1 batch number that RocksDB was caught up to
+        rocksdb_l1_batch_number: L1BatchNumber,
+    },
+}
+
+impl CachedStorage {
+    /// Specifies the number of L1 batches we allow for RocksDB to fall behind
+    /// Postgres before we stop trying to catch-up synchronously and instead
+    /// use Postgres-backed [`ReadStorage`] connections.
+    const CATCH_UP_L1_BATCHES_TOLERANCE: u32 = 100;
+
+    pub fn new(
+        pool: ConnectionPool,
+        state_keeper_db_path: String,
+        enum_index_migration_chunk_size: usize,
+    ) -> CachedStorage {
+        CachedStorage {
+            pool,
+            state_keeper_db_path,
+            enum_index_migration_chunk_size,
+            state: CachedStorageState::NotInitialized,
+        }
+    }
+
+    /// Spawns a new asynchronous task that tries to make RocksDB caught up with Postgres.
+    /// The handle to the task is returned as a part of the resulting state.
+    async fn start_catch_up(
+        pool: ConnectionPool,
+        state_keeper_db_path: &str,
+        enum_index_migration_chunk_size: usize,
+        stop_receiver: watch::Receiver<bool>,
+    ) -> anyhow::Result<CachedStorageState> {
+        tracing::debug!("Catching up RocksDB asynchronously");
+        let mut rocksdb_builder = RocksdbStorage::builder(state_keeper_db_path.as_ref())
+            .await
+            .context("Failed initializing RocksDB storage")?;
+        rocksdb_builder.enable_enum_index_migration(enum_index_migration_chunk_size);
+        let rocksdb_sync_handle = tokio::task::spawn(async move {
+            let mut storage = pool.access_storage().await?;
+            rocksdb_builder
+                .synchronize(&mut storage, &stop_receiver)
+                .await
+        });
+        Ok(CachedStorageState::CatchingUp {
+            rocksdb_sync_handle,
+        })
+    }
+
+    /// Returns a [`ReadStorage`] implementation backed by Postgres
+    async fn access_storage_pg(
+        rt_handle: Handle,
+        pool: &ConnectionPool,
+    ) -> anyhow::Result<BoxReadStorage> {
+        let mut connection = pool.access_storage().await?;
+        let postgres_l1_batch_number = connection
+            .blocks_dal()
+            .get_sealed_l1_batch_number()
+            .await?
+            .unwrap_or_default();
+        let block_number = connection
+            .blocks_dal()
+            .get_sealed_miniblock_number()
+            .await?
+            .unwrap_or_default();
+        tracing::debug!(%postgres_l1_batch_number, "Using Postgres-based storage");
+        Ok(
+            Box::new(PostgresStorage::new_async(rt_handle, connection, block_number, true).await?)
+                as BoxReadStorage,
+        )
+    }
+
+    /// Catches up RocksDB synchronously (i.e. assumes the gap is small) and
+    /// returns a [`ReadStorage`] implementation backed by caught-up RocksDB.
+    async fn access_storage_rocksdb<'a>(
+        conn: &mut StorageProcessor<'_>,
+        state_keeper_db_path: &str,
+        enum_index_migration_chunk_size: usize,
+        stop_receiver: &watch::Receiver<bool>,
+    ) -> anyhow::Result<Option<BoxReadStorage<'a>>> {
+        tracing::debug!("Catching up RocksDB synchronously");
+        let mut rocksdb_builder = RocksdbStorage::builder(state_keeper_db_path.as_ref())
+            .await
+            .context("Failed initializing RocksDB storage")?;
+        rocksdb_builder.enable_enum_index_migration(enum_index_migration_chunk_size);
+        let rocksdb = rocksdb_builder
+            .synchronize(conn, stop_receiver)
+            .await
+            .context("Failed to catch up state keeper RocksDB storage to Postgres")?;
+        if let Some(rocksdb) = &rocksdb {
+            let rocksdb_l1_batch_number = rocksdb.l1_batch_number().await.unwrap_or_default();
+            tracing::debug!(%rocksdb_l1_batch_number, "Using RocksDB-based storage");
+        } else {
+            tracing::warn!("Interrupted")
+        }
+        Ok(rocksdb.map(|rocksdb| Box::new(rocksdb) as BoxReadStorage<'a>))
+    }
+
+    /// Produces a [`ReadStorage`] implementation backed by either Postgres or
+    /// RocksDB (if it's caught up).
+    ///
+    /// # Return value
+    ///
+    /// Returns `Ok(None)` if the process is interrupted using `stop_receiver`.
+    ///
+    /// # Errors
+    ///
+    /// Propagates RocksDB and Postgres errors.
+    pub async fn access_storage(
+        &mut self,
+        rt_handle: Handle,
+        stop_receiver: watch::Receiver<bool>,
+    ) -> anyhow::Result<Option<BoxReadStorage>> {
+        // FIXME: This method can potentially be simplified by using recursion but that requires
+        // `BoxFuture` and `Pin` which IMO makes the types much more unreadable than as is
+        match self.state {
+            CachedStorageState::NotInitialized => {
+                // Presuming we are behind, (re-)starting catch up process
+                self.state = Self::start_catch_up(
+                    self.pool.clone(),
+                    &self.state_keeper_db_path,
+                    self.enum_index_migration_chunk_size,
+                    stop_receiver,
+                )
+                .await?;
+                Ok(Some(Self::access_storage_pg(rt_handle, &self.pool).await?))
+            }
+            CachedStorageState::CatchingUp {
+                ref mut rocksdb_sync_handle,
+            } => {
+                if !rocksdb_sync_handle.is_finished() {
+                    // Has not finished catching up yet
+                    return Ok(Some(Self::access_storage_pg(rt_handle, &self.pool).await?));
+                }
+                // RocksDB has finished catching up. Regardless of the outcome, we won't need the
+                // handle anymore so it's safe to replace it
+                let rocksdb_sync_handle = std::mem::replace(
+                    rocksdb_sync_handle,
+                    tokio::task::spawn_blocking(|| Ok(None)),
+                );
+                match rocksdb_sync_handle.await?? {
+                    Some(rocksdb_storage) => {
+                        // Caught-up successfully
+                        self.state = CachedStorageState::CaughtUp {
+                            rocksdb_l1_batch_number: rocksdb_storage
+                                .l1_batch_number()
+                                .await
+                                .unwrap_or_default(),
+                        };
+                        Ok(Some(Self::access_storage_pg(rt_handle, &self.pool).await?))
+                    }
+                    None => {
+                        // Interrupted
+                        self.state = CachedStorageState::NotInitialized;
+                        Ok(None)
+                    }
+                }
+            }
+            CachedStorageState::CaughtUp {
+                rocksdb_l1_batch_number,
+            } => {
+                let mut conn = self
+                    .pool
+                    .access_storage_tagged("state_keeper")
+                    .await
+                    .context("Failed getting a Postgres connection")?;
+                let Some(postgres_l1_batch_number) = conn
+                    .blocks_dal()
+                    .get_sealed_l1_batch_number()
+                    .await
+                    .context("Failed fetching sealed L1 batch number")?
+                else {
+                    return Self::access_storage_rocksdb(
+                        &mut conn,
+                        &self.state_keeper_db_path,
+                        self.enum_index_migration_chunk_size,
+                        &stop_receiver,
+                    )
+                    .await;
+                };
+                if rocksdb_l1_batch_number + Self::CATCH_UP_L1_BATCHES_TOLERANCE
+                    < postgres_l1_batch_number
+                {
+                    tracing::warn!(
+                        %rocksdb_l1_batch_number,
+                        %postgres_l1_batch_number,
+                        "RocksDB fell behind Postgres, trying to catch up asynchronously"
+                    );
+                    self.state = Self::start_catch_up(
+                        self.pool.clone(),
+                        &self.state_keeper_db_path,
+                        self.enum_index_migration_chunk_size,
+                        stop_receiver,
+                    )
+                    .await?;
+                    Ok(Some(Self::access_storage_pg(rt_handle, &self.pool).await?))
+                } else {
+                    Self::access_storage_rocksdb(
+                        &mut conn,
+                        &self.state_keeper_db_path,
+                        self.enum_index_migration_chunk_size,
+                        &stop_receiver,
+                    )
+                    .await
+                }
+            }
+        }
+    }
+}

core/lib/zksync_core/src/state_keeper/mod.rs

@@ -19,6 +19,7 @@ pub use self::{
 use crate::fee_model::BatchFeeModelInputProvider;
 
 mod batch_executor;
+mod cached_storage;
 pub(crate) mod extractors;
 pub(crate) mod io;
 mod keeper;