feat(storage_api2): load balancer

VERSION

@@ -1 +1 @@
-250702.0
+250703.0

crates/rpc/src/client2.rs

@@ -566,5 +566,3 @@ impl metrics::Enum for ErrorKind {
         self.into()
     }
 }
-
-// TODO: Vec<Connection> Load Balancer

crates/storage_api2/src/lib.rs

@@ -7,7 +7,13 @@ pub use {
 use {
     futures::FutureExt as _,
     serde::{Deserialize, Serialize},
-    std::{borrow::Cow, future::Future, str::FromStr, time::Duration},
+    std::{
+        borrow::Cow,
+        future::Future,
+        str::FromStr,
+        sync::{atomic, atomic::AtomicUsize},
+        time::Duration,
+    },
     time::OffsetDateTime as DateTime,
 };
 
@@ -55,13 +61,13 @@ pub type KeyspaceVersion = u64;
 ///   servers).
 /// - Replicas use it to finally execute the operations on their local WCN
 ///   Database instances.
-pub trait StorageApi: Clone + Send + Sync + 'static {
+pub trait StorageApi: Send + Sync + 'static {
     /// Executes the provided [`operation::Get`].
     fn get<'a>(
         &'a self,
         get: &'a operation::Get<'a>,
     ) -> impl Future<Output = Result<Option<Record<'a>>>> + Send + 'a {
-        self.execute_ref(get)
+        self.execute_ref(OperationRef::Get(get))
             .map(operation::Output::downcast_result)
     }
 
@@ -70,7 +76,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         set: &'a operation::Set<'a>,
     ) -> impl Future<Output = Result<()>> + Send + 'a {
-        self.execute_ref(set)
+        self.execute_ref(OperationRef::Set(set))
             .map(operation::Output::downcast_result)
     }
 
@@ -79,7 +85,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         del: &'a operation::Del<'a>,
     ) -> impl Future<Output = Result<()>> + Send + 'a {
-        self.execute_ref(del)
+        self.execute_ref(OperationRef::Del(del))
             .map(operation::Output::downcast_result)
     }
 
@@ -88,7 +94,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         get_exp: &'a operation::GetExp<'a>,
     ) -> impl Future<Output = Result<Option<RecordExpiration>>> + Send + 'a {
-        self.execute_ref(get_exp)
+        self.execute_ref(OperationRef::GetExp(get_exp))
             .map(operation::Output::downcast_result)
     }
 
@@ -97,7 +103,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         set_exp: &'a operation::SetExp<'a>,
     ) -> impl Future<Output = Result<()>> + Send + 'a {
-        self.execute_ref(set_exp)
+        self.execute_ref(OperationRef::SetExp(set_exp))
             .map(operation::Output::downcast_result)
     }
 
@@ -106,7 +112,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         hget: &'a operation::HGet<'a>,
     ) -> impl Future<Output = Result<Option<Record<'a>>>> + Send + 'a {
-        self.execute_ref(hget)
+        self.execute_ref(OperationRef::HGet(hget))
             .map(operation::Output::downcast_result)
     }
 
@@ -115,7 +121,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         hset: &'a operation::HSet<'a>,
     ) -> impl Future<Output = Result<()>> + Send + 'a {
-        self.execute_ref(hset)
+        self.execute_ref(OperationRef::HSet(hset))
             .map(operation::Output::downcast_result)
     }
 
@@ -124,7 +130,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         hdel: &'a operation::HDel<'a>,
     ) -> impl Future<Output = Result<()>> + Send + 'a {
-        self.execute_ref(hdel)
+        self.execute_ref(OperationRef::HDel(hdel))
             .map(operation::Output::downcast_result)
     }
 
@@ -133,7 +139,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         hget_exp: &'a operation::HGetExp<'a>,
     ) -> impl Future<Output = Result<Option<RecordExpiration>>> + Send + 'a {
-        self.execute_ref(hget_exp)
+        self.execute_ref(OperationRef::HGetExp(hget_exp))
             .map(operation::Output::downcast_result)
     }
 
