Refactor/curp client tests suits

curp/src/client_new/mod.rs

@@ -7,18 +7,27 @@
 /// Retry layer
 mod retry;
 
+/// Tests for client
+#[cfg(test)]
+mod tests;
+
 use std::{collections::HashMap, sync::Arc};
 
 use async_trait::async_trait;
 use curp_external_api::cmd::Command;
 use futures::{stream::FuturesUnordered, StreamExt};
 use utils::config::ClientConfig;
 
+use self::{
+    retry::Retry,
+    unary::{UnaryBuilder, UnaryConfig},
+};
 use crate::{
+    client_new::retry::RetryConfig,
     members::ServerId,
     rpc::{
         connect::ConnectApi, protocol_client::ProtocolClient, ConfChange, FetchClusterRequest,
-        FetchClusterResponse, Member, ReadState,
+        FetchClusterResponse, Member, Protocol, ReadState,
     },
 };
 
@@ -76,17 +85,17 @@
     #[inline]
     async fn fetch_leader_id(&self, linearizable: bool) -> Result<ServerId, Self::Error> {
         if linearizable {
             let resp = self.fetch_cluster(true).await?;
             return Ok(resp.leader_id.unwrap_or_else(|| {
                 unreachable!("linearizable fetch cluster should return a leader id")
             }));
         }
         let resp = self.fetch_cluster(false).await?;
         if let Some(id) = resp.leader_id {
             return Ok(id);
         }
         // fallback to linearizable fetch
         self.fetch_leader_id(true).await
     }
 }
 
@@ -98,94 +107,151 @@
 }
 
 /// Client builder to build a client
 #[derive(Debug, Clone)]
