Unban address in the background when ban_time exceeded

Author: magec

src/ban_service.rs

@@ -4,12 +4,20 @@ use log::{debug, error, warn};
 use parking_lot::RwLock;
 use std::sync::Arc;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use crate::config::{Address, Role};
+use crate::pool::ServerPool;
+use bb8::Pool;
+
 pub type BanList = Arc<RwLock<Vec<HashMap<Address, (BanReason, NaiveDateTime)>>>>;
 #[derive(Debug, Clone, Default)]
 pub struct BanService {
+    /// A set of addresses that are being unbanned.
+    /// This is used to prevent multiple unbanning tasks
+    /// and to prevent interfiere with normal traffic, when unbanning due to ban_time expired.
+    addresses_being_unbanned: Arc<RwLock<HashSet<Address>>>,
+
     /// List of banned addresses (see above)
     /// that should not be queried.
     banlist: BanList,
@@ -40,9 +48,14 @@ pub enum UnbanReason {
 }
 
 impl BanService {
-    pub fn new(replica_to_primary_failover_enabled: bool, ban_time: i64) -> Self {
+    pub fn new(
+        number_of_shards: usize,
+        replica_to_primary_failover_enabled: bool,
+        ban_time: i64,
+    ) -> Self {
         BanService {
-            banlist: Arc::new(RwLock::new(vec![HashMap::new()])),
+            addresses_being_unbanned: Arc::new(RwLock::new(HashSet::new())),
+            banlist: Arc::new(RwLock::new(vec![HashMap::new(); number_of_shards])),
             replica_to_primary_failover_enabled,
             ban_time,
         }
@@ -89,6 +102,52 @@ impl BanService {
         guard[address.shard].remove(address);
     }
 
+    /// Starts a tokio::task that will try to connect to the server and
+    /// if successful, it will unban the server.
+    pub fn schedule_unban(
+        &self,
+        address: &Address,
+        server_pool: &Pool<ServerPool>,
+        client_stats: ClientStats,
+    ) {
+        // If the address is in the backlog, then it means that
+        // a healthcheck is being performed and we should do nothing.
+        if self.addresses_being_unbanned.read().contains(address) {
+            return;
+        } else {
+            self.addresses_being_unbanned
+                .write()
+                .insert(address.clone());
+        }
+
+        tokio::spawn({
+            let server_pool = server_pool.clone();
+            let address = address.clone();
+            let banlist = self.banlist.clone();
+            let addresses_being_unbanned = self.addresses_being_unbanned.clone();
+
+            async move {
+                match server_pool.get().await {
+                    Ok(_) => {
+                        address.reset_error_count();
+                        warn!("Unbanning {:?}", address);
+                        let mut guard = banlist.write();
+                        guard[address.shard].remove(&address);
+                    }
+                    Err(err) => {
+                        error!(
+                            "Connection checkout error while trying to unban instance {:?}, error: {:?}",
+                            address, err
+                        );
+                        address.stats.error();
+                        client_stats.checkout_error();
+                    }
+                };
+                addresses_being_unbanned.write().remove(&address);
+            }
+        });
+    }
+
     /// Check if address is banned
     /// true if banned, false otherwise
     pub fn is_banned(&self, address: &Address) -> bool {

src/pool.rs

@@ -489,11 +489,12 @@ impl ConnectionPool {
                         pool_name, user.username
                     );
                 }
-
+                let number_of_shards = shards.len();
                 let pool = ConnectionPool {
                     databases: Arc::new(shards),
                     addresses: Arc::new(addresses),
                     ban_service: Arc::new(BanService::new(
+                        number_of_shards,
                         pool_config.replica_to_primary_failover_enabled,
                         config.general.ban_time,
                     )),
@@ -748,8 +749,12 @@ impl ConnectionPool {
                         force_healthcheck = true;
                     }
                     Some(UnbanReason::BanTimeExceeded) => {
-                        self.ban_service.unban(address);
-                        force_healthcheck = true;
+                        self.ban_service.schedule_unban(
+                            address,
+                            &self.databases[address.shard][address.address_index],
+                            client_stats.clone(),
+                        );
+                        continue;
                     }
                     Some(UnbanReason::PrimaryBanned) | Some(UnbanReason::NotBanned) => {
                         force_healthcheck = true;

src/query_router.rs

@@ -1453,7 +1453,6 @@ mod test {
             automatic_sharding_key: Some(String::from("test.id")),
             healthcheck_delay: PoolSettings::default().healthcheck_delay,
             healthcheck_timeout: PoolSettings::default().healthcheck_timeout,
-            ban_time: PoolSettings::default().ban_time,
             sharding_key_regex: None,
             shard_id_regex: None,
             default_shard: crate::config::DefaultShard::Shard(0),
@@ -1532,7 +1531,6 @@ mod test {
             automatic_sharding_key: None,
             healthcheck_delay: PoolSettings::default().healthcheck_delay,
             healthcheck_timeout: PoolSettings::default().healthcheck_timeout,
-            ban_time: PoolSettings::default().ban_time,
             sharding_key_regex: Some(Regex::new(r"/\* sharding_key: (\d+) \*/").unwrap()),
             shard_id_regex: Some(Regex::new(r"/\* shard_id: (\d+) \*/").unwrap()),
             default_shard: crate::config::DefaultShard::Shard(0),

tests/ruby/load_balancing_spec.rb

@@ -253,6 +253,7 @@
 
           expect(failed_count).to(eq(number_of_replicas))
         end
+
         context("when banned replicas are tested for availability because they expired the ban time") do
           let(:ban_time) { 2 }
           it "should be done in the background without interfering with traffic" do
@@ -266,6 +267,16 @@
             failed_count = 0
             number_of_replicas = processes[:replicas].length
 
+            # We need to allow pgcat to open connections to replicas
+            (number_of_replicas + 10).times do |n|
+              response = conn.async_exec(select_server_port)
+              expect(response[0]["port"].to_i).not_to(eq(primary_port))
+            rescue
+              conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+              failed_count += 1
+            end
+            expect(failed_count).to(eq(0))
+
             # Take down all replicas
             processes[:replicas].each(&:take_down)
 
@@ -295,6 +306,49 @@
             expect(response[0]["port"].to_i).to(eq(primary_port))
             expect(failed_count).to(eq(0))
           end
+
+          it "should unban replicas if they become available after the ban time" do
+            select_server_port = "SELECT setting AS port FROM pg_settings WHERE name = 'port';"
+            conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+            failed_count = 0
+            primary_port = processes.primary.original_port;
+
+            number_of_replicas = processes[:replicas].length
+
+            # We need to allow pgcat to open connections to replicas
+            (number_of_replicas + 10).times do |n|
+              response = conn.async_exec(select_server_port)
+              expect(response[0]["port"].to_i).not_to(eq(primary_port))
+            rescue
+              conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+              failed_count += 1
+            end
+            expect(failed_count).to(eq(0))
+
+            # Take down all replicas
+            processes[:replicas].each(&:take_down)
+
+            (number_of_replicas).times do |n|
+              conn.async_exec("SELECT 1 + 2")
+            rescue
+              conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+              failed_count += 1
+            end
+            expect(failed_count).to(eq(number_of_replicas))
+            failed_count = 0
+
+            processes[:replicas].each(&:reset)
+            sleep(ban_time + 1)
+            response = nil
+            number_of_replicas.times do
+              response = conn.async_exec("SELECT 1 + 2")
+            rescue
+              conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+              failed_count += 1
+            end
+            expect(response[0]["port"].to_i).not_to(eq(primary_port))
+            expect(failed_count).to(eq(0))
+          end
         end
       end
     end