Merge pull request #213 from cognitedata/clientbuilder-accept-backoffconfig

refactor: Clientbuilder accept backoffconfig

Author: kodiakhq[bot]

src/backoff.rs

@@ -6,6 +6,7 @@ use tracing::info;
 /// Exponential backoff with jitter
 ///
 /// See <https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/>
+#[allow(missing_copy_implementations)]
 #[derive(Debug, Clone)]
 pub struct BackoffConfig {
     pub init_backoff: Duration,

src/client/controller.rs

@@ -27,17 +27,17 @@ use super::error::RequestContext;
 pub struct ControllerClient {
     brokers: Arc<BrokerConnector>,
 
-    backoff_config: BackoffConfig,
+    backoff_config: Arc<BackoffConfig>,
 
     /// Current broker connection if any
     current_broker: Mutex<(Option<BrokerConnection>, BrokerCacheGeneration)>,
 }
 
 impl ControllerClient {
-    pub(super) fn new(brokers: Arc<BrokerConnector>) -> Self {
+    pub(super) fn new(brokers: Arc<BrokerConnector>, backoff_config: Arc<BackoffConfig>) -> Self {
         Self {
             brokers,
-            backoff_config: Default::default(),
+            backoff_config,
             current_broker: Mutex::new((None, BrokerCacheGeneration::START)),
         }
     }

src/client/mod.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use thiserror::Error;
 
 use crate::{
+    backoff::BackoffConfig,
     build_info::DEFAULT_CLIENT_ID,
     client::partition::PartitionClient,
     connection::{BrokerConnector, MetadataLookupMode, TlsConfig},
@@ -43,6 +44,7 @@ pub struct ClientBuilder {
     max_message_size: usize,
     socks5_proxy: Option<String>,
     tls_config: TlsConfig,
+    backoff_config: Arc<BackoffConfig>,
 }
 
 impl ClientBuilder {
@@ -54,6 +56,7 @@ impl ClientBuilder {
             max_message_size: 100 * 1024 * 1024, // 100MB
             socks5_proxy: None,
             tls_config: TlsConfig::default(),
+            backoff_config: Default::default(),
         }
     }
 
@@ -73,6 +76,12 @@ impl ClientBuilder {
         self
     }
 
+    /// Set up backoff configuration
+    pub fn backoff_config(mut self, backoff_config: BackoffConfig) -> Self {
+        self.backoff_config = Arc::from(backoff_config);
+        self
+    }
+
     /// Use SOCKS5 proxy.
     #[cfg(feature = "transport-socks5")]
     pub fn socks5_proxy(mut self, proxy: String) -> Self {
@@ -96,10 +105,14 @@ impl ClientBuilder {
             self.tls_config,
             self.socks5_proxy,
             self.max_message_size,
+            Arc::clone(&self.backoff_config),
         ));
         brokers.refresh_metadata().await?;
 
-        Ok(Client { brokers })
+        Ok(Client {
+            brokers,
+            backoff_config: self.backoff_config,
+        })
     }
 }
 
@@ -118,12 +131,16 @@ impl std::fmt::Debug for ClientBuilder {
 #[derive(Debug)]
 pub struct Client {
     brokers: Arc<BrokerConnector>,
+    backoff_config: Arc<BackoffConfig>,
 }
 
 impl Client {
     /// Returns a client for performing certain cluster-wide operations.
     pub fn controller_client(&self) -> Result<ControllerClient> {
-        Ok(ControllerClient::new(Arc::clone(&self.brokers)))
+        Ok(ControllerClient::new(
+            Arc::clone(&self.brokers),
+            Arc::clone(&self.backoff_config),
+        ))
     }
 
     /// Returns a client for performing operations on a specific partition
@@ -138,6 +155,7 @@ impl Client {
             partition,
             Arc::clone(&self.brokers),
             unknown_topic_handling,
+            Arc::clone(&self.backoff_config),
         )
         .await
     }

src/client/partition.rs

@@ -123,7 +123,7 @@ pub struct PartitionClient {
     partition: i32,
     brokers: Arc<BrokerConnector>,
 
-    backoff_config: BackoffConfig,
+    backoff_config: Arc<BackoffConfig>,
 
     /// Current broker connection if any
     current_broker: Mutex<CurrentBroker>,
@@ -143,12 +143,13 @@ impl PartitionClient {
         partition: i32,
         brokers: Arc<BrokerConnector>,
         unknown_topic_handling: UnknownTopicHandling,
+        backoff_config: Arc<BackoffConfig>,
     ) -> Result<Self> {
         let p = Self {
             topic,
             partition,
             brokers: Arc::clone(&brokers),
-            backoff_config: Default::default(),
+            backoff_config,
             current_broker: Mutex::new(CurrentBroker {
                 broker: None,
                 gen_broker: BrokerCacheGeneration::START,

src/connection.rs

@@ -168,7 +168,7 @@ pub struct BrokerConnector {
     cached_metadata: MetadataCache,
 
     /// The backoff configuration on error
-    backoff_config: BackoffConfig,
+    backoff_config: Arc<BackoffConfig>,
 
     /// TLS configuration if any
     tls_config: TlsConfig,
@@ -187,14 +187,15 @@ impl BrokerConnector {
         tls_config: TlsConfig,
         socks5_proxy: Option<String>,
         max_message_size: usize,
+        backoff_config: Arc<BackoffConfig>,
     ) -> Self {
         Self {
             bootstrap_brokers,
             client_id,
             topology: Default::default(),
             cached_arbitrary_broker: Mutex::new((None, BrokerCacheGeneration::START)),
             cached_metadata: Default::default(),
-            backoff_config: Default::default(),
+            backoff_config,
             tls_config,
             socks5_proxy,
             max_message_size,

src/lib.rs

@@ -19,6 +19,8 @@
 
 mod backoff;
 
+pub use backoff::BackoffConfig;
+
 pub mod build_info;
 
 pub mod client;