Added Smart driver features (yugabyte#1)

* Smart Driver Implementation

* Smart Driver Implementation

* Changes as per review comments

* Changes as per review

Author: HarshDaryani896

README.md

@@ -1,4 +1,74 @@
-# Rust-Postgres
+# YugabyteDB Rust-Postgres Smart Driver
+
+This is a Rust driver for YugabyteDB SQL. This driver is based on [sfackler/rust-postgres](https://github.com/sfackler/rust-postgres), with the following additional connection properties:
+
+- `load_balance`   - It expects **true/false** as its possible values.
+- `topology_keys`  - It takes a comma separated geo-location values. A single geo-location can be given as 'cloud.region.zone'. Multiple geo-locations too can be specified, separated by comma (`,`).
+- `yb_servers_refresh_interval` - (default value: 300 seconds) Time interval, in seconds, between two attempts to refresh the information about cluster nodes. Valid values are integers between 0 and 600. Value 0 means refresh for each connection request. Any value outside this range is ignored and the default is used.
+- `fallback_to_topology_keys_only` : (default value: false) Applicable only for TopologyAware Load Balancing. When set to true, the smart driver does not attempt to connect to servers outside of primary and fallback placements specified via property. The default behaviour is to fallback to any available server in the entire cluster.
+- `failed_host_reconnect_delay_secs` - (default value: 5 seconds) The driver marks a server as failed with a timestamp, when it cannot connect to it. Later, whenever it refreshes the server list via yb_servers(), if it sees the failed server in the response, it marks the server as UP only if failed-host-reconnect-delay-secs time has elapsed.
+
+## Connection load balancing
+
+Users can use this feature in two configurations.
+
+### Cluster-aware / Uniform connection load balancing
+
+In the cluster-aware connection load balancing, connections are distributed across all the tservers in the cluster, irrespective of their placements.
+
+To enable the cluster-aware connection load balancing, provide the parameter `load_balance` set to true as `load_balance=true` in the connection url.
+
+```
+"postgresql://127.0.0.1:5433/yugabyte?user=yugabyte&password=yugabyte&load_balance=true"
+```
+
+With this parameter specified in the url, the driver will fetch and maintain the list of tservers from the given endpoint (`127.0.0.1` in above example) available in the YugabyteDB cluster and distribute the connections equally across them.
+
+This list is refreshed every 5 minutes, when a new connection request is received.
+
+Application needs to use the same connection url to create every connection it needs, so that the distribution happens equally.
+
+### Topology-aware connection load balancing
+
+With topology-aware connnection load balancing, users can target tservers in specific zones by specifying these zones as `topology_keys` with values in the format `cloudname.regionname.zonename`. Multiple zones can be specified as comma separated values.
+
+The connections will be distributed equally with the tservers in these zones.
+
+Note that, you would still need to specify `load_balance=true` to enable the topology-aware connection load balancing.
+
+```
+"postgresql://127.0.0.1:5433/yugabyte?user=yugabyte&password=yugabyte&load_balance=true&topology_keys=cloud1.datacenter1.rack1"
+```
+### Specifying fallback zones
+
+For topology-aware load balancing, you can now specify fallback placements too. This is not applicable for cluster-aware load balancing.
+Each placement value can be suffixed with a colon (`:`) followed by a preference value between 1 and 10.
+A preference value of `:1` means it is a primary placement. A preference value of `:2` means it is the first fallback placement and so on. If no preference value is provided, it is considered to be a primary placement (equivalent to one with preference value `:1`). Example given below.
+
+```
+"postgresql://127.0.0.1:5433/yugabyte?user=yugabyte&password=yugabyte&load_balance=true&topology_keys=cloud1.region1.zone1:1,cloud1.region1.zone2:2";
+```
+
+You can also use `*` for specifying all the zones in a given region as shown below. This is not allowed for cloud or region values.
+
+```
+"postgresql://127.0.0.1:5433/yugabyte?user=yugabyte&password=yugabyte&load_balance=true&topology_keys=cloud1.region1.*:1,cloud1.region2.*:2";
+```
+
+The driver attempts to connect to a node in following order: the least loaded node in the 1) primary placement(s), else in the 2) first fallback if specified, else in the 3) second fallback if specified and so on.
+If no nodes are available either in primary placement(s) or in any of the fallback placements, then nodes in the rest of the cluster are attempted.
+And this repeats for each connection request.
+
+## Specifying Refresh Interval
+
+Users can specify Refresh Time Interval, in seconds. It is the time interval between two attempts to refresh the information about cluster nodes. Default is 300. Valid values are integers between 0 and 600. Value 0 means refresh for each connection request. Any value outside this range is ignored and the default is used.
+
+To specify Refresh Interval, use the parameter `yb_servers_refresh_interval` in the connection url.
+```
+"postgresql://127.0.0.1:5433/yugabyte?user=yugabyte&password=yugabyte&yb_servers_refresh_interval=0&load_balance=true&topology_keys=cloud1.region1.*:1,cloud1.region2.*:2";
+```
+
+For working examples which demonstrates both the configurations of connection load balancing using `tokio-postgres::connect()`, see the [driver-examples](https://github.com/yugabyte/driver-examples/tree/main/rust/rust_ysql_driver_examples) repository.
 
 PostgreSQL support for Rust.
 

postgres/src/client.rs

@@ -7,7 +7,7 @@ use std::task::Poll;
 use std::time::Duration;
 use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
 use tokio_postgres::types::{BorrowToSql, ToSql, Type};
-use tokio_postgres::{Error, Row, SimpleQueryMessage, Socket};
+use tokio_postgres::{Error, Row, SimpleQueryMessage, Socket, close};
 
 /// A synchronous PostgreSQL client.
 pub struct Client {
@@ -17,6 +17,7 @@ pub struct Client {
 
 impl Drop for Client {
     fn drop(&mut self) {
+        close(&self.client);
         let _ = self.close_inner();
     }
 }

postgres/src/config.rs

@@ -3,6 +3,7 @@
 use crate::connection::Connection;
 use crate::Client;
 use log::info;
+use std::collections::HashMap;
 use std::fmt;
 use std::net::IpAddr;
 use std::path::Path;
@@ -81,6 +82,22 @@ use tokio_postgres::{Error, Socket};
 ///     `disable`, hosts and addresses will be tried in the order provided. If set to `random`, hosts will be tried
 ///     in a random order, and the IP addresses resolved from a hostname will also be tried in a random order. Defaults
 ///     to `disable`.
+/// * `load_balance` -  It expects true/false as its possible values. Default value is true.
+/// * `topology_keys` - It takes a comma separated geo-location values. A single geo-location can be given as 'cloud.region.zone'.
+///     Multiple geo-locations too can be specified, separated by comma (,). Each placement value can be suffixed with a colon (:)
+///     followed by a preference value between 1 and 10. A preference value of :1 means it is a primary placement. A preference
+///     value of :2 means it is the first fallback placement and so on. If no preference value is provided, it is considered to
+///     be a primary placement (equivalent to one with preference value :1).
+/// * `yb_servers_refresh_interval` - Time interval, in seconds, between two attempts to refresh the information about cluster nodes.
+///     Default is 300. Valid values are integers between 0 and 600. Value 0 means refresh for each connection request. Any value
+///     outside this range is ignored and the default is used.
+/// * `fallback_to_topology_keys_only` - (default value: false) Applicable only for TopologyAware Load Balancing. When set to true,
+///     the smart driver does not attempt to connect to servers outside of primary and fallback placements specified via property.
+///     The default behaviour is to fallback to any available server in the entire cluster.
+/// * `failed_host_reconnect_delay_secs` - (default value: 5 seconds) The driver marks a server as failed with a timestamp, when it cannot
+///     connect to it. Later, whenever it refreshes the server list via yb_servers(), if it sees the failed server in the response,
+///     it marks the server as UP only if failed-host-reconnect-delay-secs time has elapsed. (The yb_servers() function does not remove
+///     a failed server immediately from its result and retains it for a while.)
 ///
 /// ## Examples
 ///
@@ -414,6 +431,103 @@ impl Config {
         self.config.get_load_balance_hosts()
     }
 
+    /// YugabyteDB Specific.
+    ///
+    /// Sets the load balance parameter.
+    ///
+    /// Defaults to false.
+    pub fn load_balance(&mut self, load_balance: bool) -> &mut Config {
+        self.config.load_balance(load_balance);
+        self
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Gets the load balance value
+    pub fn get_load_balance(&self) -> bool {
+        self.config.get_load_balance()
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Sets the topology key parameter.
+    ///
+    /// Defaults to Hashmap::new().
+    pub fn topology_keys(&mut self, topology_key: &str, priority: i64) -> &mut Config {
+        self.config.topology_keys(topology_key, priority);
+        self
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Gets the host topology keys value.
+    pub fn get_topology_keys(&self) -> HashMap<i64, Vec<String>> {
+        self.config.get_topology_keys()
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Sets the yb_servers_refresh_interval parameter.
+    ///
+    /// Defaults to 300 sec.
+    pub fn yb_servers_refresh_interval(
+        &mut self,
+        yb_servers_refresh_interval: Duration,
+    ) -> &mut Config {
+        self.config
+            .yb_servers_refresh_interval(yb_servers_refresh_interval);
+        self
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Gets the yb_servers_refresh_interval value.
+    pub fn get_yb_servers_refresh_interval(&self) -> Duration {
+        self.config.get_yb_servers_refresh_interval()
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Sets the fallback_to_topology_keys_only parameter.
+    ///
+    /// Defaults to false.
+    pub fn fallback_to_topology_keys_only(
+        &mut self,
+        fallback_to_topology_keys_only: bool,
+    ) -> &mut Config {
+        self.config
+            .fallback_to_topology_keys_only(fallback_to_topology_keys_only);
+        self
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Gets the fallback_to_topology_keys_only value.
+    pub fn get_fallback_to_topology_keys_only(&self) -> bool {
+        self.config.get_fallback_to_topology_keys_only()
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Sets the failed_host_reconnect_delay_secs parameter.
+    ///
+    /// Defaults to 5 sec.
+    pub fn failed_host_reconnect_delay_secs(
+        &mut self,
+        failed_host_reconnect_delay_secs: Duration,
+    ) -> &mut Config {
+        self.config
+            .failed_host_reconnect_delay_secs(failed_host_reconnect_delay_secs);
+        self
+    }
+
+    /// YugabyteDB Specific.
+    ///
+    /// Gets the failed_host_reconnect_delay_secs value.
+    pub fn get_failed_host_reconnect_delay_secs(&self) -> Duration {
+        self.config.get_failed_host_reconnect_delay_secs()
+    }
+
     /// Sets the notice callback.
     ///
     /// This callback will be invoked with the contents of every

tokio-postgres/Cargo.toml

@@ -60,6 +60,8 @@ tokio = { version = "1.27", features = ["io-util"] }
 tokio-util = { version = "0.7", features = ["codec"] }
 rand = "0.8.5"
 whoami = "1.4.1"
+lazy_static = "1.4.0"
+env_logger = "0.10.0"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 socket2 = { version = "0.5", features = ["all"] }

tokio-postgres/src/client.rs

@@ -214,6 +214,10 @@ impl Client {
         self.socket_config = Some(socket_config);
     }
 
+    pub(crate) fn get_socket_config(&self) -> Option<SocketConfig> {
+        self.socket_config.clone()
+    }
+
     /// Creates a new prepared statement.
     ///
     /// Prepared statements can be executed repeatedly, and may contain query parameters (indicated by `$1`, `$2`, etc),