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

Adds SHUTDOWN command as alternate option to sending SIGINT (#331)

Browse Source 
* Adds SHUTDOWN command to PgCat as alternate option to sending SIGINT

* Check if we're already in SHUTDOWN sequence

* Send signal directly from shutdown instead of using channel

* Add tests

* trigger build

* Lowercase response and boolean change

* Update tests

* Fix tests

* typo
This commit is contained in:
zainkabani
2023-02-27 01:16:30 -05:00
committed by GitHub
parent 75a7d4409a
commit eb8cfdb1f1
5 changed files with 327 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
src/main.rs
View File

@@ -232,7 +232,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
// Initiate graceful shutdown sequence on sig int
_ = interrupt_signal.recv() => {
info!("Got SIGINT, waiting for client connection drain now");
info!("Got SIGINT");
// Don't want this to happen more than once
if admin_only {
continue;
}
admin_only = true;
// Broadcast that client tasks need to finish
@@ -241,98 +247,98 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let _ = drain_tx.send(0).await;
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.
interval.tick().await;
// First tick fires immediately.
interval.tick().await;
// Second one in the interval time.
interval.tick().await;
// Second one in the interval time.
interval.tick().await;
// We're done waiting.
error!("Graceful shutdown timed out. {} active clients being closed", total_clients);
// We're done waiting.
error!("Graceful shutdown timed out. {} active clients being closed", total_clients);
let _ = exit_tx.send(()).await;
let _ = exit_tx.send(()).await;
});
},
_ = term_signal.recv() => {
info!("Got SIGTERM, closing with {} clients active", total_clients);
break;
},
new_client = listener.accept() => {
let (socket, addr) = match new_client {
Ok((socket, addr)) => (socket, addr),
Err(err) => {
error!("{:?}", err);
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)
);
}
}
},
new_client = listener.accept() => {
let (socket, addr) = match new_client {
Ok((socket, addr)) => (socket, addr),
Err(err) => {
match err {
errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
_ => warn!("Client disconnected with error {:?}", err),
}
error!("{:?}", err);
continue;
}
};
});
}
_ = exit_rx.recv() => {
break;
}
let shutdown_rx = shutdown_tx.subscribe();
let drain_tx = drain_tx.clone();
let client_server_map = client_server_map.clone();
client_ping = drain_rx.recv() => {
let client_ping = client_ping.unwrap();
total_clients += client_ping;
let tls_certificate = config.general.tls_certificate.clone();
if total_clients == 0 && admin_only {
let _ = exit_tx.send(()).await;
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...");
});