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

Constants, comments, CI fixes, dead code clean-up (#21)

Browse Source 
* constants

* server.rs docs

* client.rs comments

* dead code; comments

* comment

* query cancellation comments

* remove unnecessary cast

* move db setup up one step

* query cancellation test

* new line; good night
This commit is contained in:
Lev Kokotov
2022-02-15 22:45:45 -08:00
committed by GitHub
parent bb84dcee64
commit 7b0ceefb96
8 changed files with 208 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
src/client.rs
View File

@@ -5,12 +5,15 @@ use bytes::{Buf, BufMut, BytesMut};
use once_cell::sync::OnceCell;
use regex::Regex;
use tokio::io::{AsyncReadExt, BufReader};
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::TcpStream;
use tokio::net::{
tcp::{OwnedReadHalf, OwnedWriteHalf},
TcpStream,
};
use std::collections::HashMap;
use crate::config::Role;
use crate::constants::*;
use crate::errors::Error;
use crate::messages::*;
use crate::pool::{ClientServerMap, ConnectionPool};
@@ -97,7 +100,7 @@ impl Client {
match code {
// Client wants SSL. We don't support it at the moment.
80877103 => {
SSL_REQUEST_CODE => {
let mut no = BytesMut::with_capacity(1);
no.put_u8(b'N');
@@ -105,7 +108,7 @@ impl Client {
}
// Regular startup message.
196608 => {
PROTOCOL_VERSION_NUMBER => {
// TODO: perform actual auth.
let parameters = parse_startup(bytes.clone())?;
@@ -138,7 +141,7 @@ impl Client {
}
// Query cancel request.
80877102 => {
CANCEL_REQUEST_CODE => {
let (read, write) = stream.into_split();
let process_id = bytes.get_i32();
@@ -168,23 +171,31 @@ impl Client {
/// Client loop. We handle all messages between the client and the database here.
pub async fn handle(&mut self, mut pool: ConnectionPool) -> Result<(), Error> {
// Special: cancelling existing running query
// The client wants to cancel a query it has issued previously.
if self.cancel_mode {
let (process_id, secret_key, address, port) = {
let guard = self.client_server_map.lock().unwrap();
match guard.get(&(self.process_id, self.secret_key)) {
// Drop the mutex as soon as possible.
// We found the server the client is using for its query
// that it wants to cancel.
Some((process_id, secret_key, address, port)) => (
process_id.clone(),
secret_key.clone(),
address.clone(),
port.clone(),
),
// The client doesn't know / got the wrong server,
// we're closing the connection for security reasons.
None => return Ok(()),
}
};
// TODO: pass actual server host and port somewhere.
// Opens a new separate connection to the server, sends the backend_id
// and secret_key and then closes it for security reasons. No other interactions
// take place.
return Ok(Server::cancel(&address, &port, process_id, secret_key).await?);
}
@@ -217,7 +228,7 @@ impl Client {
};
// Parse for special server role selection command.
//
// SET SERVER ROLE TO '(primary|replica)';
match self.select_role(message.clone()) {
Some(r) => {
custom_protocol_response_ok(&mut self.write, "SET SERVER ROLE").await?;
@@ -236,15 +247,17 @@ impl Client {
}
};
let mut proxy = connection.0;
let mut reference = connection.0;
let _address = connection.1;
let server = &mut *proxy;
let server = &mut *reference;
// Claim this server as mine for query cancellation.
server.claim(self.process_id, self.secret_key);
// Transaction loop. Multiple queries can be issued by the client here.
// The connection belongs to the client until the transaction is over,
// or until the client disconnects if we are in session mode.
loop {
// No messages in the buffer, read one.
let mut message = if message.len() == 0 {
match read_message(&mut self.read).await {
Ok(message) => message,
@@ -268,19 +281,26 @@ impl Client {
msg
};
let original = message.clone(); // To be forwarded to the server
// The message will be forwarded to the server intact. We still would like to
// parse it below to figure out what to do with it.
let original = message.clone();
let code = message.get_u8() as char;
let _len = message.get_i32() as usize;
match code {
// ReadyForQuery
'Q' => {
// TODO: implement retries here for read-only transactions.
server.send(original).await?;
// Read all data the server has to offer, which can be multiple messages
// buffered in 8196 bytes chunks.
loop {
// TODO: implement retries here for read-only transactions.
let response = server.recv().await?;
// Send server reply to the client.
match write_all_half(&mut self.write, response).await {
Ok(_) => (),
Err(err) => {
@@ -294,15 +314,18 @@ impl Client {
}
}
// Send statistic
// Report query executed statistics.
self.stats.query();
// Transaction over
// The transaction is over, we can release the connection back to the pool.
if !server.in_transaction() {
// Report transaction executed statistics.
self.stats.transaction();
// Release server
// 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 {
// Report this client as idle.
self.stats.client_idle();
shard = None;
@@ -313,6 +336,7 @@ impl Client {
}
}
// Terminate
'X' => {
// Client closing. Rollback and clean up
// connection before releasing into the pool.
@@ -326,35 +350,46 @@ impl Client {
return Ok(());
}
// Parse
// The query with placeholders is here, e.g. `SELECT * FROM users WHERE email = $1 AND active = $2`.
'P' => {
// Extended protocol, let's buffer most of it
self.buffer.put(&original[..]);
}
// Bind
// The placeholder's replacements are here, e.g. 'user@email.com' and 'true'
'B' => {
self.buffer.put(&original[..]);
}
// Describe
// Command a client can issue to describe a previously prepared named statement.
'D' => {
self.buffer.put(&original[..]);
}
// Execute
// Execute a prepared statement prepared in `P` and bound in `B`.
'E' => {
self.buffer.put(&original[..]);
}
// Sync
// Frontend (client) is asking for the query result now.
'S' => {
// Extended protocol, client requests sync
self.buffer.put(&original[..]);
// TODO: retries for read-only transactions
// TODO: retries for read-only transactions.
server.send(self.buffer.clone()).await?;
self.buffer.clear();
// Read all data the server has to offer, which can be multiple messages
// buffered in 8196 bytes chunks.
loop {
// TODO: retries for read-only transactions
let response = server.recv().await?;
match write_all_half(&mut self.write, response).await {
Ok(_) => (),
Err(err) => {
@@ -368,9 +403,11 @@ impl Client {
}
}
// Report query executed statistics.
self.stats.query();
// Release server
// 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();
@@ -392,10 +429,13 @@ impl Client {
server.send(original).await?;
}
// CopyDone or CopyFail
// Copy is done, successfully or not.
'c' | 'f' => {
// Copy is done.
server.send(original).await?;
let response = server.recv().await?;
match write_all_half(&mut self.write, response).await {
Ok(_) => (),
Err(err) => {
@@ -404,24 +444,29 @@ impl Client {
}
};
// Release the server
// 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();
if self.transaction_mode {
shard = None;
role = self.default_server_role;
break;
}
}
}
// Some unexpected message. We either did not implement the protocol correctly
// or this is not a Postgres client we're talking to.
_ => {
println!(">>> Unexpected code: {}", code);
}
}
}
// The server is no longer bound to us, we can't cancel it's queries anymore.
self.release();
}
}
@@ -450,6 +495,7 @@ impl Client {
let len = buf.get_i32();
let query = String::from_utf8_lossy(&buf[..len as usize - 4 - 1]).to_ascii_uppercase(); // Don't read the ternminating null
let rgx = match SHARDING_REGEX_RE.get() {
Some(r) => r,
None => return None,
@@ -457,11 +503,13 @@ impl Client {
if rgx.is_match(&query) {
let shard = query.split("'").collect::<Vec<&str>>()[1];
match shard.parse::<i64>() {
Ok(shard) => {
let sharder = Sharder::new(shards);
Some(sharder.pg_bigint_hash(shard))
}
Err(_) => None,
}
} else {
@@ -481,6 +529,7 @@ impl Client {
let len = buf.get_i32();
let query = String::from_utf8_lossy(&buf[..len as usize - 4 - 1]).to_ascii_uppercase();
let rgx = match ROLE_REGEX_RE.get() {
Some(r) => r,
None => return None,
@@ -490,6 +539,7 @@ impl Client {
// it'll be time to abstract :).
if rgx.is_match(&query) {
let role = query.split("'").collect::<Vec<&str>>()[1];
match role {
"PRIMARY" => Some(Role::Primary),
"REPLICA" => Some(Role::Replica),