feat: replicated cache

Cargo.toml

@@ -9,6 +9,10 @@ edition = "2021"
 name = "sqrl_bench"
 harness = false
 
+[features]
+default = ["replication"]
+replication = []
+
 
 [dev-dependencies]
 assert_cmd = "0.11"
@@ -39,6 +43,7 @@ rayon = "1.9.0"
 tokio = { version = "1.36.0", features = ["full"] }
 prost = "0.12.4"
 tonic = "0.11.0"
+futures = "0.3.30"
 
 [build-dependencies]
 tonic-build = "0.11.0"

README.md

@@ -1,18 +1,20 @@
-# sqrl
+# squirrel
 
-_Pronounced squirrel, as in "to squirrel away"._
+_As in "[to squirrel away][squirrel-away]"._
 
-A distributed key-value store which uses a simple implementation of [bitcask](https://github.com/basho/bitcask/blob/develop/doc/bitcask-intro.pdf) as the underlying storage mechanism.
+A (replicated) key-value store which uses a simple implementation of [bitcask](https://github.com/basho/bitcask/blob/develop/doc/bitcask-intro.pdf) as the underlying storage mechanism.
 
 ## How it works
 
 This follows a simple structure and exposes a very common API surface: `set(k,v)`, `get(k)`, and `remove(k)`.
 
-## Implementation
+## Notes
 
-It was initially based off my implementation of the PingCAP talent plan course for building a key-value store in Rust:
+This was initially built through my implementation of the PingCAP talent plan course for building a key-value store in Rust:
 
 - [Course](https://github.com/pingcap/talent-plan/tree/master/courses/rust#the-goal-of-this-course)
 - [Lesson plan](https://github.com/pingcap/talent-plan/blob/master/courses/rust/docs/lesson-plan.md#pna-rust-lesson-plan)
 
 And has since grown into my own toy project.
+
+[squirrel-away]: https://dictionary.cambridge.org/dictionary/english/squirrel-away

proto/actions.proto

@@ -15,6 +15,7 @@ message Acknowledgement {
 message SetRequest {
     string key = 1;
     string value = 2;
+    int64 timestamp = 3;
 }
 
 message GetRequest {
@@ -23,6 +24,7 @@ message GetRequest {
 
 message GetResponse {
     optional string value = 1;
+    int64 timestamp = 2;
 }
 
 message RemoveRequest {

src/action.rs

@@ -0,0 +1,16 @@
+use clap::Subcommand;
+use serde::{Deserialize, Serialize};
+
+/// Actions that can be performed by the client.
+#[derive(Debug, Subcommand, Serialize, Deserialize)]
+pub enum Action {
+    /// Enter a key-value pair into the store.
+    Set { key: String, value: String },
+
+    /// Get a value from the store with the provided key.
+    Get { key: String },
+
+    /// Remove a value from the store with the provided key.
+    #[clap(name = "rm")]
+    Remove { key: String },
+}

src/bin/sqrl-client.rs

@@ -1,7 +1,7 @@
 use clap::Parser;
-use sqrl::actions::action_client::ActionClient;
-use sqrl::actions::{GetRequest, SetRequest};
-use sqrl::client::Action;
+
+use sqrl::action::Action;
+use sqrl::client::{Client, RemoteNodeClient};
 
 #[derive(Debug, Parser)]
 #[command(author, version, about, long_about = None)]
@@ -16,26 +16,22 @@ struct App {
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     let cli = App::parse();
-    let mut client = ActionClient::connect(format!("http://{}", cli.server)).await?;
+    let mut client = RemoteNodeClient::new(cli.server).await?;
 
     match cli.subcmd {
         Action::Set { key, value } => {
-            client
-                .set(tonic::Request::new(SetRequest { key, value }))
-                .await?;
+            client.set(key, value).await?;
         }
         Action::Get { key } => {
-            let response = client.get(tonic::Request::new(GetRequest { key })).await?;
-            match response.into_inner().value {
-                Some(v) => println!("{}", v),
+            let response = client.get(key).await?;
+            match response {
+                Some(v) => println!("{}", v.value.unwrap()),
                 None => println!("Key not found"),
             }
         }
         Action::Remove { key } => {
-            let response = client
-                .remove(tonic::Request::new(sqrl::actions::RemoveRequest { key }))
-                .await?;
-            match response.into_inner().success {
+            let response = client.remove(key).await?;
+            match response.success {
                 true => {}
                 false => {
                     eprintln!("Key not found");

src/bin/sqrl-server.rs

@@ -1,14 +1,18 @@
 use clap::Parser;
-use sqrl::actions;
-use sqrl::actions::action_server::Action as ActionSrv;
-use sqrl::actions::action_server::ActionServer;
+use futures::StreamExt;
+use sqrl::client::RemoteNodeClient;
+use sqrl::replication;
+use sqrl::replication::ReplicatedServer;
 use sqrl::KvStore;
-use sqrl::KvsEngine;
+use sqrl::StandaloneServer;
 use sqrl::ENGINE_FILE;
-use std::sync::Arc;
 use std::{ffi::OsString, path::PathBuf};
 use std::{fmt::Display, net::SocketAddr};
-use tracing::info;
+use tracing::{debug, warn};
+
+mod proto {
+    tonic::include_proto!("actions");
+}
 
 #[derive(Debug, Parser)]
 #[command(author, version, about, long_about = None)]
@@ -24,6 +28,18 @@ struct App {
 
     #[arg(long, global = true, default_value = default_log_location())]
     log_file: PathBuf,
+
+    #[cfg(feature = "replication")]
+    #[arg(long, default_value = "follower")]
+    replication_mode: replication::Mode,
+
+    #[cfg(feature = "replication")]
+    #[arg(
+        long,
+        requires_if(replication::Mode::Leader, "replication_mode"),
+        value_delimiter = ','
+    )]
+    followers: Vec<String>,
 }
 
 fn default_log_location() -> OsString {
@@ -48,76 +64,30 @@ impl Display for Engine {
     }
 }
 
-#[derive(Clone)]
-struct KvServer {
-    pub store: Arc<KvStore>,
-}
-
-impl KvServer {
-    pub fn new<P>(path: P) -> anyhow::Result<Self>
-    where
-        P: Into<std::path::PathBuf>,
-    {
-        let store = Arc::new(KvStore::open(path)?);
-        Ok(Self { store })
-    }
-}
-
-#[tonic::async_trait]
-impl ActionSrv for KvServer {
-    async fn get(
-        &self,
-        req: tonic::Request<actions::GetRequest>,
-    ) -> tonic::Result<tonic::Response<actions::GetResponse>, tonic::Status> {
-        let req = req.into_inner();
-        let value = self.store.get(req.key).await.unwrap();
-        Ok(tonic::Response::new(actions::GetResponse { value }))
-    }
-    async fn set(
-        &self,
-        req: tonic::Request<actions::SetRequest>,
-    ) -> tonic::Result<tonic::Response<actions::Acknowledgement>, tonic::Status> {
-        let req = req.into_inner();
-        self.store.set(req.key, req.value).await.unwrap();
-        Ok(tonic::Response::new(actions::Acknowledgement {
-            success: true,
-        }))
-    }
-    async fn remove(
-        &self,
-        req: tonic::Request<actions::RemoveRequest>,
-    ) -> tonic::Result<tonic::Response<actions::Acknowledgement>, tonic::Status> {
-        let req = req.into_inner();
-        match self.store.remove(req.key).await {
-            Ok(_) => Ok(tonic::Response::new(actions::Acknowledgement {
-                success: true,
-            })),
-            Err(_) => Ok(tonic::Response::new(actions::Acknowledgement {
-                success: false,
-            })),
-        }
-    }
-}
-
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     let app = App::parse();
-
     // We must error if the previous storage engine was not 'sqrl' as it is incompatible.
     KvStore::engine_is_sqrl(app.engine_name.to_string(), app.log_file.join(ENGINE_FILE))?;
-    let srv = KvServer::new(app.log_file)?;
-
-    info!(
-        "sqrl-server version: {}, engine: {}",
-        env!("CARGO_PKG_VERSION"),
-        app.engine_name
-    );
-
-    info!("Listening on {}", app.addr);
-    tonic::transport::Server::builder()
-        .add_service(ActionServer::new(srv))
-        .serve(app.addr)
-        .await?;
-
-    Ok(())
+    match app.replication_mode {
+        replication::Mode::Leader => {
+            assert_eq!(
+                app.followers.len(),
+                2,
+                "Only 2 followers are configurable at present"
+            );
+            let clients = futures::stream::iter(app.followers.iter())
+                .filter_map(|f| async move { RemoteNodeClient::new(f.to_string()).await.ok() })
+                .collect::<Vec<RemoteNodeClient>>()
+                .await;
+            debug!("Replicating to {} followers", clients.len());
+            if clients.len() > 3 {
+                warn!("Replicating to many followers can greatly impact write performance");
+            }
+            ReplicatedServer::new(clients.into(), app.log_file, app.addr)?
+                .run()
+                .await
+        }
+        replication::Mode::Follower => StandaloneServer::new(app.log_file, app.addr)?.run().await,
+    }
 }