From 5b8074d4b557fb3b244cbb40aaf4e5d4f32fb770 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20Uzarski?= <mikolaj.uzarski@scylladb.com>
Date: Wed, 30 Oct 2024 15:28:12 +0100
Subject: [PATCH] sbuilder: add hostname_resolution_timeout option

This option is responsible for setting the DNS hostname
resolution timeout.

The default timeout is set to 5 seconds.
---
 scylla/src/transport/session.rs         |  5 +++++
 scylla/src/transport/session_builder.rs | 21 +++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index 42933bdaf6..43fce28a38 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -252,6 +252,10 @@ pub struct SessionConfig {
     /// It is true by default but can be disabled if successive schema-altering statements should be performed.
     pub refresh_metadata_on_auto_schema_agreement: bool,
 
+    /// DNS hostname resolution timeout.
+    /// If `None`, the driver will wait for hostname resolution indefinitely.
+    pub hostname_resolution_timeout: Option<Duration>,
+
     /// The address translator is used to translate addresses received from ScyllaDB nodes
     /// (either with cluster metadata or with an event) to addresses that can be used to
     /// actually connect to those nodes. This may be needed e.g. when there is NAT
@@ -332,6 +336,7 @@ impl SessionConfig {
             ssl_context: None,
             authenticator: None,
             connect_timeout: Duration::from_secs(5),
+            hostname_resolution_timeout: Some(Duration::from_secs(5)),
             connection_pool_size: Default::default(),
             disallow_shard_aware_port: false,
             keyspaces_to_fetch: Vec::new(),
diff --git a/scylla/src/transport/session_builder.rs b/scylla/src/transport/session_builder.rs
index 9a7a9cbf71..50b3bdb76b 100644
--- a/scylla/src/transport/session_builder.rs
+++ b/scylla/src/transport/session_builder.rs
@@ -762,6 +762,27 @@ impl<K: SessionBuilderKind> GenericSessionBuilder<K> {
         self
     }
 
+    /// Changes DNS hostname resolution timeout.
+    /// The default is 5 seconds.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # use std::time::Duration;
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .hostname_resolution_timeout(Duration::from_secs(10))
+    ///     .build() // Turns SessionBuilder into Session
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn hostname_resolution_timeout(mut self, duration: Duration) -> Self {
+        self.config.hostname_resolution_timeout = Some(duration);
+        self
+    }
+
     /// Sets the host filter. The host filter decides whether any connections
     /// should be opened to the node or not. The driver will also avoid
     /// those nodes when re-establishing the control connection.