-pub struct ClientBuilder {
-    /// local server id
-    local_server_id: Option<ServerId>,
+pub struct ClientBuilder<P: Protocol> {
+    /// local server
+    local_server: Option<(ServerId, P)>,
     /// initial cluster version
     cluster_version: Option<u64>,
+    /// initial cluster members
+    all_members: Option<HashMap<ServerId, Vec<String>>>,
+    /// initial leader state
+    leader_state: Option<(ServerId, u64)>,
     /// client configuration
     config: ClientConfig,
-    /// initial all members
-    all_members: Option<HashMap<ServerId, Vec<String>>>,
 }
 
-impl ClientBuilder {
+impl<P: Protocol> ClientBuilder<P> {
     /// Create a client builder
     #[inline]
     #[must_use]
     pub fn new(config: ClientConfig) -> Self {
         Self {
-            local_server_id: None,
-            config,
-            all_members: None,
+            local_server: None,
             cluster_version: None,
+            all_members: None,
+            leader_state: None,
+            config,
         }
     }
 
-    /// Set the local server id
+    /// Set the local server to bypass `gRPC` request
     #[inline]
     #[must_use]
-    pub fn local_server_id(&mut self, id: ServerId) -> &mut Self {
-        self.local_server_id = Some(id);
+    pub fn bypass(&mut self, id: ServerId, server: P) -> &mut Self {
+        self.local_server = Some((id, server));
         self
     }
 
     /// Set the initial cluster version
     #[inline]
     #[must_use]
     pub fn cluster_version(&mut self, cluster_version: u64) -> &mut Self {
         self.cluster_version = Some(cluster_version);
         self
     }
 
-    /// Fetch initial all members from some endpoints if you do not know the whole members
+    /// Set the initial all members
+    #[inline]
+    #[must_use]
+    pub fn all_members(&mut self, all_members: HashMap<ServerId, Vec<String>>) -> &mut Self {
+        self.all_members = Some(all_members);
+        self
+    }
+
+    /// Set the initial leader state
+    #[inline]
+    #[must_use]
+    pub fn leader_state(&mut self, leader_id: ServerId, term: u64) -> &mut Self {
+        self.leader_state = Some((leader_id, term));
+        self
+    }
+
+    /// Discover the initial states from some endpoints
     ///
     /// # Errors
     ///
     /// Return `tonic::Status` for connection failure or some server errors.
     #[inline]
-    pub async fn fetch_all_members(
-        &mut self,
-        addrs: Vec<String>,
-    ) -> Result<&mut Self, tonic::Status> {
+    pub async fn discover_from(&mut self, addrs: Vec<String>) -> Result<&mut Self, tonic::Status> {
         let propose_timeout = *self.config.propose_timeout();
         let mut futs: FuturesUnordered<_> = addrs
             .into_iter()
             .map(|mut addr| {
                 if !addr.starts_with("http://") {
                     addr.insert_str(0, "http://");
                 }
                 async move {
                     let mut protocol_client = ProtocolClient::connect(addr).await.map_err(|e| {
                         tonic::Status::cancelled(format!("cannot connect to addr, error: {e}"))
                     })?;
                     let mut req = tonic::Request::new(FetchClusterRequest::default());
                     req.set_timeout(propose_timeout);
                     let fetch_cluster_res = protocol_client.fetch_cluster(req).await?.into_inner();
                     Ok::<FetchClusterResponse, tonic::Status>(fetch_cluster_res)
                 }
             })
             .collect();
         let mut err = tonic::Status::invalid_argument("addrs is empty");
+        // find the first one return `FetchClusterResponse`
         while let Some(r) = futs.next().await {
             match r {
                 Ok(r) => {
                     self.cluster_version = Some(r.cluster_version);
+                    if let Some(id) = r.leader_id {
+                        self.leader_state = Some((id, r.term));
+                    }
                     self.all_members = Some(r.into_members_addrs());
                     return Ok(self);
                 }
                 Err(e) => err = e,
             }
         }
         Err(err)
     }
 
-    /// Set the initial all members
+    /// Build the client
+    ///
+    /// # Errors
+    ///
+    /// Return `tonic::transport::Error` for connection failure.
     #[inline]
-    #[must_use]
-    pub fn set_all_members(&mut self, all_members: HashMap<ServerId, Vec<String>>) -> &mut Self {
-        self.all_members = Some(all_members);
-        self
+    pub async fn build<C: Command>(
+        self,
+    ) -> Result<impl ClientApi<Error = tonic::Status, Cmd = C>, tonic::transport::Error> {
+        let mut builder = UnaryBuilder::<P>::new(
+            self.all_members.unwrap_or_else(|| {
+                unreachable!("must set the initial members or discover from some endpoints")
+            }),
+            UnaryConfig::new_full(
+                *self.config.propose_timeout(),
+                *self.config.wait_synced_timeout(),
+            ),
+        );
+        if let Some(version) = self.cluster_version {
+            builder.set_cluster_version(version);
+        }
+        if let Some((id, server)) = self.local_server {
+            builder.set_local_server(id, server);
+        }
+        if let Some((id, term)) = self.leader_state {
+            builder.set_leader_state(id, term);
+        }
+        let unary = builder.build::<C>().await?;
+
+        let retry_config = if *self.config.use_backoff() {
+            RetryConfig::new_exponential(
+                *self.config.initial_retry_timeout(),
+                *self.config.max_retry_timeout(),
+                *self.config.retry_count(),
+            )
+        } else {
+            RetryConfig::new_fixed(
+                *self.config.initial_retry_timeout(),
+                *self.config.retry_count(),
+            )
+        };
+        let client = Retry::new(unary, retry_config);
+        Ok(client)
     }
 }

curp/src/client_new/retry.rs

@@ -19,7 +19,7 @@
 
 /// Backoff config
 #[derive(Debug, Clone)]
-pub(super) enum BackoffConfig {
+enum BackoffConfig {
     /// A fixed delay backoff
     Fixed,
     /// A exponential delay backoff
@@ -41,7 +41,7 @@
 }
 
 /// Backoff tool
 #[derive(Debug)]
 struct Backoff {
     /// The retry config
     config: RetryConfig,
@@ -53,7 +53,7 @@
 
 impl RetryConfig {
     /// Create a fixed retry config
-    fn new_fixed(delay: Duration, count: usize) -> Self {
+    pub(super) fn new_fixed(delay: Duration, count: usize) -> Self {
         assert!(count > 0, "retry count should be larger than 0");
         Self {
             backoff: BackoffConfig::Fixed,
@@ -63,7 +63,7 @@
     }
 
     /// Create a exponential retry config
-    fn new_exponential(delay: Duration, max_delay: Duration, count: usize) -> Self {
+    pub(super) fn new_exponential(delay: Duration, max_delay: Duration, count: usize) -> Self {
         assert!(count > 0, "retry count should be larger than 0");
         Self {
             backoff: BackoffConfig::Exponential { max_delay },
@@ -103,7 +103,7 @@
 
 /// The retry client automatically retry the requests of the inner client api
 /// which raises the [`tonic::Status`] error
 #[derive(Debug)]
 pub(super) struct Retry<Api> {
     /// Inner client
     inner: Api,
@@ -116,7 +116,7 @@
     Api: ClientApi<Error = CurpError> + LeaderStateUpdate + Send + Sync + 'static,
 {
     /// Create a retry client
-    fn new(inner: Api, config: RetryConfig) -> Self {
+    pub(super) fn new(inner: Api, config: RetryConfig) -> Self {
         Self { inner, config }
     }
 
@@ -128,7 +128,7 @@
         let mut backoff = self.config.init_backoff();
         while let Some(delay) = backoff.next_delay() {
             let err = match f(&self.inner).await {
                 Ok(res) => return Ok(res),
                 Err(err) => err,
             };
 
@@ -152,15 +152,15 @@
                 // update the cluster state if got WrongClusterVersion
                 CurpError::WrongClusterVersion(_) => {
                     // the inner client should automatically update cluster state when fetch_cluster
                     if let Err(e) = self.inner.fetch_cluster(false).await {
                         warn!("fetch cluster failed, error {e:?}");
                     }
                 }
 
                 // update the leader state if got Redirect
                 CurpError::Redirect(Redirect { leader_id, term }) => {
                     self.inner.update_leader(leader_id, term);
                 }
             }
 
             warn!("retry on {} seconds later", delay.as_secs_f32());
@@ -183,9 +183,9 @@
     type Cmd = Api::Cmd;
 
     /// Get the local connection when the client is on the server node.
     async fn local_connect(&self) -> Option<Arc<dyn ConnectApi>> {
         self.inner.local_connect().await
     }
 
     /// Send propose to the whole cluster, `use_fast_path` set to `false` to fallback into ordered
     /// requests (event the requests are commutative).
@@ -199,49 +199,79 @@
     }
 
     /// Send propose configuration changes to the cluster
     async fn propose_conf_change(
         &self,
         changes: Vec<ConfChange>,
     ) -> Result<Vec<Member>, tonic::Status> {
         self.retry::<_, _>(|client| {
             let changes_c = changes.clone();
             async move { client.propose_conf_change(changes_c).await }
         })
         .await
     }
 
     /// Send propose to shutdown cluster
     async fn propose_shutdown(&self) -> Result<(), tonic::Status> {
         self.retry::<_, _>(ClientApi::propose_shutdown).await
     }
 
     /// Send propose to publish a node id and name
     async fn propose_publish(
         &self,
         node_id: ServerId,
         node_name: String,
     ) -> Result<(), Self::Error> {
         self.retry::<_, _>(|client| {
             let name_c = node_name.clone();
             async move { client.propose_publish(node_id, name_c).await }
         })
         .await
     }
 
     /// Send fetch read state from leader
     async fn fetch_read_state(&self, cmd: &Self::Cmd) -> Result<ReadState, tonic::Status> {
         self.retry::<_, _>(|client| client.fetch_read_state(cmd))
             .await
     }
 
     /// Send fetch cluster requests to all servers (That's because initially, we didn't
     /// know who the leader is.)
     /// Note: The fetched cluster may still be outdated if `linearizable` is false
     async fn fetch_cluster(
         &self,
         linearizable: bool,
     ) -> Result<FetchClusterResponse, tonic::Status> {
         self.retry::<_, _>(|client| client.fetch_cluster(linearizable))
             .await
     }
 }
+
+/// Tests for backoff
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use super::RetryConfig;
+
+    #[test]
+    fn test_fixed_backoff_works() {
+        let config = RetryConfig::new_fixed(Duration::from_secs(1), 3);
+        let mut backoff = config.init_backoff();
+        assert_eq!(backoff.next_delay(), Some(Duration::from_secs(1)));
+        assert_eq!(backoff.next_delay(), Some(Duration::from_secs(1)));
+        assert_eq!(backoff.next_delay(), Some(Duration::from_secs(1)));
+        assert_eq!(backoff.next_delay(), None);
+    }
+
+    #[test]
+    fn test_exponential_backoff_works() {
+        let config =
+            RetryConfig::new_exponential(Duration::from_secs(1), Duration::from_secs(5), 4);
+        let mut backoff = config.init_backoff();
+        assert_eq!(backoff.next_delay(), Some(Duration::from_secs(1)));
+        assert_eq!(backoff.next_delay(), Some(Duration::from_secs(2)));
+        assert_eq!(backoff.next_delay(), Some(Duration::from_secs(4)));
+        assert_eq!(backoff.next_delay(), Some(Duration::from_secs(5))); // 8 > 5
+        assert_eq!(backoff.next_delay(), None);
+    }
+}