Remove RocksDb writer thread and disable WAL

crates/storage-query-datafusion/src/idempotency/tests.rs

@@ -27,7 +27,7 @@ async fn get_idempotency_key() {
         .default_runtime_handle(tokio::runtime::Handle::current())
         .build()
         .expect("task_center builds");
-    let (mut engine, shutdown) = tc
+    let mut engine = tc
         .run_in_scope("mock-query-engine", None, MockQueryEngine::create())
         .await;
 
@@ -96,6 +96,4 @@ async fn get_idempotency_key() {
             )
         )
     );
-
-    shutdown.await;
 }

crates/storage-query-datafusion/src/journal/tests.rs

@@ -34,7 +34,7 @@ async fn get_entries() {
         .default_runtime_handle(tokio::runtime::Handle::current())
         .build()
         .expect("task_center builds");
-    let (mut engine, shutdown) = tc
+    let mut engine = tc
         .run_in_scope("mock-query-engine", None, MockQueryEngine::create())
         .await;
 
@@ -127,6 +127,4 @@ async fn get_entries() {
             )
         )
     );
-
-    shutdown.await;
 }

crates/storage-query-datafusion/src/mocks.rs

@@ -28,7 +28,6 @@ use restate_types::config::{CommonOptions, QueryEngineOptions, WorkerOptions};
 use restate_types::identifiers::{DeploymentId, ServiceRevision};
 use restate_types::invocation::ServiceType;
 use std::fmt::Debug;
-use std::future::Future;
 use std::marker::PhantomData;
 
 #[derive(Default, Clone, Debug)]
@@ -87,35 +86,27 @@ impl MockQueryEngine {
             + Debug
             + Clone
             + 'static,
-    ) -> (Self, impl Future<Output = ()>) {
+    ) -> Self {
         // Prepare Rocksdb
         task_center().run_in_scope_sync("db-manager-init", None, || {
             RocksDbManager::init(Constant::new(CommonOptions::default()))
         });
         let worker_options = WorkerOptions::default();
-        let (rocksdb, writer) = RocksDBStorage::open(
+        let rocksdb = RocksDBStorage::open(
             Constant::new(worker_options.storage.clone()),
             Constant::new(worker_options.storage.rocksdb),
         )
         .await
         .expect("RocksDB storage creation should succeed");
-        let (signal, watch) = drain::channel();
-        let writer_join_handle = writer.run(watch);
 
-        let query_engine = Self(
+        Self(
             rocksdb.clone(),
             QueryContext::from_options(&QueryEngineOptions::default(), rocksdb, status, schemas)
                 .unwrap(),
-        );
-
-        // Return shutdown future
-        (query_engine, async {
-            signal.drain().await;
-            writer_join_handle.await.unwrap().unwrap();
-        })
+        )
     }
 
-    pub async fn create() -> (Self, impl Future<Output = ()>) {
+    pub async fn create() -> Self {
         Self::create_with(MockStatusHandle::default(), MockSchemas::default()).await
     }
 

crates/storage-query-datafusion/src/tests.rs

@@ -39,7 +39,7 @@ async fn query_sys_invocation() {
         .default_runtime_handle(tokio::runtime::Handle::current())
         .build()
         .expect("task_center builds");
-    let (mut engine, shutdown) = tc
+    let mut engine = tc
         .run_in_scope(
             "mock-query-engine",
             None,
@@ -117,6 +117,4 @@ async fn query_sys_invocation() {
             }
         ))
     );
-
-    shutdown.await;
 }

crates/storage-rocksdb/benches/basic_benchmark.rs