@@ -142,7 +148,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         hset_exp: &'a operation::HSetExp<'a>,
     ) -> impl Future<Output = Result<()>> + Send + 'a {
-        self.execute_ref(hset_exp)
+        self.execute_ref(OperationRef::HSetExp(hset_exp))
             .map(operation::Output::downcast_result)
     }
 
@@ -151,7 +157,7 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         hcard: &'a operation::HCard<'a>,
     ) -> impl Future<Output = Result<u64>> + Send + 'a {
-        self.execute_ref(hcard)
+        self.execute_ref(OperationRef::HCard(hcard))
             .map(operation::Output::downcast_result)
     }
 
@@ -160,16 +166,16 @@ pub trait StorageApi: Clone + Send + Sync + 'static {
         &'a self,
         hscan: &'a operation::HScan<'a>,
     ) -> impl Future<Output = Result<MapPage<'a>>> + Send + 'a {
-        self.execute_ref(hscan)
+        self.execute_ref(OperationRef::HScan(hscan))
             .map(operation::Output::downcast_result)
     }
 
     /// Executes the provided [`StorageApi`] [`OperationRef`].
     fn execute_ref<'a>(
         &'a self,
-        operation: impl Into<OperationRef<'a>> + Send + 'a,
+        operation: OperationRef<'a>,
     ) -> impl Future<Output = Result<operation::Output<'a>>> + Send + 'a {
-        self.execute(operation.into().to_owned())
+        self.execute(operation.to_owned())
     }
 
     /// Executes the provided [`StorageApi`] [`Operation`].
@@ -410,6 +416,58 @@ impl From<ErrorKind> for Error {
     }
 }
 
+/// [`StorageApi`] Load balancer.
+///
+/// Load balances execution of [`Operation`]s across a list of [`StorageApi`]
+/// implementors using a round-robin strategy.
+pub struct LoadBalancer<S> {
+    /// List of [`StorageApi`]s of this [`LoadBalancer`].
+    ///
+    /// Can be safely modified at any point.
+    pub apis: Vec<S>,
+
+    counter: AtomicUsize,
+}
+
+impl<S: StorageApi> LoadBalancer<S> {
+    /// Creates a new [`LoadBalancer`].
+    pub fn new(connections: impl IntoIterator<Item = S>) -> Self {
+        Self {
+            apis: connections.into_iter().collect(),
+            // overflows and starts from `0`
+            counter: usize::MAX.into(),
+        }
+    }
+
+    fn next_api(&self) -> Result<&S> {
+        if self.apis.is_empty() {
+            return Err(Error::new(
+                ErrorKind::Internal,
+                Some("LoadBalancer::apis is empty".to_string()),
+            ));
+        }
+
+        let n = self.counter.fetch_add(1, atomic::Ordering::Relaxed);
+        Ok(&self.apis[n % self.apis.len()])
+    }
+}
+
+impl<S: StorageApi> StorageApi for LoadBalancer<S> {
+    async fn execute_ref<'a>(
+        &'a self,
+        operation: OperationRef<'a>,
+    ) -> Result<operation::Output<'a>> {
+        self.next_api()?.execute_ref(operation).await
+    }
+
+    async fn execute<'a>(
+        &'a self,
+        operation: crate::Operation<'a>,
+    ) -> Result<operation::Output<'a>> {
+        self.next_api()?.execute(operation).await
+    }
+}
+
 #[cfg(test)]
 #[test]
 fn test_namespace_from_str() {

crates/storage_api2/src/rpc/client.rs

@@ -105,9 +105,9 @@ where
 
     async fn execute_ref<'a>(
         &'a self,
-        operation: impl Into<crate::OperationRef<'a>> + Send + 'a,
+        operation: crate::OperationRef<'a>,
     ) -> Result<operation::Output<'a>> {
