cluster: define and implement cass_cluster_set_load_balance_rack_awar…

…e[_n]

This is an extension introduced by Scylla's fork of cpp-driver.
See: scylladb/cpp-driver@9691ec0

Note:
I removed this part of the docstring:
```
 * With empty local_rack and local_dc,  default local_dc and local_rack
 * is chosen from the first connected contact point,
 * and no remote hosts are considered in query plans.
 * If relying on this mechanism, be sure to use only contact
 * points from the local rack.
```

There are multiple reasons for this:
- this behaviour does not make much sense, and we should not mimic it IMO
- rust-driver does not behave like this
- this is not even true for cpp-driver

Why it's not true for cpp-driver:
If you carefully study the changes introduced to cpp-driver in the aforementioned
commit, you will notice that it's not possible for the driver to use rack aware
policy with an empty strings. This is because API functions reject
empty string, thus RackAwarePolicy object is never constructed in such case.
```
CassError cass_cluster_set_load_balance_rack_aware_n(CassCluster* cluster, const char* local_dc,
                                                   size_t local_dc_length,
                                                   const char* local_rack,
                                                   size_t local_rack_length) {
  if (local_dc == NULL || local_dc_length == 0 || local_rack == NULL || local_rack_length == 0) {
    return CASS_ERROR_LIB_BAD_PARAMS;
  }
  cluster->config().set_load_balancing_policy(new RackAwarePolicy(
      String(local_dc, local_dc_length), String(local_rack, local_rack_length)));
  return CASS_OK;
}
```

Why is this part of docstring included in cpp-driver then? No idea. Maybe,
because `cass_cluster_set_load_balance_dc_aware` mentions something similar
for empty (non-specified) dc. However, in this case it's true, since dc awareness is enabled
by default in cpp-driver. See the docstring:
```
 * Configures the cluster to use DC-aware load balancing.
 * For each query, all live nodes in a primary 'local' DC are tried first,
 * followed by any node from other DCs.
 *
 * <b>Note:</b> This is the default, and does not need to be called unless
 * switching an existing from another policy or changing settings.
 * Without further configuration, a default local_dc is chosen from the
 * first connected contact point, and no remote hosts are considered in
 * query plans. If relying on this mechanism, be sure to use only contact
 * points from the local DC.
```

include/cassandra.h

@@ -2219,6 +2219,46 @@ cass_cluster_set_load_balance_dc_aware_n(CassCluster* cluster,
                                          unsigned used_hosts_per_remote_dc,
                                          cass_bool_t allow_remote_dcs_for_local_cl);
 
