Biased load balancing

pgdog/src/backend/pool/cluster.rs

@@ -98,7 +98,9 @@ impl<'a> ClusterConfig<'a> {
             user: &user.name,
             replication_sharding: user.replication_sharding.clone(),
             pooler_mode: user.pooler_mode.unwrap_or(general.pooler_mode),
-            lb_strategy: general.load_balancing_strategy,
+            lb_strategy: user
+                .load_balancing_strategy
+                .unwrap_or(general.load_balancing_strategy),
             shards,
             sharded_tables,
             mirror_of,

pgdog/src/backend/pool/replicas.rs

@@ -134,6 +134,15 @@ impl Replicas {
                 LeastActiveConnections => {
                     candidates.sort_by_cached_key(|pool| pool.lock().idle());
                 }
+                PrimaryOnlyWithFailover => (), // Leave the current order. Primary will be attempted first.
+                ReplicasOnlyWithFailover => {
+                    if primary.is_some() {
+                        let mut reshuffled = vec![];
+                        reshuffled.extend_from_slice(&candidates[1..]);
+                        reshuffled.push(candidates[0]);
+                        candidates = reshuffled;
+                    }
+                }
             }
 
             let mut banned = 0;

pgdog/src/config/mod.rs

@@ -574,13 +574,17 @@ impl std::fmt::Display for PoolerMode {
     }
 }
 
-#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)]
+#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy, Eq, PartialOrd, Ord)]
 #[serde(rename_all = "snake_case")]
 pub enum LoadBalancingStrategy {
     #[default]
     Random,
     RoundRobin,
     LeastActiveConnections,
+    /// Use replicas for queries unless they are all down.
+    ReplicasOnlyWithFailover,
+    /// Use primary for queries unless it's banned.
+    PrimaryOnlyWithFailover,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)]
@@ -744,6 +748,9 @@ pub struct User {
     pub idle_timeout: Option<u64>,
     /// Read-only mode.
     pub read_only: Option<bool>,
+    /// Load balancing strategy.
+    #[serde(default)]
+    pub load_balancing_strategy: Option<LoadBalancingStrategy>,
 }
 
 impl User {