-        match operation.into() {
+        match operation {
             OperationRef::Get(get) => self.get(get).await.map(Into::into),
             OperationRef::Set(set) => self.set(set).await.map(Into::into),
             OperationRef::Del(del) => self.del(del).await.map(Into::into),

crates/storage_api2/src/rpc/server.rs

@@ -10,21 +10,21 @@ use {
 };
 
 /// Creates a new [`CoordinatorApi`] RPC server.
-pub fn coordinator(storage_api: impl StorageApi) -> impl wcn_rpc::server2::Server {
+pub fn coordinator(storage_api: impl StorageApi + Clone) -> impl wcn_rpc::server2::Server {
     new::<api_kind::Coordinator>(storage_api)
 }
 
 /// Creates a new [`ReplicaApi`] RPC server.
-pub fn replica(storage_api: impl StorageApi) -> impl wcn_rpc::server2::Server {
+pub fn replica(storage_api: impl StorageApi + Clone) -> impl wcn_rpc::server2::Server {
     new::<api_kind::Replica>(storage_api)
 }
 
 /// Creates a new [`DatabaseApi`] RPC server.
-pub fn database(storage_api: impl StorageApi) -> impl wcn_rpc::server2::Server {
+pub fn database(storage_api: impl StorageApi + Clone) -> impl wcn_rpc::server2::Server {
     new::<api_kind::Database>(storage_api)
 }
 
-fn new<Kind>(storage_api: impl StorageApi) -> impl wcn_rpc::server2::Server
+fn new<Kind>(storage_api: impl StorageApi + Clone) -> impl wcn_rpc::server2::Server
 where
     Kind: Clone + Send + Sync + 'static,
     Api<Kind>: wcn_rpc::Api<RpcId = RpcId>,
@@ -43,7 +43,7 @@ struct ConnectionHandler<S: StorageApi, Kind> {
 
 impl<S, Kind> HandleConnection for ConnectionHandler<S, Kind>
 where
-    S: StorageApi,
+    S: StorageApi + Clone,
     Kind: Clone + Send + Sync + 'static,
     Api<Kind>: wcn_rpc::Api<RpcId = RpcId>,
 {

crates/storage_api2/tests/integration.rs

@@ -16,6 +16,7 @@ use {
         operation,
         Bytes,
         KeyspaceVersion,
+        LoadBalancer,
         MapEntry,
         MapPage,
         Namespace,
@@ -75,9 +76,9 @@ async fn test_rpc_api<API, S>(
         port: server_port,
         keypair: server_keypair,
         connection_timeout: Duration::from_secs(10),
-        max_connections: 1,
-        max_connections_per_ip: 1,
-        max_connection_rate_per_ip: 1,
+        max_connections: 2,
+        max_connections_per_ip: 2,
+        max_connection_rate_per_ip: 2,
         max_concurrent_rpcs: 10,
         priority: transport::Priority::High,
     };
@@ -95,14 +96,18 @@ async fn test_rpc_api<API, S>(
     let client = client(client_config).unwrap();
 
     let server_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, server_port);
-    let client_conn = client
+    let client_conn1 = client
+        .connect(server_addr, &server_peer_id, ())
+        .await
+        .unwrap();
+    let client_conn2 = client
         .connect(server_addr, &server_peer_id, ())
         .await
         .unwrap();
 
     let ctx = &TestContext {
         storage,
-        client_conn,
+        lb: LoadBalancer::new([client_conn1, client_conn2]),
     };
 
     ctx.test_operations().await;
@@ -112,7 +117,7 @@ async fn test_rpc_api<API, S>(
 
 struct TestContext<API: Api> {
     storage: TestStorage,
-    client_conn: Connection<API>,
+    lb: LoadBalancer<Connection<API>>,
 }
 
 impl<API: Api> TestContext<API>
@@ -146,7 +151,7 @@ where
 
         let expect = (operation.clone(), result.clone());
         let _ = self.storage.expect.lock().unwrap().insert(expect);
-        assert_eq!(self.client_conn.execute(operation).await, result);
+        assert_eq!(self.lb.execute(operation).await, result);
     }
 
     async fn test_get(&self) {