feat: multi-thread reads and batched writes (#56)

Write batching and read multi-threading improves throughput by 5-18x,
depending on workload type.

Author: losfair

Cargo.lock

@@ -31,7 +31,7 @@ env_logger = "0.10.0"
 futures = "0.3.28"
 hex = "0.4"
 hyper = { version = "0.14", features = ["client"] }
-hyper-proxy = "0.9.1"
+hyper-proxy = { version = "0.9.1", default-features = false }
 log = "0.4.20"
 num-bigint = "0.4"
 prost = "0.11"
@@ -44,6 +44,7 @@ serde_json = "1.0.107"
 tempfile = "3"
 thiserror = "1"
 tokio = { version = "1.33.0", features = ["full"] }
+tokio-stream = "0.1"
 tokio-util = "0.7"
 url = "2"
 uuid = { version = "1.4.1", features = ["v4", "serde"] }

Cargo.toml

@@ -11,6 +11,10 @@ authors.workspace = true
 path = "main.rs"
 name = "denokv"
 
+[features]
+default = ["bundled-sqlite"]
+bundled-sqlite = ["rusqlite/bundled"]
+
 [dependencies]
 anyhow.workspace = true
 aws-config.workspace = true

denokv/Cargo.toml

@@ -39,6 +39,14 @@ pub struct ServeOptions {
   #[clap(long, conflicts_with = "read_only")]
   pub sync_from_s3: bool,
 
+  /// Atomic write batch timeout. Batching is disabled if this is not set.
+  #[clap(long, env = "DENO_KV_ATOMIC_WRITE_BATCH_TIMEOUT_MS")]
+  pub atomic_write_batch_timeout_ms: Option<u64>,
+
+  /// Number of worker threads.
+  #[clap(long, env = "DENO_KV_NUM_WORKERS", default_value = "1")]
+  pub num_workers: usize,
+
   #[command(flatten)]
   pub replica: ReplicaOptions,
 }

denokv/config.rs

@@ -47,6 +47,7 @@ use denokv_proto::SnapshotReadOptions;
 use denokv_sqlite::Connection;
 use denokv_sqlite::Sqlite;
 use denokv_sqlite::SqliteBackendError;
+use denokv_sqlite::SqliteConfig;
 use denokv_sqlite::SqliteNotifier;
 use denokv_timemachine::backup_source_s3::DatabaseBackupSourceS3;
 use denokv_timemachine::backup_source_s3::DatabaseBackupSourceS3Config;
@@ -212,11 +213,18 @@ async fn run_serve(
 
   let path = Path::new(&config.sqlite_path);
   let read_only = options.read_only || options.sync_from_s3;
-  let (sqlite, path) = open_sqlite(path, read_only)?;
+  let sqlite_config = SqliteConfig {
+    batch_timeout: options
+      .atomic_write_batch_timeout_ms
+      .map(std::time::Duration::from_millis),
+    num_workers: options.num_workers,
+  };
+  let sqlite = open_sqlite(path, read_only, sqlite_config.clone())?;
   info!(
-    "Opened{} database at {}",
+    "Opened{} database at {}. Batch timeout: {:?}",
     if read_only { " read only" } else { "" },
-    path,
+    path.to_string_lossy(),
+    sqlite_config.batch_timeout,
   );
 
   let access_token = options.access_token.as_str();
@@ -274,7 +282,7 @@ async fn run_sync(
       let proxy_uri = https_proxy.parse().unwrap();
       let proxy = Proxy::new(Intercept::All, proxy_uri);
       let connector = HttpConnector::new();
-      ProxyConnector::from_proxy(connector, proxy).unwrap()
+      ProxyConnector::from_proxy_unsecured(connector, proxy)
     };
     let hyper_client =
       aws_smithy_client::hyper_ext::Adapter::builder().build(proxy);
@@ -327,20 +335,27 @@ async fn run_sync(
 fn open_sqlite(
   path: &Path,
   read_only: bool,
-) -> Result<(Sqlite, String), anyhow::Error> {
+  config: SqliteConfig,
+) -> Result<Sqlite, anyhow::Error> {
   let flags = if read_only {
     OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX
   } else {
     OpenFlags::SQLITE_OPEN_READ_WRITE
       | OpenFlags::SQLITE_OPEN_CREATE
       | OpenFlags::SQLITE_OPEN_NO_MUTEX
   };
-  let conn = Connection::open_with_flags(path, flags)?;
   let notifier = SqliteNotifier::default();
-  let rng: Box<_> = Box::new(rand::rngs::StdRng::from_entropy());
-  let path = conn.path().unwrap().to_owned();
-  let sqlite = Sqlite::new(conn, notifier, rng)?;
-  Ok((sqlite, path))
+  let sqlite = Sqlite::new(
+    || {
+      Ok((
+        Connection::open_with_flags(path, flags)?,
+        Box::new(rand::rngs::StdRng::from_entropy()),
+      ))
+    },
+    notifier,
+    config,
+  )?;
+  Ok(sqlite)
 }
 
 #[axum::debug_handler]

denokv/main.rs

@@ -3,15 +3,15 @@
 [package]
 authors = []
 edition = "2021"
-name    = "deno-kv-napi"
-version = "0.0.0" # never published
+name = "deno-kv-napi"
+version = "0.0.0"     # never published
 
 [lib]
 crate-type = ["cdylib"]
 
 [dependencies]
 futures = "0.3.28"
-napi = {version="2", features=["async","anyhow"]}
+napi = { version = "2", features = ["async", "anyhow"] }
 napi-derive = "2"
 denokv_sqlite = { path = "../../sqlite" }
 denokv_proto = { path = "../../proto" }
@@ -21,7 +21,7 @@ tokio = { version = "1.33.0", features = ["full"] }
 anyhow = "1"
 prost = "0.11"
 once_cell = "1.18.0"
-rusqlite = "0.29.0"
+rusqlite = { version = "0.29.0", features = ["bundled"] }
 
 [build-dependencies]
 napi-build = "2"

npm/napi/Cargo.toml

@@ -10,6 +10,7 @@ use denokv_proto::WatchKeyOutput;
 use denokv_sqlite::Connection;
 use denokv_sqlite::Sqlite;
 use denokv_sqlite::SqliteBackendError;
+use denokv_sqlite::SqliteConfig;
 use denokv_sqlite::SqliteMessageHandle;
 use denokv_sqlite::SqliteNotifier;
 use futures::stream::BoxStream;
@@ -38,14 +39,11 @@ pub fn open(path: String, in_memory: Option<bool>, debug: bool) -> Result<u32> {
   } else {
     OpenFlags::default()
   };
-  let conn = Connection::open_with_flags(Path::new(&path), flags)
-    .map_err(anyhow::Error::from)?;
-  conn
-    .pragma_update(None, "journal_mode", "wal")
-    .map_err(anyhow::Error::from)?;
-  let rng = Box::new(rand::rngs::StdRng::from_entropy());
-
-  let opened_path = conn.path().map(|p| p.to_string());
+
+  let opened_path = Connection::open_with_flags(Path::new(&path), flags)
+    .map_err(anyhow::Error::from)?
+    .path()
+    .map(|p| p.to_string());
   let notifier = match &opened_path {
     Some(opened_path) if !opened_path.is_empty() => NOTIFIERS_MAP
       .lock()
@@ -56,7 +54,20 @@ pub fn open(path: String, in_memory: Option<bool>, debug: bool) -> Result<u32> {
     _ => SqliteNotifier::default(),
   };
 
-  let sqlite = Sqlite::new(conn, notifier, rng)?;
+  let sqlite = Sqlite::new(
+    || {
+      Ok((
+        Connection::open_with_flags(Path::new(&path), flags)
+          .map_err(anyhow::Error::from)?,
+        Box::new(rand::rngs::StdRng::from_entropy()),
+      ))
+    },
+    notifier,
+    SqliteConfig {
+      num_workers: 1,
+      batch_timeout: None,
+    },
+  )?;
 
   let db_id = DB_ID.fetch_add(1, Ordering::Relaxed);
   DBS.lock().unwrap().insert(db_id, sqlite);