+
+/**
+ * Configures the cluster to use Rack-aware load balancing.
+ * For each query, all live nodes in a primary 'local' rack are tried first,
+ * followed by nodes from local DC and then nodes from other DCs.
+ *
+ * @public @memberof CassCluster
+ *
+ * @param[in] cluster
+ * @param[in] local_dc The primary data center to try first
+ * @param[in] local_rack The primary rack to try first
+ * @return CASS_OK if successful, otherwise an error occurred
+ */
+CASS_EXPORT CassError
+cass_cluster_set_load_balance_rack_aware(CassCluster* cluster,
+                                       const char* local_dc,
+                                       const char* local_rack);
+
+
+/**
+ * Same as cass_cluster_set_load_balance_rack_aware(), but with lengths for string
+ * parameters.
+ *
+ * @public @memberof CassCluster
+ *
+ * @param[in] cluster
+ * @param[in] local_dc
+ * @param[in] local_dc_length
+ * @return same as cass_cluster_set_load_balance_dc_aware()
+ *
+ * @see cass_cluster_set_load_balance_dc_aware()
+ */
+CASS_EXPORT CassError
+cass_cluster_set_load_balance_rack_aware_n(CassCluster* cluster,
+                                         const char* local_dc,
+                                         size_t local_dc_length,
+                                         const char* local_rack,
+                                         size_t local_rack_length);
+
+
 /**
  * Configures the cluster to use token-aware request routing or not.
  *

scylla-rust-wrapper/extern/cassandra.h

@@ -1087,6 +1087,7 @@ cass_execution_profile_set_load_balance_dc_aware_n(CassExecProfile* profile,
                                                    unsigned used_hosts_per_remote_dc,
                                                    cass_bool_t allow_remote_dcs_for_local_cl);
 
+
 /**
  * Configures the execution profile to use token-aware request routing or not.
  *
@@ -2219,6 +2220,46 @@ cass_cluster_set_load_balance_dc_aware_n(CassCluster* cluster,
                                          unsigned used_hosts_per_remote_dc,
                                          cass_bool_t allow_remote_dcs_for_local_cl);
 
+
+/**
+ * Configures the cluster to use Rack-aware load balancing.
+ * For each query, all live nodes in a primary 'local' rack are tried first,
+ * followed by nodes from local DC and then nodes from other DCs.
+ *
+ * @public @memberof CassCluster
+ *
+ * @param[in] cluster
+ * @param[in] local_dc The primary data center to try first
+ * @param[in] local_rack The primary rack to try first
+ * @return CASS_OK if successful, otherwise an error occurred
+ */
+CASS_EXPORT CassError
+cass_cluster_set_load_balance_rack_aware(CassCluster* cluster,
+                                       const char* local_dc,
+                                       const char* local_rack);
+
+
+/**
+ * Same as cass_cluster_set_load_balance_rack_aware(), but with lengths for string
+ * parameters.
+ *
+ * @public @memberof CassCluster
+ *
+ * @param[in] cluster
+ * @param[in] local_dc
+ * @param[in] local_dc_length
+ * @return same as cass_cluster_set_load_balance_dc_aware()
+ *
+ * @see cass_cluster_set_load_balance_dc_aware()
+ */
+CASS_EXPORT CassError
+cass_cluster_set_load_balance_rack_aware_n(CassCluster* cluster,
+                                         const char* local_dc,
+                                         size_t local_dc_length,
+                                         const char* local_rack,
+                                         size_t local_rack_length);
+
+
 /**
  * Configures the cluster to use token-aware request routing or not.
  *

scylla-rust-wrapper/src/cluster.rs

@@ -55,9 +55,22 @@ impl LoadBalancingConfig {
             builder =
                 builder.enable_shuffling_replicas(self.token_aware_shuffling_replicas_enabled);
         }
-        if let NodeLocationPreference::Datacenter { local_dc } = self.node_location_preference {
-            builder = builder.prefer_datacenter(local_dc).permit_dc_failover(true)
+
+        match self.node_location_preference {
+            NodeLocationPreference::Datacenter { local_dc } => {
+                builder = builder.prefer_datacenter(local_dc).permit_dc_failover(true)
+            }
+            NodeLocationPreference::DatacenterAndRack {
+                local_dc,
+                local_rack,
+            } => {
+                builder = builder
+                    .prefer_datacenter_and_rack(local_dc, local_rack)
+                    .permit_dc_failover(true)
+            }
+            NodeLocationPreference::Any => {}
         }
+
         if self.latency_awareness_enabled {
             builder = builder.latency_awareness(self.latency_awareness_builder);
         }
@@ -79,7 +92,13 @@ impl Default for LoadBalancingConfig {
 #[derive(Clone, Debug)]
 pub(crate) enum NodeLocationPreference {
     Any,
-    Datacenter { local_dc: String },
+    Datacenter {
+        local_dc: String,
+    },
+    DatacenterAndRack {
+        local_dc: String,
+        local_rack: String,
+    },
 }
 
 #[derive(Clone)]
@@ -477,6 +496,68 @@ pub unsafe extern "C" fn cass_cluster_set_load_balance_dc_aware_n(
     )
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn cass_cluster_set_load_balance_rack_aware(
+    cluster_raw: *mut CassCluster,
+    local_dc_raw: *const c_char,
+    local_rack_raw: *const c_char,
+) -> CassError {
+    cass_cluster_set_load_balance_rack_aware_n(
+        cluster_raw,
+        local_dc_raw,
+        strlen(local_dc_raw),
+        local_rack_raw,
+        strlen(local_rack_raw),
+    )
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn cass_cluster_set_load_balance_rack_aware_n(
+    cluster_raw: *mut CassCluster,
+    local_dc_raw: *const c_char,
+    local_dc_length: size_t,
+    local_rack_raw: *const c_char,
+    local_rack_length: size_t,
+) -> CassError {
+    let cluster = ptr_to_ref_mut(cluster_raw);
+
+    set_load_balance_rack_aware_n(
+        &mut cluster.load_balancing_config,
+        local_dc_raw,
+        local_dc_length,
+        local_rack_raw,
+        local_rack_length,
+    )
+}
+
+unsafe fn set_load_balance_rack_aware_n(
+    load_balancing_config: &mut LoadBalancingConfig,
+    local_dc_raw: *const c_char,
+    local_dc_length: size_t,
+    local_rack_raw: *const c_char,
+    local_rack_length: size_t,
+) -> CassError {
+    let (local_dc, local_rack) = match (
+        ptr_to_cstr_n(local_dc_raw, local_dc_length),
+        ptr_to_cstr_n(local_rack_raw, local_rack_length),
+    ) {
+        (Some(local_dc_str), Some(local_rack_str))
+            if local_dc_length > 0 && local_rack_length > 0 =>
+        {
+            (local_dc_str.to_owned(), local_rack_str.to_owned())
+        }
+        // One of them is either null pointer, is an empty string or is not a proper utf-8.
+        _ => return CassError::CASS_ERROR_LIB_BAD_PARAMS,
+    };
+
+    load_balancing_config.node_location_preference = NodeLocationPreference::DatacenterAndRack {
+        local_dc,
+        local_rack,
+    };
+
+    CassError::CASS_OK
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn cass_cluster_set_cloud_secure_connection_bundle_n(
     _cluster_raw: *mut CassCluster,