docker/pgcat
1
0
Fork 0
mirror of https://github.com/postgresml/pgcat.git synced 2026-03-23 01:16:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

name servers & dont leave open servers with bad state

Browse Source
This commit is contained in:
Lev Kokotov
2022-02-03 17:06:19 -08:00
parent abb71b6c4d
commit 6deb7b1162
6 changed files with 182 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
Cargo.lock generated
View File

@@ -118,6 +118,17 @@ dependencies = [
"version_check",
]
[[package]]
name = "getrandom"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "418d37c8b1d42553c93648be529cb70f920d3baf8ef469b74b9638df426e0b4c"
dependencies = [
"cfg-if",
"libc",
"wasi",
]
[[package]]
name = "hermit-abi"
version = "0.1.19"
@@ -259,6 +270,12 @@ version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
[[package]]
name = "ppv-lite86"
version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eb9f9e6e233e5c4a35559a617bf40a4ec447db2e84c20b55a6f83167b7e57872"
[[package]]
name = "proc-macro2"
version = "1.0.36"
@@ -285,9 +302,50 @@ dependencies = [
"bb8",
"bytes",
"md-5",
"rand",
"tokio",
]
[[package]]
name = "rand"
version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2e7573632e6454cf6b99d7aac4ccca54be06da05aca2ef7423d22d27d4d4bcd8"
dependencies = [
"libc",
"rand_chacha",
"rand_core",
"rand_hc",
]
[[package]]
name = "rand_chacha"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
dependencies = [
"ppv-lite86",
"rand_core",
]
[[package]]
name = "rand_core"
version = "0.6.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d34f1408f55294453790c48b2f1ebbb1c5b4b7563eb1f418bcfcfdbb06ebb4e7"
dependencies = [
"getrandom",
]
[[package]]
name = "rand_hc"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d51e9f596de227fda2ea6c84607f5558e196eeaf43c986b724ba4fb8fdf497e7"
dependencies = [
"rand_core",
]
[[package]]
name = "redox_syscall"
version = "0.2.10"
@@ -383,6 +441,12 @@ version = "0.9.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
[[package]]
name = "wasi"
version = "0.10.2+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6"
[[package]]
name = "winapi"
version = "0.3.9"

3
Cargo.toml
View File

@@ -10,4 +10,5 @@ tokio = { version = "1", features = ["full"] }
bytes = "1"
md-5 = "*"
bb8 = "*"
async-trait = "*"
async-trait = "*"
rand = "*"

46
src/client.rs
View File

