mirror of
https://github.com/postgresml/pgcat.git
synced 2026-03-23 01:16:30 +00:00
Compare commits
2 Commits
sven_md5_a
...
v0.1.0-bet
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
341ebf4123 | ||
|
|
35828a0a8c |
@@ -13,6 +13,9 @@ function start_pgcat() {
|
||||
|
||||
# Setup the database with shards and user
|
||||
psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_routing_setup.sql
|
||||
PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard0 -i
|
||||
PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard1 -i
|
||||
PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard2 -i
|
||||
|
||||
# Install Toxiproxy to simulate a downed/slow database
|
||||
wget -O toxiproxy-2.1.4.deb https://github.com/Shopify/toxiproxy/releases/download/v2.1.4/toxiproxy_2.1.4_amd64.deb
|
||||
@@ -28,9 +31,9 @@ toxiproxy-cli create -l 127.0.0.1:5433 -u 127.0.0.1:5432 postgres_replica
|
||||
start_pgcat "info"
|
||||
|
||||
# pgbench test
|
||||
pgbench -i -h 127.0.0.1 -p 6432 && \
|
||||
pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple && \
|
||||
pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
|
||||
pgbench -i -h 127.0.0.1 -p 6432
|
||||
pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple -f tests/pgbench/simple.sql
|
||||
pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
|
||||
|
||||
# COPY TO STDOUT test
|
||||
psql -h 127.0.0.1 -p 6432 -c 'COPY (SELECT * FROM pgbench_accounts LIMIT 15) TO STDOUT;' > /dev/null
|
||||
|
||||
3
.gitignore
vendored
3
.gitignore
vendored
@@ -1,5 +1,2 @@
|
||||
/target
|
||||
*.deb
|
||||
.idea/*
|
||||
tests/ruby/.bundle/*
|
||||
tests/ruby/vendor/*
|
||||
34
Cargo.lock
generated
34
Cargo.lock
generated
@@ -220,12 +220,6 @@ dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1aab8fc367588b89dcee83ab0fd66b72b50b72fa1904d7095045ace2b0c81c35"
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.117"
|
||||
@@ -375,10 +369,8 @@ dependencies = [
|
||||
"regex",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
"sha-1",
|
||||
"sqlparser",
|
||||
"statsd",
|
||||
"tokio",
|
||||
"toml",
|
||||
]
|
||||
@@ -485,12 +477,6 @@ version = "0.6.25"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f497285884f3fcff424ffc933e56d7cbca511def0c9831a7f9b5f6153e3cc89b"
|
||||
|
||||
[[package]]
|
||||
name = "ryu"
|
||||
version = "1.0.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "73b4b750c782965c211b42f022f59af1fbceabdd026623714f104152f1ec149f"
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
version = "1.1.0"
|
||||
@@ -514,17 +500,6 @@ dependencies = [
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.79"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e8d9fa5c3b304765ce1fd9c4c8a3de2c8db365a5b91be52f186efc675681d95"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"ryu",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha-1"
|
||||
version = "0.10.0"
|
||||
@@ -566,15 +541,6 @@ dependencies = [
|
||||
"log",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "statsd"
|
||||
version = "0.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "df1efceb4bf2c0b5ebec94354285a43bbbed1375605bdf2ebe4132299434a330"
|
||||
dependencies = [
|
||||
"rand",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "1.0.86"
|
||||
|
||||
@@ -17,11 +17,9 @@ sha-1 = "0.10"
|
||||
toml = "0.5"
|
||||
serde = "1"
|
||||
serde_derive = "1"
|
||||
serde_json = "1"
|
||||
regex = "1"
|
||||
num_cpus = "1"
|
||||
once_cell = "1"
|
||||
statsd = "0.15"
|
||||
sqlparser = "0.14"
|
||||
log = "0.4"
|
||||
arc-swap = "1"
|
||||
|
||||
12
pgcat.toml
12
pgcat.toml
@@ -48,8 +48,8 @@ password = "sharding_user"
|
||||
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
["127.0.0.1", 5432, "primary"],
|
||||
["localhost", 5432, "replica"],
|
||||
[ "127.0.0.1", 5432, "primary" ],
|
||||
[ "localhost", 5432, "replica" ],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
# Database name (e.g. "postgres")
|
||||
@@ -58,8 +58,8 @@ database = "shard0"
|
||||
[shards.1]
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
["127.0.0.1", 5432, "primary"],
|
||||
["localhost", 5432, "replica"],
|
||||
[ "127.0.0.1", 5432, "primary" ],
|
||||
[ "localhost", 5432, "replica" ],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
database = "shard1"
|
||||
@@ -67,8 +67,8 @@ database = "shard1"
|
||||
[shards.2]
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
["127.0.0.1", 5432, "primary"],
|
||||
["localhost", 5432, "replica"],
|
||||
[ "127.0.0.1", 5432, "primary" ],
|
||||
[ "localhost", 5432, "replica" ],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
database = "shard2"
|
||||
|
||||
75
src/admin.rs
75
src/admin.rs
@@ -31,7 +31,7 @@ pub async fn handle_admin(
|
||||
|
||||
if query.starts_with("SHOW STATS") {
|
||||
trace!("SHOW STATS");
|
||||
show_stats(stream).await
|
||||
show_stats(stream, &pool).await
|
||||
} else if query.starts_with("RELOAD") {
|
||||
trace!("RELOAD");
|
||||
reload(stream).await
|
||||
@@ -77,11 +77,19 @@ async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
|
||||
])); // but admin tools that work with pgbouncer want this
|
||||
res.put(data_row(&vec![
|
||||
"free_clients".to_string(),
|
||||
stats["cl_idle"].to_string(),
|
||||
stats
|
||||
.keys()
|
||||
.map(|address_id| stats[&address_id]["cl_idle"])
|
||||
.sum::<i64>()
|
||||
.to_string(),
|
||||
]));
|
||||
res.put(data_row(&vec![
|
||||
"used_clients".to_string(),
|
||||
stats["cl_active"].to_string(),
|
||||
stats
|
||||
.keys()
|
||||
.map(|address_id| stats[&address_id]["cl_active"])
|
||||
.sum::<i64>()
|
||||
.to_string(),
|
||||
]));
|
||||
res.put(data_row(&vec![
|
||||
"login_clients".to_string(),
|
||||
@@ -89,11 +97,19 @@ async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
|
||||
]));
|
||||
res.put(data_row(&vec![
|
||||
"free_servers".to_string(),
|
||||
stats["sv_idle"].to_string(),
|
||||
stats
|
||||
.keys()
|
||||
.map(|address_id| stats[&address_id]["sv_idle"])
|
||||
.sum::<i64>()
|
||||
.to_string(),
|
||||
]));
|
||||
res.put(data_row(&vec![
|
||||
"used_servers".to_string(),
|
||||
stats["sv_active"].to_string(),
|
||||
stats
|
||||
.keys()
|
||||
.map(|address_id| stats[&address_id]["sv_active"])
|
||||
.sum::<i64>()
|
||||
.to_string(),
|
||||
]));
|
||||
res.put(data_row(&vec!["dns_names".to_string(), "0".to_string()]));
|
||||
res.put(data_row(&vec!["dns_zones".to_string(), "0".to_string()]));
|
||||
@@ -125,7 +141,7 @@ async fn show_version(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
}
|
||||
|
||||
/// SHOW POOLS
|
||||
async fn show_pools(stream: &mut OwnedWriteHalf, _pool: &ConnectionPool) -> Result<(), Error> {
|
||||
async fn show_pools(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
|
||||
let stats = get_stats();
|
||||
let config = {
|
||||
let guard = get_config();
|
||||
@@ -151,16 +167,26 @@ async fn show_pools(stream: &mut OwnedWriteHalf, _pool: &ConnectionPool) -> Resu
|
||||
let mut res = BytesMut::new();
|
||||
res.put(row_description(&columns));
|
||||
|
||||
let mut row = vec![String::from("all"), config.user.name.clone()];
|
||||
for shard in 0..pool.shards() {
|
||||
for server in 0..pool.servers(shard) {
|
||||
let address = pool.address(shard, server);
|
||||
let stats = match stats.get(&address.id) {
|
||||
Some(stats) => stats.clone(),
|
||||
None => HashMap::new(),
|
||||
};
|
||||
|
||||
for column in &columns[2..columns.len() - 1] {
|
||||
let value = stats.get(column.0).unwrap_or(&0).to_string();
|
||||
row.push(value);
|
||||
let mut row = vec![address.name(), config.user.name.clone()];
|
||||
|
||||
for column in &columns[2..columns.len() - 1] {
|
||||
let value = stats.get(column.0).unwrap_or(&0).to_string();
|
||||
row.push(value);
|
||||
}
|
||||
|
||||
row.push(config.general.pool_mode.to_string());
|
||||
res.put(data_row(&row));
|
||||
}
|
||||
}
|
||||
|
||||
row.push(config.general.pool_mode.to_string());
|
||||
|
||||
res.put(data_row(&row));
|
||||
res.put(command_complete("SHOW"));
|
||||
|
||||
res.put_u8(b'Z');
|
||||
@@ -309,7 +335,7 @@ async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
}
|
||||
|
||||
/// SHOW STATS
|
||||
async fn show_stats(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
async fn show_stats(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
|
||||
let columns = vec![
|
||||
("database", DataType::Text),
|
||||
("total_xact_count", DataType::Numeric),
|
||||
@@ -332,15 +358,24 @@ async fn show_stats(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
let mut res = BytesMut::new();
|
||||
res.put(row_description(&columns));
|
||||
|
||||
let mut row = vec![
|
||||
String::from("all"), // TODO: per-database stats,
|
||||
];
|
||||
for shard in 0..pool.shards() {
|
||||
for server in 0..pool.servers(shard) {
|
||||
let address = pool.address(shard, server);
|
||||
let stats = match stats.get(&address.id) {
|
||||
Some(stats) => stats.clone(),
|
||||
None => HashMap::new(),
|
||||
};
|
||||
|
||||
for column in &columns[1..] {
|
||||
row.push(stats.get(column.0).unwrap_or(&0).to_string());
|
||||
let mut row = vec![address.name()];
|
||||
|
||||
for column in &columns[1..] {
|
||||
row.push(stats.get(column.0).unwrap_or(&0).to_string());
|
||||
}
|
||||
|
||||
res.put(data_row(&row));
|
||||
}
|
||||
}
|
||||
|
||||
res.put(data_row(&row));
|
||||
res.put(command_complete("SHOW"));
|
||||
|
||||
res.put_u8(b'Z');
|
||||
|
||||
@@ -58,6 +58,12 @@ pub struct Client {
|
||||
|
||||
// Clients want to talk to admin
|
||||
admin: bool,
|
||||
|
||||
// Last address the client talked to
|
||||
last_address_id: Option<usize>,
|
||||
|
||||
// Last server process id we talked to
|
||||
last_server_id: Option<i32>,
|
||||
}
|
||||
|
||||
impl Client {
|
||||
@@ -108,18 +114,15 @@ impl Client {
|
||||
// Regular startup message.
|
||||
PROTOCOL_VERSION_NUMBER => {
|
||||
trace!("Got StartupMessage");
|
||||
|
||||
// TODO: perform actual auth.
|
||||
let parameters = parse_startup(bytes.clone())?;
|
||||
let mut user_name: String = String::new();
|
||||
match parameters.get(&"user") {
|
||||
Some(&user) => user_name = user,
|
||||
None => return Err(Error::ClientBadStartup),
|
||||
}
|
||||
start_auth(&mut stream, &user_name).await?;
|
||||
|
||||
// Generate random backend ID and secret key
|
||||
let process_id: i32 = rand::random();
|
||||
let secret_key: i32 = rand::random();
|
||||
|
||||
auth_ok(&mut stream).await?;
|
||||
write_all(&mut stream, server_info).await?;
|
||||
backend_key_data(&mut stream, process_id, secret_key).await?;
|
||||
ready_for_query(&mut stream).await?;
|
||||
@@ -150,6 +153,8 @@ impl Client {
|
||||
parameters: parameters,
|
||||
stats: stats,
|
||||
admin: admin,
|
||||
last_address_id: None,
|
||||
last_server_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -172,6 +177,8 @@ impl Client {
|
||||
parameters: HashMap::new(),
|
||||
stats: stats,
|
||||
admin: false,
|
||||
last_address_id: None,
|
||||
last_server_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -222,9 +229,6 @@ impl Client {
|
||||
loop {
|
||||
trace!("Client idle, waiting for message");
|
||||
|
||||
// Client idle, waiting for messages.
|
||||
self.stats.client_idle(self.process_id);
|
||||
|
||||
// Read a complete message from the client, which normally would be
|
||||
// either a `Q` (query) or `P` (prepare, extended protocol).
|
||||
// We can parse it here before grabbing a server from the pool,
|
||||
@@ -295,13 +299,13 @@ impl Client {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Waiting for server connection.
|
||||
self.stats.client_waiting(self.process_id);
|
||||
|
||||
debug!("Waiting for connection from pool");
|
||||
|
||||
// Grab a server from the pool: the client issued a regular query.
|
||||
let connection = match pool.get(query_router.shard(), query_router.role()).await {
|
||||
let connection = match pool
|
||||
.get(query_router.shard(), query_router.role(), self.process_id)
|
||||
.await
|
||||
{
|
||||
Ok(conn) => {
|
||||
debug!("Got connection from pool");
|
||||
conn
|
||||
@@ -315,15 +319,23 @@ impl Client {
|
||||
};
|
||||
|
||||
let mut reference = connection.0;
|
||||
let _address = connection.1;
|
||||
let address = connection.1;
|
||||
let server = &mut *reference;
|
||||
|
||||
// Claim this server as mine for query cancellation.
|
||||
server.claim(self.process_id, self.secret_key);
|
||||
|
||||
// "disconnect" from the previous server stats-wise
|
||||
if let Some(last_address_id) = self.last_address_id {
|
||||
self.stats
|
||||
.client_disconnecting(self.process_id, last_address_id);
|
||||
}
|
||||
|
||||
// Client active & server active
|
||||
self.stats.client_active(self.process_id);
|
||||
self.stats.server_active(server.process_id());
|
||||
self.stats.client_active(self.process_id, address.id);
|
||||
self.stats.server_active(server.process_id(), address.id);
|
||||
self.last_address_id = Some(address.id);
|
||||
self.last_server_id = Some(server.process_id());
|
||||
|
||||
debug!(
|
||||
"Client {:?} talking to server {:?}",
|
||||
@@ -395,17 +407,17 @@ impl Client {
|
||||
}
|
||||
|
||||
// Report query executed statistics.
|
||||
self.stats.query();
|
||||
self.stats.query(self.process_id, address.id);
|
||||
|
||||
// The transaction is over, we can release the connection back to the pool.
|
||||
if !server.in_transaction() {
|
||||
// Report transaction executed statistics.
|
||||
self.stats.transaction();
|
||||
self.stats.transaction(self.process_id, address.id);
|
||||
|
||||
// Release server back to the pool if we are in transaction mode.
|
||||
// If we are in session mode, we keep the server until the client disconnects.
|
||||
if self.transaction_mode {
|
||||
self.stats.server_idle(server.process_id());
|
||||
self.stats.server_idle(server.process_id(), address.id);
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -481,15 +493,15 @@ impl Client {
|
||||
}
|
||||
|
||||
// Report query executed statistics.
|
||||
self.stats.query();
|
||||
self.stats.query(self.process_id, address.id);
|
||||
|
||||
// Release server back to the pool if we are in transaction mode.
|
||||
// If we are in session mode, we keep the server until the client disconnects.
|
||||
if !server.in_transaction() {
|
||||
self.stats.transaction();
|
||||
self.stats.transaction(self.process_id, address.id);
|
||||
|
||||
if self.transaction_mode {
|
||||
self.stats.server_idle(server.process_id());
|
||||
self.stats.server_idle(server.process_id(), address.id);
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -520,10 +532,10 @@ impl Client {
|
||||
// Release server back to the pool if we are in transaction mode.
|
||||
// If we are in session mode, we keep the server until the client disconnects.
|
||||
if !server.in_transaction() {
|
||||
self.stats.transaction();
|
||||
self.stats.transaction(self.process_id, address.id);
|
||||
|
||||
if self.transaction_mode {
|
||||
self.stats.server_idle(server.process_id());
|
||||
self.stats.server_idle(server.process_id(), address.id);
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -540,6 +552,7 @@ impl Client {
|
||||
// The server is no longer bound to us, we can't cancel it's queries anymore.
|
||||
debug!("Releasing server back into the pool");
|
||||
self.release();
|
||||
self.stats.client_idle(self.process_id, address.id);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -552,6 +565,14 @@ impl Client {
|
||||
|
||||
impl Drop for Client {
|
||||
fn drop(&mut self) {
|
||||
self.stats.client_disconnecting(self.process_id);
|
||||
// Disconnect the client
|
||||
if let Some(address_id) = self.last_address_id {
|
||||
self.stats.client_disconnecting(self.process_id, address_id);
|
||||
|
||||
// The server is now idle
|
||||
if let Some(process_id) = self.last_server_id {
|
||||
self.stats.server_idle(process_id, address_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -48,6 +48,7 @@ impl PartialEq<Role> for Option<Role> {
|
||||
|
||||
#[derive(Clone, PartialEq, Hash, std::cmp::Eq, Debug)]
|
||||
pub struct Address {
|
||||
pub id: usize,
|
||||
pub host: String,
|
||||
pub port: String,
|
||||
pub shard: usize,
|
||||
@@ -58,6 +59,7 @@ pub struct Address {
|
||||
impl Default for Address {
|
||||
fn default() -> Address {
|
||||
Address {
|
||||
id: 0,
|
||||
host: String::from("127.0.0.1"),
|
||||
port: String::from("5432"),
|
||||
shard: 0,
|
||||
|
||||
@@ -8,7 +8,5 @@ pub enum Error {
|
||||
// ServerTimeout,
|
||||
// DirtyServer,
|
||||
BadConfig,
|
||||
BadUserList,
|
||||
AllServersDown,
|
||||
AuthenticationError
|
||||
}
|
||||
|
||||
25
src/main.rs
25
src/main.rs
@@ -31,7 +31,6 @@ extern crate once_cell;
|
||||
extern crate serde;
|
||||
extern crate serde_derive;
|
||||
extern crate sqlparser;
|
||||
extern crate statsd;
|
||||
extern crate tokio;
|
||||
extern crate toml;
|
||||
|
||||
@@ -50,7 +49,6 @@ use std::sync::Arc;
|
||||
mod admin;
|
||||
mod client;
|
||||
mod config;
|
||||
mod userlist;
|
||||
mod constants;
|
||||
mod errors;
|
||||
mod messages;
|
||||
@@ -95,15 +93,6 @@ async fn main() {
|
||||
}
|
||||
};
|
||||
|
||||
// Prepare user list
|
||||
match userlist::parse("userlist.json").await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
error!("Userlist parse error: {:?}", err);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let config = get_config();
|
||||
|
||||
let addr = format!("{}:{}", config.general.host, config.general.port);
|
||||
@@ -123,15 +112,19 @@ async fn main() {
|
||||
|
||||
// Collect statistics and send them to StatsD
|
||||
let (tx, rx) = mpsc::channel(100);
|
||||
let collector_tx = tx.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let mut stats_collector = Collector::new(rx, collector_tx);
|
||||
stats_collector.collect().await;
|
||||
});
|
||||
|
||||
// Connection pool for all shards and replicas
|
||||
let mut pool =
|
||||
ConnectionPool::from_config(client_server_map.clone(), Reporter::new(tx.clone())).await;
|
||||
|
||||
let collector_tx = tx.clone();
|
||||
let addresses = pool.databases();
|
||||
tokio::task::spawn(async move {
|
||||
let mut stats_collector = Collector::new(rx, collector_tx);
|
||||
stats_collector.collect(addresses).await;
|
||||
});
|
||||
|
||||
// Connect to all servers and validate their versions.
|
||||
let server_info = match pool.validate().await {
|
||||
Ok(info) => info,
|
||||
Err(err) => {
|
||||
|
||||
100
src/messages.rs
100
src/messages.rs
@@ -1,7 +1,5 @@
|
||||
/// Helper functions to send one-off protocol messages
|
||||
/// and handle TcpStream (TCP socket).
|
||||
|
||||
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use md5::{Digest, Md5};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
|
||||
@@ -9,16 +7,10 @@ use tokio::net::{
|
||||
tcp::{OwnedReadHalf, OwnedWriteHalf},
|
||||
TcpStream,
|
||||
};
|
||||
use log::{error};
|
||||
|
||||
use crate::errors::Error;
|
||||
use std::collections::HashMap;
|
||||
|
||||
use rand::Rng;
|
||||
|
||||
use crate::userlist::get_user_list;
|
||||
|
||||
|
||||
/// Postgres data type mappings
|
||||
/// used in RowDescription ('T') message.
|
||||
pub enum DataType {
|
||||
@@ -37,98 +29,6 @@ impl From<&DataType> for i32 {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
1. Generate salt (4 bytes of random data)
|
||||
md5(concat(md5(concat(password, username)), random-salt)))
|
||||
2. Send md5 auth request
|
||||
3. recieve PasswordMessage with salt.
|
||||
4. refactor md5_password function to be reusable
|
||||
5. check username hash combo against file
|
||||
6. AuthenticationOk or ErrorResponse
|
||||
**/
|
||||
pub async fn start_auth(stream: &mut TcpStream, user_name: &String) -> Result<(), Error> {
|
||||
let mut rng = rand::thread_rng();
|
||||
|
||||
//Generate random 4 byte salt
|
||||
let salt = rng.gen::<u32>();
|
||||
|
||||
// Send AuthenticationMD5Password request
|
||||
send_md5_request(stream, salt).await?;
|
||||
|
||||
let code = match stream.read_u8().await {
|
||||
Ok(code) => code as char,
|
||||
Err(_) => return Err(Error::AuthenticationError),
|
||||
};
|
||||
|
||||
match code {
|
||||
// Password response
|
||||
'p' => {
|
||||
fetch_password_and_authenticate(stream, &user_name, &salt).await?;
|
||||
Ok(auth_ok(stream).await?)
|
||||
}
|
||||
_ => {
|
||||
error!("Unknown code: {}", code);
|
||||
return Err(Error::AuthenticationError);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn send_md5_request(stream: &mut TcpStream, salt: u32) -> Result<(), Error> {
|
||||
let mut authentication_md5password = BytesMut::with_capacity(12);
|
||||
authentication_md5password.put_u8(b'R');
|
||||
authentication_md5password.put_i32(12);
|
||||
authentication_md5password.put_i32(5);
|
||||
authentication_md5password.put_u32(salt);
|
||||
|
||||
// Send AuthenticationMD5Password request
|
||||
Ok(write_all(stream, authentication_md5password).await?)
|
||||
}
|
||||
|
||||
pub async fn fetch_password_and_authenticate(stream: &mut TcpStream, user_name: &String, salt: &u32) -> Result<(), Error> {
|
||||
/**
|
||||
1. How do I store the lists of users and paswords? clear text or hash?? wtf
|
||||
2. Add auth to tests
|
||||
**/
|
||||
|
||||
let len = match stream.read_i32().await {
|
||||
Ok(len) => len,
|
||||
Err(_) => return Err(Error::AuthenticationError),
|
||||
};
|
||||
|
||||
// Read whatever is left.
|
||||
let mut password_hash = vec![0u8; len as usize - 4];
|
||||
|
||||
match stream.read_exact(&mut password_hash).await {
|
||||
Ok(_) => (),
|
||||
Err(_) => return Err(Error::AuthenticationError),
|
||||
};
|
||||
|
||||
let user_list = get_user_list();
|
||||
let mut password: String = String::new();
|
||||
match user_list.get(&user_name) {
|
||||
Some(&p) => password = p,
|
||||
None => return Err(Error::AuthenticationError),
|
||||
}
|
||||
|
||||
let mut md5 = Md5::new();
|
||||
|
||||
// concat('md5', md5(concat(md5(concat(password, username)), random-salt)))
|
||||
// First pass
|
||||
md5.update(&password.as_bytes());
|
||||
md5.update(&user_name.as_bytes());
|
||||
let output = md5.finalize_reset();
|
||||
// Second pass
|
||||
md5.update(format!("{:x}", output));
|
||||
md5.update(salt.to_be_bytes().to_vec());
|
||||
|
||||
|
||||
let password_string: String = String::from_utf8(password_hash).expect("Could not get password hash");
|
||||
match format!("md5{:x}", md5.finalize()) == password_string {
|
||||
true => Ok(()),
|
||||
_ => Err(Error::AuthenticationError)
|
||||
}
|
||||
}
|
||||
|
||||
/// Tell the client that authentication handshake completed successfully.
|
||||
pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
|
||||
let mut auth_ok = BytesMut::with_capacity(9);
|
||||
|
||||
40
src/pool.rs
40
src/pool.rs
@@ -38,6 +38,7 @@ impl ConnectionPool {
|
||||
let mut shards = Vec::new();
|
||||
let mut addresses = Vec::new();
|
||||
let mut banlist = Vec::new();
|
||||
let mut address_id = 0;
|
||||
let mut shard_ids = config
|
||||
.shards
|
||||
.clone()
|
||||
@@ -63,6 +64,7 @@ impl ConnectionPool {
|
||||
};
|
||||
|
||||
let address = Address {
|
||||
id: address_id,
|
||||
host: server.0.clone(),
|
||||
port: server.1.to_string(),
|
||||
role: role,
|
||||
@@ -70,6 +72,8 @@ impl ConnectionPool {
|
||||
shard: shard_idx.parse::<usize>().unwrap(),
|
||||
};
|
||||
|
||||
address_id += 1;
|
||||
|
||||
if role == Role::Replica {
|
||||
replica_number += 1;
|
||||
}
|
||||
@@ -121,9 +125,13 @@ impl ConnectionPool {
|
||||
pub async fn validate(&mut self) -> Result<BytesMut, Error> {
|
||||
let mut server_infos = Vec::new();
|
||||
|
||||
let stats = self.stats.clone();
|
||||
for shard in 0..self.shards() {
|
||||
for _ in 0..self.servers(shard) {
|
||||
let connection = match self.get(shard, None).await {
|
||||
// To keep stats consistent.
|
||||
let fake_process_id = 0;
|
||||
|
||||
let connection = match self.get(shard, None, fake_process_id).await {
|
||||
Ok(conn) => conn,
|
||||
Err(err) => {
|
||||
error!("Shard {} down or misconfigured: {:?}", shard, err);
|
||||
@@ -137,6 +145,8 @@ impl ConnectionPool {
|
||||
|
||||
let server_info = server.server_info();
|
||||
|
||||
stats.client_disconnecting(fake_process_id, address.id);
|
||||
|
||||
if server_infos.len() > 0 {
|
||||
// Compare against the last server checked.
|
||||
if server_info != server_infos[server_infos.len() - 1] {
|
||||
@@ -165,6 +175,7 @@ impl ConnectionPool {
|
||||
&mut self,
|
||||
shard: usize,
|
||||
role: Option<Role>,
|
||||
process_id: i32,
|
||||
) -> Result<(PooledConnection<'_, ServerPool>, Address), Error> {
|
||||
let now = Instant::now();
|
||||
let addresses = &self.addresses[shard];
|
||||
@@ -200,6 +211,8 @@ impl ConnectionPool {
|
||||
let index = self.round_robin % addresses.len();
|
||||
let address = &addresses[index];
|
||||
|
||||
self.stats.client_waiting(process_id, address.id);
|
||||
|
||||
// Make sure you're getting a primary or a replica
|
||||
// as per request. If no specific role is requested, the first
|
||||
// available will be chosen.
|
||||
@@ -219,6 +232,9 @@ impl ConnectionPool {
|
||||
Err(err) => {
|
||||
error!("Banning replica {}, error: {:?}", index, err);
|
||||
self.ban(address, shard);
|
||||
self.stats.client_disconnecting(process_id, address.id);
|
||||
self.stats
|
||||
.checkout_time(now.elapsed().as_micros(), process_id, address.id);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
@@ -227,7 +243,7 @@ impl ConnectionPool {
|
||||
let server = &mut *conn;
|
||||
let healthcheck_timeout = get_config().general.healthcheck_timeout;
|
||||
|
||||
self.stats.server_tested(server.process_id());
|
||||
self.stats.server_tested(server.process_id(), address.id);
|
||||
|
||||
match tokio::time::timeout(
|
||||
tokio::time::Duration::from_millis(healthcheck_timeout),
|
||||
@@ -238,8 +254,9 @@ impl ConnectionPool {
|
||||
// Check if health check succeeded
|
||||
Ok(res) => match res {
|
||||
Ok(_) => {
|
||||
self.stats.checkout_time(now.elapsed().as_micros());
|
||||
self.stats.server_idle(conn.process_id());
|
||||
self.stats
|
||||
.checkout_time(now.elapsed().as_micros(), process_id, address.id);
|
||||
self.stats.server_idle(conn.process_id(), address.id);
|
||||
return Ok((conn, address.clone()));
|
||||
}
|
||||
Err(_) => {
|
||||
@@ -248,6 +265,9 @@ impl ConnectionPool {
|
||||
server.mark_bad();
|
||||
|
||||
self.ban(address, shard);
|
||||
self.stats.client_disconnecting(process_id, address.id);
|
||||
self.stats
|
||||
.checkout_time(now.elapsed().as_micros(), process_id, address.id);
|
||||
continue;
|
||||
}
|
||||
},
|
||||
@@ -258,6 +278,9 @@ impl ConnectionPool {
|
||||
server.mark_bad();
|
||||
|
||||
self.ban(address, shard);
|
||||
self.stats.client_disconnecting(process_id, address.id);
|
||||
self.stats
|
||||
.checkout_time(now.elapsed().as_micros(), process_id, address.id);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
@@ -395,13 +418,14 @@ impl ManageConnection for ServerPool {
|
||||
async fn connect(&self) -> Result<Self::Connection, Self::Error> {
|
||||
info!(
|
||||
"Creating a new connection to {:?} using user {:?}",
|
||||
self.address, self.user.name
|
||||
self.address.name(),
|
||||
self.user.name
|
||||
);
|
||||
|
||||
// Put a temporary process_id into the stats
|
||||
// for server login.
|
||||
let process_id = rand::random::<i32>();
|
||||
self.stats.server_login(process_id);
|
||||
self.stats.server_login(process_id, self.address.id);
|
||||
|
||||
match Server::startup(
|
||||
&self.address,
|
||||
@@ -414,12 +438,12 @@ impl ManageConnection for ServerPool {
|
||||
{
|
||||
Ok(conn) => {
|
||||
// Remove the temporary process_id from the stats.
|
||||
self.stats.server_disconnecting(process_id);
|
||||
self.stats.server_disconnecting(process_id, self.address.id);
|
||||
Ok(conn)
|
||||
}
|
||||
Err(err) => {
|
||||
// Remove the temporary process_id from the stats.
|
||||
self.stats.server_disconnecting(process_id);
|
||||
self.stats.server_disconnecting(process_id, self.address.id);
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,7 +12,7 @@ use sqlparser::parser::Parser;
|
||||
|
||||
const CUSTOM_SQL_REGEXES: [&str; 5] = [
|
||||
r"(?i)^ *SET SHARDING KEY TO '?([0-9]+)'? *;? *$",
|
||||
r"(?i)^ *SET SHARD TO '?([0-9]+)'? *;? *$",
|
||||
r"(?i)^ *SET SHARD TO '?([0-9]+|ANY)'? *;? *$",
|
||||
r"(?i)^ *SHOW SHARD *;? *$",
|
||||
r"(?i)^ *SET SERVER ROLE TO '(PRIMARY|REPLICA|ANY|AUTO|DEFAULT)' *;? *$",
|
||||
r"(?i)^ *SHOW SERVER ROLE *;? *$",
|
||||
@@ -192,7 +192,10 @@ impl QueryRouter {
|
||||
}
|
||||
|
||||
Command::SetShard => {
|
||||
self.active_shard = Some(value.parse::<usize>().unwrap());
|
||||
self.active_shard = match value.to_ascii_uppercase().as_ref() {
|
||||
"ANY" => Some(rand::random::<usize>() % self.shards),
|
||||
_ => Some(value.parse::<usize>().unwrap()),
|
||||
};
|
||||
}
|
||||
|
||||
Command::SetServerRole => {
|
||||
|
||||
@@ -268,7 +268,8 @@ impl Server {
|
||||
|
||||
/// Send messages to the server from the client.
|
||||
pub async fn send(&mut self, messages: BytesMut) -> Result<(), Error> {
|
||||
self.stats.data_sent(messages.len());
|
||||
self.stats
|
||||
.data_sent(messages.len(), self.process_id, self.address.id);
|
||||
|
||||
match write_all_half(&mut self.write, messages).await {
|
||||
Ok(_) => Ok(()),
|
||||
@@ -374,7 +375,8 @@ impl Server {
|
||||
let bytes = self.buffer.clone();
|
||||
|
||||
// Keep track of how much data we got from the server for stats.
|
||||
self.stats.data_received(bytes.len());
|
||||
self.stats
|
||||
.data_received(bytes.len(), self.process_id, self.address.id);
|
||||
|
||||
// Clear the buffer for next query.
|
||||
self.buffer.clear();
|
||||
@@ -470,7 +472,8 @@ impl Drop for Server {
|
||||
/// the socket is in non-blocking mode, so it may not be ready
|
||||
/// for a write.
|
||||
fn drop(&mut self) {
|
||||
self.stats.server_disconnecting(self.process_id());
|
||||
self.stats
|
||||
.server_disconnecting(self.process_id(), self.address.id);
|
||||
|
||||
let mut bytes = BytesMut::with_capacity(4);
|
||||
bytes.put_u8(b'X');
|
||||
|
||||
275
src/stats.rs
275
src/stats.rs
@@ -1,17 +1,20 @@
|
||||
use log::{debug, info};
|
||||
/// Statistics and reporting.
|
||||
use log::info;
|
||||
use once_cell::sync::Lazy;
|
||||
use parking_lot::Mutex;
|
||||
use statsd::Client;
|
||||
use std::collections::HashMap;
|
||||
use tokio::sync::mpsc::{Receiver, Sender};
|
||||
|
||||
use std::collections::HashMap;
|
||||
// Latest stats updated every second; used in SHOW STATS and other admin commands.
|
||||
static LATEST_STATS: Lazy<Mutex<HashMap<usize, HashMap<String, i64>>>> =
|
||||
Lazy::new(|| Mutex::new(HashMap::new()));
|
||||
|
||||
use crate::config::get_config;
|
||||
|
||||
// Stats used in SHOW STATS
|
||||
static LATEST_STATS: Lazy<Mutex<HashMap<String, i64>>> = Lazy::new(|| Mutex::new(HashMap::new()));
|
||||
static STAT_PERIOD: u64 = 15000; //15 seconds
|
||||
// Statistics period used for average calculations.
|
||||
// 15 seconds.
|
||||
static STAT_PERIOD: u64 = 15000;
|
||||
|
||||
/// The names for the events reported
|
||||
/// to the statistics collector.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum EventName {
|
||||
CheckoutTime,
|
||||
@@ -28,186 +31,253 @@ enum EventName {
|
||||
ServerTested,
|
||||
ServerLogin,
|
||||
ServerDisconnecting,
|
||||
FlushStatsToStatsD,
|
||||
UpdateStats,
|
||||
UpdateAverages,
|
||||
}
|
||||
|
||||
/// Event data sent to the collector
|
||||
/// from clients and servers.
|
||||
#[derive(Debug)]
|
||||
pub struct Event {
|
||||
/// The name of the event being reported.
|
||||
name: EventName,
|
||||
|
||||
/// The value being reported. Meaning differs based on event name.
|
||||
value: i64,
|
||||
process_id: Option<i32>,
|
||||
|
||||
/// The client or server connection reporting the event.
|
||||
process_id: i32,
|
||||
|
||||
/// The server the client is connected to.
|
||||
address_id: usize,
|
||||
}
|
||||
|
||||
/// The statistics reporter. An instance is given
|
||||
/// to each possible source of statistics,
|
||||
/// e.g. clients, servers, connection pool.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Reporter {
|
||||
tx: Sender<Event>,
|
||||
}
|
||||
|
||||
impl Reporter {
|
||||
/// Create a new Reporter instance.
|
||||
pub fn new(tx: Sender<Event>) -> Reporter {
|
||||
Reporter { tx: tx }
|
||||
}
|
||||
|
||||
pub fn query(&self) {
|
||||
/// Report a query executed by a client against
|
||||
/// a server identified by the `address_id`.
|
||||
pub fn query(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::Query,
|
||||
value: 1,
|
||||
process_id: None,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn transaction(&self) {
|
||||
/// Report a transaction executed by a client against
|
||||
/// a server identified by the `address_id`.
|
||||
pub fn transaction(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::Transaction,
|
||||
value: 1,
|
||||
process_id: None,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn data_sent(&self, amount: usize) {
|
||||
/// Report data sent to a server identified by `address_id`.
|
||||
/// The `amount` is measured in bytes.
|
||||
pub fn data_sent(&self, amount: usize, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::DataSent,
|
||||
value: amount as i64,
|
||||
process_id: None,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn data_received(&self, amount: usize) {
|
||||
/// Report data received from a server identified by `address_id`.
|
||||
/// The `amount` is measured in bytes.
|
||||
pub fn data_received(&self, amount: usize, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::DataReceived,
|
||||
value: amount as i64,
|
||||
process_id: None,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn checkout_time(&self, ms: u128) {
|
||||
/// Time spent waiting to get a healthy connection from the pool
|
||||
/// for a server identified by `address_id`.
|
||||
/// Measured in milliseconds.
|
||||
pub fn checkout_time(&self, ms: u128, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::CheckoutTime,
|
||||
value: ms as i64,
|
||||
process_id: None,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn client_waiting(&self, process_id: i32) {
|
||||
/// Reports a client identified by `process_id` waiting for a connection
|
||||
/// to a server identified by `address_id`.
|
||||
pub fn client_waiting(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ClientWaiting,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn client_active(&self, process_id: i32) {
|
||||
/// Reports a client identified by `process_id` is done waiting for a connection
|
||||
/// to a server identified by `address_id` and is about to query the server.
|
||||
pub fn client_active(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ClientActive,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn client_idle(&self, process_id: i32) {
|
||||
/// Reports a client identified by `process_id` is done querying the server
|
||||
/// identified by `address_id` and is no longer active.
|
||||
pub fn client_idle(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ClientIdle,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn client_disconnecting(&self, process_id: i32) {
|
||||
/// Reports a client identified by `process_id` is disconecting from the pooler.
|
||||
/// The last server it was connected to is identified by `address_id`.
|
||||
pub fn client_disconnecting(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ClientDisconnecting,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn server_active(&self, process_id: i32) {
|
||||
/// Reports a server connection identified by `process_id` for
|
||||
/// a configured server identified by `address_id` is actively used
|
||||
/// by a client.
|
||||
pub fn server_active(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ServerActive,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn server_idle(&self, process_id: i32) {
|
||||
/// Reports a server connection identified by `process_id` for
|
||||
/// a configured server identified by `address_id` is no longer
|
||||
/// actively used by a client and is now idle.
|
||||
pub fn server_idle(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ServerIdle,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn server_login(&self, process_id: i32) {
|
||||
/// Reports a server connection identified by `process_id` for
|
||||
/// a configured server identified by `address_id` is attempting
|
||||
/// to login.
|
||||
pub fn server_login(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ServerLogin,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn server_tested(&self, process_id: i32) {
|
||||
/// Reports a server connection identified by `process_id` for
|
||||
/// a configured server identified by `address_id` is being
|
||||
/// tested before being given to a client.
|
||||
pub fn server_tested(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ServerTested,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
pub fn server_disconnecting(&self, process_id: i32) {
|
||||
/// Reports a server connection identified by `process_id` is disconecting from the pooler.
|
||||
/// The configured server it was connected to is identified by `address_id`.
|
||||
pub fn server_disconnecting(&self, process_id: i32, address_id: usize) {
|
||||
let event = Event {
|
||||
name: EventName::ServerDisconnecting,
|
||||
value: 1,
|
||||
process_id: Some(process_id),
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
}
|
||||
|
||||
/// The statistics collector which is receiving statistics
|
||||
/// from clients, servers, and the connection pool. There is
|
||||
/// only one collector (kind of like a singleton).
|
||||
/// The collector can trigger events on its own, e.g.
|
||||
/// it updates aggregates every second and averages every
|
||||
/// 15 seconds.
|
||||
pub struct Collector {
|
||||
rx: Receiver<Event>,
|
||||
tx: Sender<Event>,
|
||||
client: Client,
|
||||
}
|
||||
|
||||
impl Collector {
|
||||
/// Create a new collector instance. There should only be one instance
|
||||
/// at a time. This is ensured by mpsc which allows only one receiver.
|
||||
pub fn new(rx: Receiver<Event>, tx: Sender<Event>) -> Collector {
|
||||
Collector {
|
||||
rx,
|
||||
tx,
|
||||
client: Client::new(&get_config().general.statsd_address, "pgcat").unwrap(),
|
||||
}
|
||||
Collector { rx, tx }
|
||||
}
|
||||
|
||||
pub async fn collect(&mut self) {
|
||||
/// The statistics collection handler. It will collect statistics
|
||||
/// for `address_id`s starting at 0 up to `addresses`.
|
||||
pub async fn collect(&mut self, addresses: usize) {
|
||||
info!("Events reporter started");
|
||||
|
||||
let mut stats = HashMap::from([
|
||||
let stats_template = HashMap::from([
|
||||
("total_query_count", 0),
|
||||
("total_xact_count", 0),
|
||||
("total_sent", 0),
|
||||
@@ -232,25 +302,47 @@ impl Collector {
|
||||
("sv_tested", 0),
|
||||
]);
|
||||
|
||||
let mut stats = HashMap::new();
|
||||
|
||||
// Stats saved after each iteration of the flush event. Used in calculation
|
||||
// of averages in the last flush period.
|
||||
let mut old_stats: HashMap<String, i64> = HashMap::new();
|
||||
let mut old_stats: HashMap<usize, HashMap<String, i64>> = HashMap::new();
|
||||
|
||||
// Track which state the client and server are at any given time.
|
||||
let mut client_server_states: HashMap<i32, EventName> = HashMap::new();
|
||||
let mut client_server_states: HashMap<usize, HashMap<i32, EventName>> = HashMap::new();
|
||||
|
||||
// Flush stats to StatsD and calculate averages every 15 seconds.
|
||||
let tx = self.tx.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let mut interval =
|
||||
tokio::time::interval(tokio::time::Duration::from_millis(STAT_PERIOD / 15));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
for address_id in 0..addresses {
|
||||
let _ = tx.try_send(Event {
|
||||
name: EventName::UpdateStats,
|
||||
value: 0,
|
||||
process_id: -1,
|
||||
address_id: address_id,
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let tx = self.tx.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let mut interval =
|
||||
tokio::time::interval(tokio::time::Duration::from_millis(STAT_PERIOD));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
let _ = tx.try_send(Event {
|
||||
name: EventName::FlushStatsToStatsD,
|
||||
value: 0,
|
||||
process_id: None,
|
||||
});
|
||||
for address_id in 0..addresses {
|
||||
let _ = tx.try_send(Event {
|
||||
name: EventName::UpdateAverages,
|
||||
value: 0,
|
||||
process_id: -1,
|
||||
address_id: address_id,
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -264,6 +356,14 @@ impl Collector {
|
||||
}
|
||||
};
|
||||
|
||||
let stats = stats
|
||||
.entry(stat.address_id)
|
||||
.or_insert(stats_template.clone());
|
||||
let client_server_states = client_server_states
|
||||
.entry(stat.address_id)
|
||||
.or_insert(HashMap::new());
|
||||
let old_stats = old_stats.entry(stat.address_id).or_insert(HashMap::new());
|
||||
|
||||
// Some are counters, some are gauges...
|
||||
match stat.name {
|
||||
EventName::Query => {
|
||||
@@ -313,16 +413,16 @@ impl Collector {
|
||||
| EventName::ServerIdle
|
||||
| EventName::ServerTested
|
||||
| EventName::ServerLogin => {
|
||||
client_server_states.insert(stat.process_id.unwrap(), stat.name);
|
||||
client_server_states.insert(stat.process_id, stat.name);
|
||||
}
|
||||
|
||||
EventName::ClientDisconnecting | EventName::ServerDisconnecting => {
|
||||
client_server_states.remove(&stat.process_id.unwrap());
|
||||
client_server_states.remove(&stat.process_id);
|
||||
}
|
||||
|
||||
EventName::FlushStatsToStatsD => {
|
||||
EventName::UpdateStats => {
|
||||
// Calculate connection states
|
||||
for (_, state) in &client_server_states {
|
||||
for (_, state) in client_server_states.iter() {
|
||||
match state {
|
||||
EventName::ClientActive => {
|
||||
let counter = stats.entry("cl_active").or_insert(0);
|
||||
@@ -334,11 +434,6 @@ impl Collector {
|
||||
*counter += 1;
|
||||
}
|
||||
|
||||
EventName::ClientIdle => {
|
||||
let counter = stats.entry("cl_idle").or_insert(0);
|
||||
*counter += 1;
|
||||
}
|
||||
|
||||
EventName::ServerIdle => {
|
||||
let counter = stats.entry("sv_idle").or_insert(0);
|
||||
*counter += 1;
|
||||
@@ -359,39 +454,20 @@ impl Collector {
|
||||
*counter += 1;
|
||||
}
|
||||
|
||||
EventName::ClientIdle => {
|
||||
let counter = stats.entry("cl_idle").or_insert(0);
|
||||
*counter += 1;
|
||||
}
|
||||
|
||||
_ => unreachable!(),
|
||||
};
|
||||
}
|
||||
|
||||
// Calculate averages
|
||||
for stat in &[
|
||||
"avg_query_count",
|
||||
"avgxact_count",
|
||||
"avg_sent",
|
||||
"avg_received",
|
||||
"avg_wait_time",
|
||||
] {
|
||||
let total_name = stat.replace("avg_", "total_");
|
||||
let old_value = old_stats.entry(total_name.clone()).or_insert(0);
|
||||
let new_value = stats.get(total_name.as_str()).unwrap_or(&0).to_owned();
|
||||
let avg = (new_value - *old_value) / (STAT_PERIOD as i64 / 1_000); // Avg / second
|
||||
|
||||
stats.insert(stat, avg);
|
||||
*old_value = new_value;
|
||||
}
|
||||
|
||||
debug!("{:?}", stats);
|
||||
|
||||
// Update latest stats used in SHOW STATS
|
||||
let mut guard = LATEST_STATS.lock();
|
||||
for (key, value) in &stats {
|
||||
guard.insert(key.to_string(), value.clone());
|
||||
}
|
||||
|
||||
let mut pipeline = self.client.pipeline();
|
||||
|
||||
for (key, value) in stats.iter() {
|
||||
pipeline.gauge(key, *value as f64);
|
||||
let entry = guard.entry(stat.address_id).or_insert(HashMap::new());
|
||||
entry.insert(key.to_string(), value.clone());
|
||||
}
|
||||
|
||||
// These are re-calculated every iteration of the loop, so we don't want to add values
|
||||
@@ -409,14 +485,33 @@ impl Collector {
|
||||
] {
|
||||
stats.insert(stat, 0);
|
||||
}
|
||||
}
|
||||
|
||||
pipeline.send(&self.client);
|
||||
EventName::UpdateAverages => {
|
||||
// Calculate averages
|
||||
for stat in &[
|
||||
"avg_query_count",
|
||||
"avgxact_count",
|
||||
"avg_sent",
|
||||
"avg_received",
|
||||
"avg_wait_time",
|
||||
] {
|
||||
let total_name = stat.replace("avg_", "total_");
|
||||
let old_value = old_stats.entry(total_name.clone()).or_insert(0);
|
||||
let new_value = stats.get(total_name.as_str()).unwrap_or(&0).to_owned();
|
||||
let avg = (new_value - *old_value) / (STAT_PERIOD as i64 / 1_000); // Avg / second
|
||||
|
||||
stats.insert(stat, avg);
|
||||
*old_value = new_value;
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_stats() -> HashMap<String, i64> {
|
||||
/// Get a snapshot of statistics. Updated once a second
|
||||
/// by the `Collector`.
|
||||
pub fn get_stats() -> HashMap<usize, HashMap<String, i64>> {
|
||||
LATEST_STATS.lock().clone()
|
||||
}
|
||||
|
||||
@@ -1,4 +0,0 @@
|
||||
{
|
||||
"sven": "clear_text_password",
|
||||
"sharding_user": "sharding_user"
|
||||
}
|
||||
@@ -1,57 +0,0 @@
|
||||
use arc_swap::{ArcSwap, Guard};
|
||||
use log::{error};
|
||||
use once_cell::sync::Lazy;
|
||||
use tokio::fs::File;
|
||||
use tokio::io::AsyncReadExt;
|
||||
|
||||
use std::collections::{HashMap};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::errors::Error;
|
||||
|
||||
pub type UserList = HashMap<String, String>;
|
||||
static USER_LIST: Lazy<ArcSwap<UserList>> = Lazy::new(|| ArcSwap::from_pointee(HashMap::new()));
|
||||
|
||||
pub fn get_user_list() -> Guard<Arc<UserList>> {
|
||||
USER_LIST.load()
|
||||
}
|
||||
|
||||
/// Parse the user list.
|
||||
pub async fn parse(path: &str) -> Result<(), Error> {
|
||||
let mut contents = String::new();
|
||||
let mut file = match File::open(path).await {
|
||||
Ok(file) => file,
|
||||
Err(err) => {
|
||||
error!("Could not open '{}': {}", path, err.to_string());
|
||||
return Err(Error::BadConfig);
|
||||
}
|
||||
};
|
||||
|
||||
match file.read_to_string(&mut contents).await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
error!("Could not read config file: {}", err.to_string());
|
||||
return Err(Error::BadConfig);
|
||||
}
|
||||
};
|
||||
|
||||
let map: HashMap<String, String> = serde_json::from_str(&contents).expect("JSON was not well-formatted");
|
||||
|
||||
|
||||
|
||||
USER_LIST.store(Arc::new(map.clone()));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_config() {
|
||||
parse("userlist.json").await.unwrap();
|
||||
assert_eq!(get_user_list()["sven"], "clear_text_password");
|
||||
assert_eq!(get_user_list()["sharding_user"], "sharding_user");
|
||||
}
|
||||
}
|
||||
28
tests/pgbench/simple.sql
Normal file
28
tests/pgbench/simple.sql
Normal file
@@ -0,0 +1,28 @@
|
||||
|
||||
-- \setrandom aid 1 :naccounts
|
||||
\set aid random(1, 100000)
|
||||
-- \setrandom bid 1 :nbranches
|
||||
\set bid random(1, 100000)
|
||||
-- \setrandom tid 1 :ntellers
|
||||
\set tid random(1, 100000)
|
||||
-- \setrandom delta -5000 5000
|
||||
\set delta random(-5000,5000)
|
||||
|
||||
\set shard random(0, 2)
|
||||
|
||||
SET SHARD TO :shard;
|
||||
|
||||
BEGIN;
|
||||
|
||||
UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;
|
||||
|
||||
SELECT abalance FROM pgbench_accounts WHERE aid = :aid;
|
||||
|
||||
UPDATE pgbench_tellers SET tbalance = tbalance + :delta WHERE tid = :tid;
|
||||
|
||||
UPDATE pgbench_branches SET bbalance = bbalance + :delta WHERE bid = :bid;
|
||||
|
||||
INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);
|
||||
|
||||
END;
|
||||
|
||||
Reference in New Issue
Block a user