perf(store): Shard sqlite

Author: john-z-yang

benches/store_bench.rs

@@ -5,7 +5,7 @@ use criterion::{Criterion, criterion_group, criterion_main};
 use rand::Rng;
 use taskbroker::{
     store::inflight_activation::{
-        InflightActivationStatus, InflightActivationStore, InflightActivationStoreConfig,
+        InflightActivationShard, InflightActivationStatus, InflightActivationStoreConfig,
     },
     test_utils::{generate_temp_filename, make_activations},
 };
@@ -23,9 +23,10 @@ async fn process_activations(num_activations: u32, num_workers: u32) {
         generate_temp_filename()
     };
     let store = Arc::new(
-        InflightActivationStore::new(
+        InflightActivationShard::new(
             &url,
             InflightActivationStoreConfig {
+                sharding_factor: 8,
                 max_processing_attempts: 1,
             },
         )
@@ -80,7 +81,7 @@ fn store_bench(c: &mut Criterion) {
     let num_activations: u32 = 4_096;
     let num_workers = 64;
 
-    c.benchmark_group("bench_InflightActivationStore")
+    c.benchmark_group("bench_InflightActivationShard")
         .sample_size(256)
         .throughput(criterion::Throughput::Elements(num_activations.into()))
         .bench_function("process_activations", |b| {

src/config.rs

@@ -75,6 +75,9 @@ pub struct Config {
     /// The path to the sqlite database
     pub db_path: String,
 
+    /// The number of physical files to shard the database by
+    pub db_sharding_factor: u8,
+
     /// The maximum number of pending records that can be
     /// in the InflightTaskStore (sqlite)
     pub max_pending_count: usize,
@@ -115,7 +118,8 @@ impl Default for Config {
             kafka_auto_commit_interval_ms: 5000,
             kafka_auto_offset_reset: "latest".to_owned(),
             kafka_send_timeout_ms: 500,
-            db_path: "./taskbroker-inflight.sqlite".to_owned(),
+            db_path: "./taskbroker-inflight".to_owned(),
+            db_sharding_factor: 8,
             max_pending_count: 2048,
             max_pending_buffer_count: 128,
             max_processing_attempts: 5,
@@ -197,7 +201,7 @@ mod tests {
         assert_eq!(config.log_format, LogFormat::Text);
         assert_eq!(config.grpc_port, 50051);
         assert_eq!(config.kafka_topic, "task-worker");
-        assert_eq!(config.db_path, "./taskbroker-inflight.sqlite");
+        assert_eq!(config.db_path, "./taskbroker-inflight");
         assert_eq!(config.max_pending_count, 2048);
     }
 
@@ -284,7 +288,7 @@ mod tests {
             assert_eq!(config.log_filter, "error");
             assert_eq!(config.kafka_topic, "task-worker".to_owned());
             assert_eq!(config.kafka_deadletter_topic, "task-worker-dlq".to_owned());
-            assert_eq!(config.db_path, "./taskbroker-inflight.sqlite".to_owned());
+            assert_eq!(config.db_path, "./taskbroker-inflight".to_owned());
             assert_eq!(config.max_pending_count, 2048);
             assert_eq!(config.max_processing_attempts, 5);
             assert_eq!(

src/grpc/server.rs

@@ -9,12 +9,12 @@ use std::time::Instant;
 use tonic::{Request, Response, Status};
 
 use crate::store::inflight_activation::{
-    InflightActivation, InflightActivationStatus, InflightActivationStore,
+    InflightActivation, InflightActivationShard, InflightActivationStatus,
 };
 use tracing::{error, instrument};
 
 pub struct TaskbrokerServer {
-    pub store: Arc<InflightActivationStore>,
+    pub store: Arc<InflightActivationShard>,
 }
 
 #[tonic::async_trait]

src/kafka/inflight_activation_writer.rs

@@ -6,7 +6,7 @@ use tracing::{debug, instrument};
 
 use crate::{
     config::Config,
-    store::inflight_activation::{InflightActivation, InflightActivationStore},
+    store::inflight_activation::{InflightActivation, InflightActivationShard},
 };
 
 use super::consumer::{
@@ -31,12 +31,12 @@ impl ActivationWriterConfig {
 
 pub struct InflightActivationWriter {
     config: ActivationWriterConfig,
-    store: Arc<InflightActivationStore>,
+    store: Arc<InflightActivationShard>,
     batch: Option<Vec<InflightActivation>>,
 }
 
 impl InflightActivationWriter {
-    pub fn new(store: Arc<InflightActivationStore>, config: ActivationWriterConfig) -> Self {
+    pub fn new(store: Arc<InflightActivationShard>, config: ActivationWriterConfig) -> Self {
         Self {
             config,
             store,
@@ -87,14 +87,13 @@ impl Reducer for InflightActivationWriter {
                 .min_by_key(|item| item.timestamp())
                 .unwrap();
 
-        let res = self.store.store(take(&mut self.batch).unwrap()).await?;
+        let rows_affected = self.store.store(take(&mut self.batch).unwrap()).await?;
         metrics::histogram!("consumer.inflight_activation_writer.insert_lag")
             .record(lag.num_seconds() as f64);
-        metrics::counter!("consumer.inflight_activation_writer.stored")
-            .increment(res.rows_affected);
+        metrics::counter!("consumer.inflight_activation_writer.stored").increment(rows_affected);
         debug!(
             "Inserted {:?} entries with max lag: {:?}s",
-            res.rows_affected,
+            rows_affected,
             lag.num_seconds()
         );
 

src/main.rs

@@ -28,7 +28,7 @@ use taskbroker::logging;
 use taskbroker::metrics;
 use taskbroker::processing_strategy;
 use taskbroker::store::inflight_activation::{
-    InflightActivationStore, InflightActivationStoreConfig,
+    InflightActivationShard, InflightActivationStoreConfig,
 };
 use taskbroker::{Args, get_version};
 
@@ -57,7 +57,7 @@ async fn main() -> Result<(), Error> {
     logging::init(logging::LoggingConfig::from_config(&config));
     metrics::init(metrics::MetricsConfig::from_config(&config));
     let store = Arc::new(
-        InflightActivationStore::new(
+        InflightActivationShard::new(
             &config.db_path,
             InflightActivationStoreConfig::from_config(&config),
         )