fix: don't accept new requests when shutting down

This commit is contained in:
Timo Kösters 2023-03-18 08:58:20 +01:00
parent da3871f39a
commit 2a7c4693b8
No known key found for this signature in database
GPG Key ID: 0B25E636FBA7E4CB
2 changed files with 22 additions and 10 deletions

View File

@ -7,7 +7,7 @@
#![allow(clippy::suspicious_else_formatting)]
#![deny(clippy::dbg_macro)]
use std::{future::Future, io, net::SocketAddr, time::Duration};
use std::{future::Future, io, net::SocketAddr, sync::atomic, time::Duration};
use axum::{
extract::{DefaultBodyLimit, FromRequest, MatchedPath},
@ -212,13 +212,6 @@ async fn run_server() -> io::Result<()> {
}
}
// On shutdown
info!(target: "shutdown-sync", "Received shutdown notification, notifying sync helpers...");
services().globals.rotate.fire();
#[cfg(feature = "systemd")]
let _ = sd_notify::notify(true, &[sd_notify::NotifyState::Stopping]);
Ok(())
}
@ -226,6 +219,9 @@ async fn spawn_task<B: Send + 'static>(
req: axum::http::Request<B>,
next: axum::middleware::Next<B>,
) -> std::result::Result<axum::response::Response, StatusCode> {
if services().globals.shutdown.load(atomic::Ordering::Relaxed) {
return Err(StatusCode::SERVICE_UNAVAILABLE);
}
tokio::spawn(next.run(req))
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
@ -452,6 +448,11 @@ async fn shutdown_signal(handle: ServerHandle) {
warn!("Received {}, shutting down...", sig);
handle.graceful_shutdown(Some(Duration::from_secs(30)));
services().globals.shutdown();
#[cfg(feature = "systemd")]
let _ = sd_notify::notify(true, &[sd_notify::NotifyState::Stopping]);
}
async fn not_found(uri: Uri) -> impl IntoResponse {

View File

@ -6,7 +6,7 @@ use ruma::{
use crate::api::server_server::FedDest;
use crate::{Config, Error, Result};
use crate::{services, Config, Error, Result};
use ruma::{
api::{
client::sync::sync_events,
@ -14,6 +14,7 @@ use ruma::{
},
DeviceId, RoomVersionId, ServerName, UserId,
};
use std::sync::atomic::{self, AtomicBool};
use std::{
collections::{BTreeMap, HashMap},
fs,
@ -24,7 +25,7 @@ use std::{
time::{Duration, Instant},
};
use tokio::sync::{broadcast, watch::Receiver, Mutex as TokioMutex, Semaphore};
use tracing::error;
use tracing::{error, info};
use trust_dns_resolver::TokioAsyncResolver;
type WellKnownMap = HashMap<OwnedServerName, (FedDest, String)>;
@ -58,6 +59,8 @@ pub struct Service {
pub roomid_federationhandletime: RwLock<HashMap<OwnedRoomId, (OwnedEventId, Instant)>>,
pub stateres_mutex: Arc<Mutex<()>>,
pub rotate: RotationHandler,
pub shutdown: AtomicBool,
}
/// Handles "rotation" of long-polling requests. "Rotation" in this context is similar to "rotation" of log files and the like.
@ -160,6 +163,7 @@ impl Service {
stateres_mutex: Arc::new(Mutex::new(())),
sync_receivers: RwLock::new(HashMap::new()),
rotate: RotationHandler::new(),
shutdown: AtomicBool::new(false),
};
fs::create_dir_all(s.get_media_folder())?;
@ -341,6 +345,13 @@ impl Service {
r.push(base64::encode_config(key, base64::URL_SAFE_NO_PAD));
r
}
pub fn shutdown(&self) {
self.shutdown.store(true, atomic::Ordering::Relaxed);
// On shutdown
info!(target: "shutdown-sync", "Received shutdown notification, notifying sync helpers...");
services().globals.rotate.fire();
}
}
fn reqwest_client_builder(config: &Config) -> Result<reqwest::ClientBuilder> {