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

Fix formatting (#299)

Browse Source
This commit is contained in:
Lev Kokotov
2023-01-28 09:17:49 -08:00
committed by GitHub
parent fbe256cc4e
commit 2e3eb2663e
1 changed files with 139 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

277
src/main.rs
View File

@@ -123,103 +123,105 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
runtime.block_on(async move { runtime.block_on(async move {
if let Some(true) = config.general.enable_prometheus_exporter { if let Some(true) = config.general.enable_prometheus_exporter {
let http_addr_str = format!( let http_addr_str = format!(
"{}:{}", "{}:{}",
config.general.host, config.general.prometheus_exporter_port config.general.host, config.general.prometheus_exporter_port
); );
let http_addr = match SocketAddr::from_str(&http_addr_str) { let http_addr = match SocketAddr::from_str(&http_addr_str) {
Ok(addr) => addr, Ok(addr) => addr,
Err(err) => { Err(err) => {
error!("Invalid http address: {}", err); error!("Invalid http address: {}", err);
std::process::exit(exitcode::CONFIG); std::process::exit(exitcode::CONFIG);
} }
}; };
tokio::task::spawn(async move { tokio::task::spawn(async move {
start_metric_server(http_addr).await; start_metric_server(http_addr).await;
}); });
} }
let addr = format!("{}:{}", config.general.host, config.general.port); let addr = format!("{}:{}", config.general.host, config.general.port);
let listener = match TcpListener::bind(&addr).await { let listener = match TcpListener::bind(&addr).await {
Ok(sock) => sock, Ok(sock) => sock,
Err(err) => { Err(err) => {
error!("Listener socket error: {:?}", err); error!("Listener socket error: {:?}", err);
std::process::exit(exitcode::CONFIG); std::process::exit(exitcode::CONFIG);
} }
}; };
info!("Running on {}", addr); info!("Running on {}", addr);
config.show(); config.show();
// Tracks which client is connected to which server for query cancellation. // Tracks which client is connected to which server for query cancellation.
let client_server_map: ClientServerMap = Arc::new(Mutex::new(HashMap::new())); let client_server_map: ClientServerMap = Arc::new(Mutex::new(HashMap::new()));
// Statistics reporting. // Statistics reporting.
let (stats_tx, stats_rx) = mpsc::channel(100_000); let (stats_tx, stats_rx) = mpsc::channel(100_000);
REPORTER.store(Arc::new(Reporter::new(stats_tx.clone()))); REPORTER.store(Arc::new(Reporter::new(stats_tx.clone())));
// Connection pool that allows to query all shards and replicas. // Connection pool that allows to query all shards and replicas.
match ConnectionPool::from_config(client_server_map.clone()).await { match ConnectionPool::from_config(client_server_map.clone()).await {
Ok(_) => (), Ok(_) => (),
Err(err) => { Err(err) => {
error!("Pool error: {:?}", err); error!("Pool error: {:?}", err);
std::process::exit(exitcode::CONFIG); std::process::exit(exitcode::CONFIG);
} }
}; };
tokio::task::spawn(async move { tokio::task::spawn(async move {
let mut stats_collector = Collector::new(stats_rx, stats_tx.clone()); let mut stats_collector = Collector::new(stats_rx, stats_tx.clone());
stats_collector.collect().await; stats_collector.collect().await;
}); });
info!("Config autoreloader: {}", config.general.autoreload); info!("Config autoreloader: {}", config.general.autoreload);
let mut autoreload_interval = tokio::time::interval(tokio::time::Duration::from_millis(15_000)); let mut autoreload_interval = tokio::time::interval(tokio::time::Duration::from_millis(15_000));
let autoreload_client_server_map = client_server_map.clone(); let autoreload_client_server_map = client_server_map.clone();
tokio::task::spawn(async move {
tokio::task::spawn(async move {
loop { loop {
autoreload_interval.tick().await; autoreload_interval.tick().await;
if config.general.autoreload { if config.general.autoreload {
info!("Automatically reloading config"); info!("Automatically reloading config");
if let Ok(changed) = reload_config(autoreload_client_server_map.clone()).await { if let Ok(changed) = reload_config(autoreload_client_server_map.clone()).await {
if changed { if changed {
get_config().show() get_config().show()
} }
}; };
} }
} }
}); });
let mut term_signal = unix_signal(SignalKind::terminate()).unwrap(); let mut term_signal = unix_signal(SignalKind::terminate()).unwrap();
let mut interrupt_signal = unix_signal(SignalKind::interrupt()).unwrap(); let mut interrupt_signal = unix_signal(SignalKind::interrupt()).unwrap();
let mut sighup_signal = unix_signal(SignalKind::hangup()).unwrap(); let mut sighup_signal = unix_signal(SignalKind::hangup()).unwrap();
let (shutdown_tx, _) = broadcast::channel::<()>(1); let (shutdown_tx, _) = broadcast::channel::<()>(1);
let (drain_tx, mut drain_rx) = mpsc::channel::<i32>(2048); let (drain_tx, mut drain_rx) = mpsc::channel::<i32>(2048);
let (exit_tx, mut exit_rx) = mpsc::channel::<()>(1); let (exit_tx, mut exit_rx) = mpsc::channel::<()>(1);
let mut admin_only = false;
let mut total_clients = 0;
info!("Waiting for clients"); info!("Waiting for clients");
let mut admin_only = false; loop {
let mut total_clients = 0;
loop {
tokio::select! { tokio::select! {
// Reload config: // Reload config:
// kill -SIGHUP $(pgrep pgcat) // kill -SIGHUP $(pgrep pgcat)
_ = sighup_signal.recv() => { _ = sighup_signal.recv() => {
info!("Reloading config"); info!("Reloading config");
_ = reload_config(client_server_map.clone()).await; _ = reload_config(client_server_map.clone()).await;
get_config().show(); get_config().show();
}, },
// Initiate graceful shutdown sequence on sig int // Initiate graceful shutdown sequence on sig int
_ = interrupt_signal.recv() => { _ = interrupt_signal.recv() => {
info!("Got SIGINT, waiting for client connection drain now"); info!("Got SIGINT, waiting for client connection drain now");
admin_only = true; admin_only = true;
@@ -229,101 +231,100 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let _ = drain_tx.send(0).await; let _ = drain_tx.send(0).await;
tokio::task::spawn(async move { tokio::task::spawn(async move {
let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout)); let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout));
// First tick fires immediately. // First tick fires immediately.
interval.tick().await; interval.tick().await;
// Second one in the interval time. // Second one in the interval time.
interval.tick().await; interval.tick().await;
// We're done waiting. // We're done waiting.
error!("Graceful shutdown timed out. {} active clients being closed", total_clients); error!("Graceful shutdown timed out. {} active clients being closed", total_clients);
let _ = exit_tx.send(()).await; let _ = exit_tx.send(()).await;
}); });
}, },
_ = term_signal.recv() => { _ = term_signal.recv() => {
info!("Got SIGTERM, closing with {} clients active", total_clients); info!("Got SIGTERM, closing with {} clients active", total_clients);
break; break;
}, },
new_client = listener.accept() => { new_client = listener.accept() => {
let (socket, addr) = match new_client { let (socket, addr) = match new_client {
Ok((socket, addr)) => (socket, addr), Ok((socket, addr)) => (socket, addr),
Err(err) => { Err(err) => {
error!("{:?}", err); error!("{:?}", err);
continue; continue;
}
};
let shutdown_rx = shutdown_tx.subscribe();
let drain_tx = drain_tx.clone();
let client_server_map = client_server_map.clone();
let tls_certificate = config.general.tls_certificate.clone();
tokio::task::spawn(async move {
let start = chrono::offset::Utc::now().naive_utc();
match client::client_entrypoint(
socket,
client_server_map,
shutdown_rx,
drain_tx,
admin_only,
tls_certificate.clone(),
config.general.log_client_connections,
)
.await
{
Ok(()) => {
let duration = chrono::offset::Utc::now().naive_utc() - start;
if config.general.log_client_disconnections {
info!(
"Client {:?} disconnected, session duration: {}",
addr,
format_duration(&duration)
);
} else {
debug!(
"Client {:?} disconnected, session duration: {}",
addr,
format_duration(&duration)
);
}
}
Err(err) => {
match err {
errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
_ => warn!("Client disconnected with error {:?}", err),
}
}
};
});
}
_ = exit_rx.recv() => {
break;
}
client_ping = drain_rx.recv() => {
let client_ping = client_ping.unwrap();
total_clients += client_ping;
if total_clients == 0 && admin_only {
let _ = exit_tx.send(()).await;
} }
} };
}
}
info!("Shutting down..."); let shutdown_rx = shutdown_tx.subscribe();
let drain_tx = drain_tx.clone();
let client_server_map = client_server_map.clone();
let tls_certificate = config.general.tls_certificate.clone();
tokio::task::spawn(async move {
let start = chrono::offset::Utc::now().naive_utc();
match client::client_entrypoint(
socket,
client_server_map,
shutdown_rx,
drain_tx,
admin_only,
tls_certificate.clone(),
config.general.log_client_connections,
)
.await
{
Ok(()) => {
let duration = chrono::offset::Utc::now().naive_utc() - start;
if config.general.log_client_disconnections {
info!(
"Client {:?} disconnected, session duration: {}",
addr,
format_duration(&duration)
);
} else {
debug!(
"Client {:?} disconnected, session duration: {}",
addr,
format_duration(&duration)
);
}
}
Err(err) => {
match err {
errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
_ => warn!("Client disconnected with error {:?}", err),
}
}
};
});
}
_ = exit_rx.recv() => {
break;
}
client_ping = drain_rx.recv() => {
let client_ping = client_ping.unwrap();
total_clients += client_ping;
if total_clients == 0 && admin_only {
let _ = exit_tx.send(()).await;
}
}
}
}
info!("Shutting down...");
}); });
Ok(()) Ok(())
} }