@@ -24,16 +24,13 @@ async fn writing_to_rocksdb(worker_options: WorkerOptions) {
     //
     // setup
     //
-    let (mut rocksdb, writer) = RocksDBStorage::open(
+    let mut rocksdb = RocksDBStorage::open(
         Constant::new(worker_options.storage.clone()),
         Constant::new(worker_options.storage.rocksdb),
     )
     .await
     .expect("RocksDB storage creation should succeed");
 
-    let (signal, watch) = drain::channel();
-    let writer_join_handler = writer.run(watch);
-
     //
     // write
     //
@@ -45,9 +42,6 @@ async fn writing_to_rocksdb(worker_options: WorkerOptions) {
         }
         txn.commit().await.unwrap();
     }
-
-    signal.drain().await;
-    writer_join_handler.await.unwrap().unwrap();
 }
 
 fn basic_writing_reading_benchmark(c: &mut Criterion) {

crates/storage-rocksdb/src/lib.rs

@@ -21,11 +21,9 @@ pub mod scan;
 pub mod service_status_table;
 pub mod state_table;
 pub mod timer_table;
-mod writer;
 
 use crate::keys::TableKey;
 use crate::scan::{PhysicalScan, TableScan};
-use crate::writer::{Writer, WriterHandle};
 use crate::TableKind::{
     Deduplication, Idempotency, Inbox, InvocationStatus, Journal, Outbox, PartitionStateMachine,
     ServiceStatus, State, Timers,
@@ -49,17 +47,12 @@ use rocksdb::PrefixRange;
 use rocksdb::ReadOptions;
 use std::sync::Arc;
 
-pub use writer::JoinHandle as RocksDBWriterJoinHandle;
-pub use writer::Writer as RocksDBWriter;
-
 pub type DB = rocksdb::OptimisticTransactionDB<MultiThreaded>;
 type TransactionDB<'a> = rocksdb::Transaction<'a, DB>;
 
 pub type DBIterator<'b> = DBRawIteratorWithThreadMode<'b, DB>;
 pub type DBIteratorTransaction<'b> = DBRawIteratorWithThreadMode<'b, rocksdb::Transaction<'b, DB>>;
 
-type WriteBatch = rocksdb::WriteBatchWithTransaction<true>;
-
 // matches the default directory name
 const DB_NAME: &str = "db";
 
@@ -156,7 +149,6 @@ pub enum BuildError {
 
 pub struct RocksDBStorage {
     db: Arc<DB>,
-    writer_handle: WriterHandle,
     key_buffer: BytesMut,
     value_buffer: BytesMut,
 }
@@ -165,7 +157,6 @@ impl std::fmt::Debug for RocksDBStorage {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("RocksDBStorage")
             .field("db", &self.db)
-            .field("writer_handle", &self.writer_handle)
             .field("key_buffer", &self.key_buffer)
             .field("value_buffer", &self.value_buffer)
             .finish()
@@ -176,7 +167,6 @@ impl Clone for RocksDBStorage {
     fn clone(&self) -> Self {
         RocksDBStorage {
             db: self.db.clone(),
-            writer_handle: self.writer_handle.clone(),
             key_buffer: BytesMut::default(),
             value_buffer: BytesMut::default(),
         }
@@ -231,7 +221,7 @@ impl RocksDBStorage {
     pub async fn open(
         mut storage_opts: impl Updateable<StorageOptions> + Send + 'static,
         updateable_opts: impl Updateable<RocksDbOptions> + Send + 'static,
-    ) -> std::result::Result<(Self, Writer), BuildError> {
+    ) -> std::result::Result<Self, BuildError> {
         let cfs = vec![
             //
             // keyed by partition key + user key
@@ -272,18 +262,11 @@ impl RocksDBStorage {
         .await
         .map_err(|_| ShutdownError)??;
 
-        let writer = Writer::new(rdb.clone(), storage_opts);
-        let writer_handle = writer.create_writer_handle();
-
-        Ok((
-            Self {
-                db: rdb,
-                writer_handle,
-                key_buffer: BytesMut::default(),
-                value_buffer: BytesMut::default(),
-            },
-            writer,
-        ))
+        Ok(Self {
+            db: rdb,
+            key_buffer: BytesMut::default(),
+            value_buffer: BytesMut::default(),
+        })
     }
 
     fn table_handle(&self, table_kind: TableKind) -> Arc<BoundColumnFamily> {
@@ -341,7 +324,6 @@ impl RocksDBStorage {
             db,
             key_buffer: &mut self.key_buffer,
             value_buffer: &mut self.value_buffer,
-            writer_handle: &self.writer_handle,
         }
     }
 }
@@ -406,7 +388,6 @@ pub struct RocksDBTransaction<'a> {
     db: Arc<DB>,
     key_buffer: &'a mut BytesMut,
     value_buffer: &'a mut BytesMut,
-    writer_handle: &'a WriterHandle,
 }
 
 impl<'a> RocksDBTransaction<'a> {
@@ -446,8 +427,16 @@ impl<'a> Transaction for RocksDBTransaction<'a> {
         // writes to RocksDB. However, it is safe to write the WriteBatch for a given partition,
         // because there can only be a single writer (the leading PartitionProcessor).
         let write_batch = self.txn.get_writebatch();
-
-        self.writer_handle.write(write_batch).await
+        // todo: make async and use configuration to control use of WAL
+        if write_batch.is_empty() {
+            return Ok(());
+        }
+        let mut opts = rocksdb::WriteOptions::default();
+        // We disable WAL since bifrost is our durable distributed log.
+        opts.disable_wal(true);
+        self.db
+            .write_opt(&write_batch, &rocksdb::WriteOptions::default())
+            .map_err(|error| StorageError::Generic(error.into()))
     }
 }
 

crates/storage-rocksdb/tests/idempotency_table_test/mod.rs

@@ -35,7 +35,7 @@ const IDEMPOTENCY_ID_3: IdempotencyId =
 
 #[tokio::test]
 async fn test_idempotency_key() {
-    let (mut rocksdb, close) = storage_test_environment().await;
+    let mut rocksdb = storage_test_environment().await;
 
     // Fill in some data
     let mut txn = rocksdb.transaction();
@@ -114,6 +114,4 @@ async fn test_idempotency_key() {
             .unwrap(),
         None
     );
-
-    close.await;
 }

crates/storage-rocksdb/tests/integration_test.rs

@@ -21,7 +21,6 @@ use restate_types::invocation::{InvocationTarget, ServiceInvocation, Source, Spa
 use restate_types::state_mut::ExternalStateMutation;
 use std::collections::HashMap;
 use std::fmt::Debug;
-use std::future::Future;
 use std::pin::pin;
 use tokio_stream::StreamExt;
 
@@ -34,7 +33,7 @@ mod state_table_test;
 mod timer_table_test;
 mod virtual_object_status_table_test;
 
-async fn storage_test_environment() -> (RocksDBStorage, impl Future<Output = ()>) {
+async fn storage_test_environment() -> RocksDBStorage {
     //
     // create a rocksdb storage from options
     //
@@ -46,25 +45,17 @@ async fn storage_test_environment() -> (RocksDBStorage, impl Future<Output = ()>
         RocksDbManager::init(Constant::new(CommonOptions::default()))
     });
     let worker_options = WorkerOptions::default();
-    let (rocksdb, writer) = RocksDBStorage::open(
+    RocksDBStorage::open(
         Constant::new(worker_options.storage.clone()),
         Constant::new(worker_options.storage.rocksdb),
     )
     .await
-    .expect("RocksDB storage creation should succeed");
-
-    let (signal, watch) = drain::channel();
-    let writer_join_handle = writer.run(watch);
-
-    (rocksdb, async {
-        signal.drain().await;
-        writer_join_handle.await.unwrap().unwrap();
-    })
+    .expect("RocksDB storage creation should succeed")
 }
 
 #[tokio::test]
 async fn test_read_write() {
-    let (rocksdb, close) = storage_test_environment().await;
+    let rocksdb = storage_test_environment().await;
 
     //
     // run the tests
@@ -76,8 +67,6 @@ async fn test_read_write() {
     invocation_status_table_test::run_tests(rocksdb.clone()).await;
     virtual_object_status_table_test::run_tests(rocksdb.clone()).await;
     timer_table_test::run_tests(rocksdb).await;
-
-    close.await;
 }
 
 pub(crate) fn mock_service_invocation(service_id: ServiceId) -> ServiceInvocation {

crates/storage-rocksdb/tests/state_table_test/mod.rs

@@ -112,7 +112,7 @@ pub(crate) async fn run_tests(mut rocksdb: RocksDBStorage) {
 
 #[tokio::test]
 async fn test_delete_all() {
-    let (mut rocksdb, close) = storage_test_environment().await;
+    let mut rocksdb = storage_test_environment().await;
 
     let mut txn = rocksdb.transaction();
 
@@ -143,6 +143,4 @@ async fn test_delete_all() {
         .await
         .expect("should not fail")
         .is_some());
-
-    close.await;
 }