@@ -9,13 +9,15 @@ use bytes::{Buf, BufMut, BytesMut};
use crate::errors::Error;
use crate::messages::*;
use bb8::Pool;
use crate::pool::ServerPool;
use bb8::Pool;
use rand::{distributions::Alphanumeric, Rng};
pub struct Client {
read: BufReader<OwnedReadHalf>,
write: OwnedWriteHalf,
buffer: BytesMut,
name: String,
}
impl Client {
@@ -57,10 +59,17 @@ impl Client {
let (read, write) = stream.into_split();
let name: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(7)
.map(char::from)
.collect();
return Ok(Client {
read: BufReader::new(read),
write: write,
buffer: BytesMut::with_capacity(8196),
name: name,
});
}
@@ -72,7 +81,6 @@ impl Client {
}
pub async fn handle(&mut self, pool: Pool<ServerPool>) -> Result<(), Error> {
loop {
// Only grab a connection once we have some traffic on the socket
// TODO: this is not the most optimal way to share servers.
@@ -84,8 +92,20 @@ impl Client {
let mut proxy = pool.get().await.unwrap();
let server = &mut *proxy;
server.set_name(&self.name).await?;
loop {
let mut message = read_message(&mut self.read).await?;
let mut message = match read_message(&mut self.read).await {
Ok(message) => message,
Err(err) => {
if server.in_transaction() {
server.mark_bad();
}
return Err(err);
}
};
let original = message.clone(); // To be forwarded to the server
let code = message.get_u8() as char;
let _len = message.get_i32() as usize;
@@ -94,10 +114,17 @@ impl Client {
'Q' => {
server.send(original).await?;
let response = server.recv().await?;
write_all_half(&mut self.write, response).await?;
match write_all_half(&mut self.write, response).await {
Ok(_) => (),
Err(err) => {
server.mark_bad();
return Err(err);
}
};
// Release server
// Release server
if !server.in_transaction() {
drop(server);
break;
}
}
@@ -131,10 +158,17 @@ impl Client {
self.buffer.clear();
let response = server.recv().await?;
write_all_half(&mut self.write, response).await?;
match write_all_half(&mut self.write, response).await {
Ok(_) => (),
Err(err) => {
server.mark_bad();
return Err(err);
}
};
// Release server
if !server.in_transaction() {
drop(server);
break;
}
}

10
src/main.rs
View File

@@ -1,17 +1,17 @@
extern crate async_trait;
extern crate bb8;
extern crate bytes;
extern crate md5;
extern crate tokio;
extern crate async_trait;
extern crate bb8;
use tokio::net::TcpListener;
use bb8::Pool;
use tokio::net::TcpListener;
mod client;
mod errors;
mod messages;
mod server;
mod pool;
mod server;
#[tokio::main]
async fn main() {
@@ -42,7 +42,7 @@ async fn main() {
// Client goes to another thread, bye.
tokio::task::spawn(async move {
println!(">> Client {:?} connected.", addr);
let pool = pool.clone();
match client::Client::startup(socket).await {

20
src/pool.rs
View File

@@ -1,9 +1,8 @@
use async_trait::async_trait;
use bb8::{ManageConnection, PooledConnection};
use crate::server::Server;
use crate::errors::Error;
use crate::server::Server;
pub struct ServerPool {
host: String,
@@ -30,10 +29,17 @@ impl ManageConnection for ServerPool {
type Connection = Server;
type Error = Error;
/// Attempts to create a new connection.
/// Attempts to create a new connection.
async fn connect(&self) -> Result<Self::Connection, Self::Error> {
println!(">> Getting connetion from pool");
Ok(Server::startup(&self.host, &self.port, &self.user, &self.password, &self.database).await?)
Ok(Server::startup(
&self.host,
&self.port,
&self.user,
&self.password,
&self.database,
)
.await?)
}
/// Determines if the connection is still connected to the database.
@@ -42,7 +48,7 @@ impl ManageConnection for ServerPool {
}
/// Synchronously determine if the connection is no longer usable, if possible.
fn has_broken(&self, _conn: &mut Self::Connection) -> bool {
false
fn has_broken(&self, conn: &mut Self::Connection) -> bool {
conn.is_bad()
}
}
}

66
src/server.rs
View File

@@ -11,6 +11,7 @@ pub struct Server {
write: OwnedWriteHalf,
buffer: BytesMut,
in_transaction: bool,
bad: bool,
}
impl Server {
@@ -114,6 +115,7 @@ impl Server {
write: write,
buffer: BytesMut::with_capacity(8196),
in_transaction: false,
bad: false,
});
}
@@ -126,12 +128,26 @@ impl Server {
}
pub async fn send(&mut self, messages: BytesMut) -> Result<(), Error> {
Ok(write_all_half(&mut self.write, messages).await?)
match write_all_half(&mut self.write, messages).await {
Ok(_) => Ok(()),
Err(err) => {
println!(">> Terminating server because of: {:?}", err);
self.bad = true;
Err(err)
}
}
}
pub async fn recv(&mut self) -> Result<BytesMut, Error> {
loop {
let mut message = read_message(&mut self.read).await?;
let mut message = match read_message(&mut self.read).await {
Ok(message) => message,
Err(err) => {
println!(">> Terminating server because of: {:?}", err);
self.bad = true;
return Err(err);
}
};
// Buffer the message we'll forward to the client in a bit.
self.buffer.put(&message[..]);
@@ -143,23 +159,30 @@ impl Server {
'Z' => {
// Ready for query, time to forward buffer to client.
let transaction_state = message.get_u8() as char;
match transaction_state {
'T' => {
self.in_transaction = true;
},
}
'I' => {
self.in_transaction = false;
},
}
// Error client didn't clean up!
// We shuold drop this server
'E' => {
self.bad = true;
}
_ => {
self.in_transaction = false;
},
self.bad = true;
return Err(Error::ProtocolSyncError);
}
};
break;
},
}
_ => {
// Keep buffering,
@@ -176,4 +199,31 @@ impl Server {
pub fn in_transaction(&self) -> bool {
self.in_transaction
}
pub fn is_bad(&self) -> bool {
self.bad
}
pub fn mark_bad(&mut self) {
println!(">> Server marked bad");
self.bad = true;
}
pub async fn set_name(&mut self, name: &str) -> Result<(), Error> {
let mut query = BytesMut::from(&format!("SET application_name = {}", name).as_bytes()[..]);
query.put_u8(0);
let len = query.len() as i32 + 4;
let mut msg = BytesMut::with_capacity(len as usize + 1);
msg.put_u8(b'Q');
msg.put_i32(len);
msg.put_slice(&query[..]);
self.send(msg).await?;
let _ = self.recv().await?;
Ok(())
}
}