docker/pgcat
1
0
Fork 0
mirror of https://github.com/postgresml/pgcat.git synced 2026-03-25 02:06:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add Manual host banning to PgCat (#340)

Browse Source 
Sometimes we want an admin to be able to ban a host for some time to route traffic away from that host for reasons like partial outages, replication lag, and scheduled maintenance.

We can achieve this today using a configuration update but a quicker approach is to send a control command to PgCat that bans the replica for some specified duration.

This command does not change the current banning rules like

Primaries cannot be banned
When all replicas are banned, all replicas are unbanned
This commit is contained in:
Mostafa Abdelraouf
2023-03-06 06:10:59 -06:00
committed by GitHub
parent 8a0da10a87
commit 2cc6a09fba
5 changed files with 300 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/client.rs
View File

@@ -1,6 +1,9 @@
use crate::errors::Error;
use crate::pool::BanReason;
/// Handle clients by pretending to be a PostgreSQL server.
use bytes::{Buf, BufMut, BytesMut};
use log::{debug, error, info, trace, warn};
use std::collections::HashMap;
use std::time::Instant;
use tokio::io::{split, AsyncReadExt, BufReader, ReadHalf, WriteHalf};
@@ -11,7 +14,7 @@ use tokio::sync::mpsc::Sender;
use crate::admin::{generate_server_info_for_admin, handle_admin};
use crate::config::{get_config, Address, PoolMode};
use crate::constants::*;
use crate::errors::Error;
use crate::messages::*;
use crate::pool::{get_pool, ClientServerMap, ConnectionPool};
use crate::query_router::{Command, QueryRouter};
@@ -1111,7 +1114,7 @@ where
match server.send(message).await {
Ok(_) => Ok(()),
Err(err) => {
pool.ban(address, self.process_id);
pool.ban(address, BanReason::MessageSendFailed, self.process_id);
Err(err)
}
}
@@ -1133,7 +1136,7 @@ where
Ok(result) => match result {
Ok(message) => Ok(message),
Err(err) => {
pool.ban(address, self.process_id);
pool.ban(address, BanReason::MessageReceiveFailed, self.process_id);
error_response_terminal(
&mut self.write,
&format!("error receiving data from server: {:?}", err),
@@ -1148,7 +1151,7 @@ where
address, pool.settings.user.username
);
server.mark_bad();
pool.ban(address, self.process_id);
pool.ban(address, BanReason::StatementTimeout, self.process_id);
error_response_terminal(&mut self.write, "pool statement timeout").await?;
Err(Error::StatementTimeout)
}
@@ -1157,7 +1160,7 @@ where
match server.recv().await {
Ok(message) => Ok(message),
Err(err) => {
pool.ban(address, self.process_id);
pool.ban(address, BanReason::MessageReceiveFailed, self.process_id);
error_response_terminal(
&mut self.write,
&format!("error receiving data from server: {:?}", err),