From a03d9cd1002b52ae85ebbd2dca452a078de8e40e Mon Sep 17 00:00:00 2001
From: Lev Kokotov <lev.kokotov@gmail.com>
Date: Mon, 16 Jun 2025 18:05:34 -0700
Subject: [PATCH] Biased load balancing

---
 pgdog/src/backend/pool/cluster.rs  | 4 +++-
 pgdog/src/backend/pool/replicas.rs | 9 +++++++++
 pgdog/src/config/mod.rs            | 9 ++++++++-
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs
index b2c5f098..3eb1843f 100644
--- a/pgdog/src/backend/pool/cluster.rs
+++ b/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,
diff --git a/pgdog/src/backend/pool/replicas.rs b/pgdog/src/backend/pool/replicas.rs
index 39919762..6419b0c4 100644
--- a/pgdog/src/backend/pool/replicas.rs
+++ b/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;
diff --git a/pgdog/src/config/mod.rs b/pgdog/src/config/mod.rs
index ea4d8cca..f373a6f3 100644
--- a/pgdog/src/config/mod.rs
+++ b/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 {