this file contains shutdown handler pattern.
<file>src/common/http/src/http_shutdown_handlers.rs</file>
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::net::SocketAddr;
use std::time::Duration;
use anyerror::AnyError;
use databend_common_base::base::tokio;
use databend_common_base::base::tokio::sync::broadcast;
use databend_common_base::base::tokio::sync::oneshot;
use databend_common_base::base::tokio::task::JoinHandle;
use futures::future::Either;
use futures::FutureExt;
use log::error;
use log::info;
use poem::listener::Acceptor;
use poem::listener::AcceptorExt;
use poem::listener::IntoTlsConfigStream;
use poem::listener::Listener;
use poem::listener::OpensslTlsConfig;
use poem::listener::TcpListener;
use poem::Endpoint;
use crate::HttpError;
pub struct HttpShutdownHandler {
service_name: String,
join_handle: Option<JoinHandle<std::io::Result<()>>>,
abort_handle: Option<oneshot::Sender<()>>,
}
impl HttpShutdownHandler {
pub fn create(service_name: String) -> HttpShutdownHandler {
HttpShutdownHandler {
service_name,
join_handle: None,
abort_handle: None,
}
}
pub async fn start_service(
&mut self,
listening: SocketAddr,
tls_config: Option<OpensslTlsConfig>,
ep: impl Endpoint + 'static,
graceful_shutdown_timeout: Option<Duration>,
) -> Result<SocketAddr, HttpError> {
assert!(self.join_handle.is_none());
assert!(self.abort_handle.is_none());
let mut acceptor = TcpListener::bind(listening)
.into_acceptor()
.await
.map_err(|err| HttpError::listen_error(listening, err))?
.boxed();
let addr = acceptor
.local_addr()
.pop()
.and_then(|addr| addr.0.as_socket_addr().cloned())
.expect("socket addr");
if let Some(tls_config) = tls_config {
let conf_stream = tls_config
.into_stream()
.map_err(|err| HttpError::TlsConfigError(AnyError::new(&err)))?;
acceptor = acceptor.openssl_tls(conf_stream).boxed();
}
let (tx, rx) = oneshot::channel();
let join_handle = databend_common_base::runtime::spawn(
poem::Server::new_with_acceptor(acceptor)
.name(self.service_name.clone())
.idle_timeout(Duration::from_secs(20))
.run_with_graceful_shutdown(ep, rx.map(|_| ()), graceful_shutdown_timeout),
);
self.join_handle = Some(join_handle);
self.abort_handle = Some(tx);
Ok(addr)
}
/// Shutdown in graceful mode and returns a join handle.
/// To force shutdown: call the `abort()` method of the returned handle.
fn send_stop_signal(&mut self) -> JoinHandle<std::io::Result<()>> {
info!("{}: graceful stop", self.service_name);
if let Some(abort_handle) = self.abort_handle.take() {
info!("{}: send signal to abort_handle", self.service_name);
let res = abort_handle.send(());
info!(
"Done: {}: send signal to abort_handle, res: {:?}",
self.service_name, res
);
}
let join_handle = self.join_handle.take();
join_handle.unwrap()
}
/// Stop service gracefully. If `force` is ready, force shutdown the service.
pub async fn stop(&mut self, force: Option<broadcast::Receiver<()>>) {
let join_handle = self.send_stop_signal();
if let Some(mut force) = force {
let h = Box::pin(join_handle);
let f = Box::pin(force.recv());
match futures::future::select(f, h).await {
Either::Left((_x, h)) => {
info!("{}: received force shutdown signal", self.service_name);
h.abort();
}
Either::Right((_, _)) => {
info!("Done: {}: graceful shutdown", self.service_name);
}
}
} else {
info!(
"{}: force is None, wait for join handle for ever",
self.service_name
);
let res = join_handle.await;
info!(
"Done: {}: waiting for join handle for ever, res: {:?}",
self.service_name, res
);
}
}
pub async fn shutdown(&mut self, graceful: bool) {
if graceful {
if let Some(abort_handle) = self.abort_handle.take() {
let _ = abort_handle.send(());
}
if let Some(join_handle) = self.join_handle.take() {
if let Err(error) = join_handle.await {
error!(
"Unexpected error during shutdown Http Server {}. cause {}",
self.service_name, error
);
}
}
if let Some(join_handle) = self.join_handle.take() {
if let Err(_err) = tokio::time::timeout(Duration::from_secs(5), join_handle).await {
error!("Timeout during shutdown Http Server {}", self.service_name);
}
}
} else if let Some(join_handle) = self.join_handle.take() {
join_handle.abort();
}
}
}
<file>src/query/service/src/servers/http/http_services.rs</file>
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::net::SocketAddr;
use std::time::Duration;
use databend_common_config::GlobalConfig;
use databend_common_config::InnerConfig;
use databend_common_exception::ErrorCode;
use databend_common_http::HttpError;
use databend_common_http::HttpShutdownHandler;
use databend_common_meta_types::anyerror::AnyError;
use http::StatusCode;
use log::info;
use poem::get;
use poem::listener::OpensslTlsConfig;
use poem::middleware::CatchPanic;
use poem::middleware::CookieJarManager;
use poem::middleware::NormalizePath;
use poem::middleware::TrailingSlash;
use poem::Endpoint;
use poem::EndpointExt;
use poem::IntoResponse;
use poem::Route;
use super::v1::HttpQueryContext;
use crate::servers::http::middleware::json_response;
use crate::servers::http::middleware::EndpointKind;
use crate::servers::http::middleware::HTTPSessionMiddleware;
use crate::servers::http::middleware::PanicHandler;
use crate::servers::http::v1::clickhouse_router;
use crate::servers::http::v1::query_route;
use crate::servers::Server;
#[derive(Copy, Clone)]
pub enum HttpHandlerKind {
Query,
Clickhouse,
}
impl HttpHandlerKind {
pub fn usage(&self, sock: SocketAddr) -> String {
match self {
HttpHandlerKind::Query => {
format!(
r#" curl -u${{USER}} -p${{PASSWORD}}: --request POST '{:?}/v1/query/' --header 'Content-Type: application/json' --data-raw '{{"sql": "SELECT avg(number) FROM numbers(100000000)"}}'
"#,
sock,
)
}
HttpHandlerKind::Clickhouse => {
let json = r#"{"foo": "bar"}"#;
format!(
r#" echo 'create table test(foo string)' | curl -u${{USER}} -p${{PASSWORD}}: '{:?}' --data-binary @-
)
}
}
}
}
pub struct HttpHandler {
shutdown_handler: HttpShutdownHandler,
kind: HttpHandlerKind,
}
#[poem::handler]
#[async_backtrace::framed]
pub async fn verify_handler(_ctx: &HttpQueryContext) -> poem::Result<impl IntoResponse> {
Ok(StatusCode::OK)
}
impl HttpHandler {
pub fn create(kind: HttpHandlerKind) -> Box<dyn Server> {
Box::new(HttpHandler {
kind,
shutdown_handler: HttpShutdownHandler::create("http handler".to_string()),
})
}
#[allow(clippy::let_with_type_underscore)]
#[async_backtrace::framed]
async fn build_router(&self, sock: SocketAddr) -> impl Endpoint {
let ep_clickhouse = Route::new()
.nest("/", clickhouse_router())
.with(HTTPSessionMiddleware::create(
self.kind,
EndpointKind::Clickhouse,
))
.with(CookieJarManager::new());
let ep_usage = Route::new().at(
"/",
get(poem::endpoint::make_sync(move |_| {
HttpHandlerKind::Query.usage(sock)
})),
);
let ep_health = Route::new().at("/", get(poem::endpoint::make_sync(move |_| "ok")));
let ep = match self.kind {
HttpHandlerKind::Query => Route::new()
.at("/", ep_usage)
.nest("/health", ep_health)
.nest("/v1", query_route())
.nest("/clickhouse", ep_clickhouse),
HttpHandlerKind::Clickhouse => Route::new()
.nest("/", ep_clickhouse)
.nest("/health", ep_health),
};
ep.with(NormalizePath::new(TrailingSlash::Trim))
.with(CatchPanic::new().with_handler(PanicHandler::new()))
.around(json_response)
.boxed()
}
fn build_tls(config: &InnerConfig) -> Result<OpensslTlsConfig, std::io::Error> {
let cfg = OpensslTlsConfig::new()
.cert_from_file(config.query.http_handler_tls_server_cert.as_str())
.key_from_file(config.query.http_handler_tls_server_key.as_str());
// if Path::new(&config.query.http_handler_tls_server_root_ca_cert).exists() {
// cfg = cfg.client_auth_required(std::fs::read(
// config.query.http_handler_tls_server_root_ca_cert.as_str(),
// )?);
// }
Ok(cfg)
}
#[async_backtrace::framed]
async fn start_with_tls(&mut self, listening: SocketAddr) -> Result<SocketAddr, HttpError> {
info!("Http Handler TLS enabled");
let config = GlobalConfig::instance();
let tls_config = Self::build_tls(config.as_ref())
.map_err(|e: std::io::Error| HttpError::TlsConfigError(AnyError::new(&e)))?;
let router = self.build_router(listening).await;
self.shutdown_handler
.start_service(
listening,
Some(tls_config),
router,
Some(Duration::from_millis(1000)),
)
.await
}
#[async_backtrace::framed]
async fn start_without_tls(&mut self, listening: SocketAddr) -> Result<SocketAddr, HttpError> {
let router = self.build_router(listening).await;
self.shutdown_handler
.start_service(listening, None, router, Some(Duration::from_millis(1000)))
.await
}
}
#[async_trait::async_trait]
impl Server for HttpHandler {
#[async_backtrace::framed]
async fn shutdown(&mut self, graceful: bool) {
self.shutdown_handler.shutdown(graceful).await;
}
#[async_backtrace::framed]
async fn start(&mut self, listening: SocketAddr) -> Result<SocketAddr, ErrorCode> {
let config = GlobalConfig::instance();
let res = match config.query.http_handler_tls_server_key.is_empty()
|| config.query.http_handler_tls_server_cert.is_empty()
{
true => self.start_without_tls(listening).await,
false => self.start_with_tls(listening).await,
};
res.map_err(|e: HttpError| match e {
HttpError::BadAddressFormat(any_err) => {
ErrorCode::BadAddressFormat(any_err.to_string())
}
le @ HttpError::ListenError { .. } => ErrorCode::CannotListenerPort(le.to_string()),
HttpError::TlsConfigError(any_err) => {
ErrorCode::TLSConfigurationFailure(any_err.to_string())
}
})
}
}
---
Usage is in http_service file.
<file>src/query/service/src/servers/admin/admin_service.rs</file>
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::net::SocketAddr;
use std::time::Duration;
use databend_common_config::InnerConfig;
use databend_common_exception::ErrorCode;
use databend_common_http::health_handler;
use databend_common_http::home::debug_home_handler;
#[cfg(feature = "memory-profiling")]
use databend_common_http::jeprof::debug_jeprof_dump_handler;
use databend_common_http::pprof::debug_pprof_handler;
use databend_common_http::stack::debug_dump_stack;
use databend_common_http::HttpError;
use databend_common_http::HttpShutdownHandler;
use databend_common_meta_types::anyerror::AnyError;
use log::info;
use log::warn;
use poem::get;
use poem::listener::OpensslTlsConfig;
use poem::post;
use poem::Endpoint;
use poem::Route;
use crate::servers::Server;
pub struct AdminService {
config: InnerConfig,
shutdown_handler: HttpShutdownHandler,
}
impl AdminService {
pub fn create(config: &InnerConfig) -> Box<AdminService> {
Box::new(AdminService {
config: config.clone(),
shutdown_handler: HttpShutdownHandler::create("http api".to_string()),
})
}
fn build_router(&self) -> impl Endpoint {
#[cfg_attr(not(feature = "memory-profiling"), allow(unused_mut))]
let mut route = Route::new()
.at("/v1/health", get(health_handler))
.at("/v1/config", get(super::v1::config::config_handler))
.at("/v1/system", get(super::v1::system::system_handler))
.at(
"/v1/status",
get(super::v1::instance_status::instance_status_handler),
)
.at(
"/v1/processlist",
get(super::v1::processes::processlist_handler),
)
.at(
"/v1/tables",
get(super::v1::tenant_tables::list_tables_handler),
)
.at(
"/v1/tables/stats",
get(super::v1::tenant_table_stats::get_tables_stats_handler),
)
.at(
"/v1/cluster/list",
get(super::v1::cluster::cluster_list_handler),
)
.at(
"v1/queries/:query_id/profiling",
get(super::v1::query_profiling::query_profiling_handler),
)
.at("/debug/home", get(debug_home_handler))
.at("/debug/pprof/profile", get(debug_pprof_handler))
.at("/debug/async_tasks/dump", get(debug_dump_stack));
// Multiple tenants admin api
if self.config.query.management_mode {
route = route
.at(
"/v1/tenants/:tenant/tables",
get(super::v1::tenant_tables::list_tenant_tables_handler),
)
.at(
"/v1/tenants/:tenant/tables/stats",
get(super::v1::tenant_table_stats::get_tenant_tables_stats_handler),
)
.at(
"v1/tenants/:tenant/stream_status",
get(super::v1::stream_status::stream_status_handler),
)
.at(
"/v1/tenants/:tenant/settings",
get(super::v1::settings::list_settings),
)
.at(
"/v1/tenants/:tenant/settings/:key",
post(super::v1::settings::set_settings)
.delete(super::v1::settings::unset_settings),
)
.at(
"/v1/tenants/:tenant/user_functions",
get(super::v1::user_functions::user_functions),
);
}
#[cfg(feature = "memory-profiling")]
{
route = route.at(
// to follow the conversions of jeprof, we arrange the path in
// this way, so that jeprof could be invoked like:
// `jeprof ./target/debug/databend-query http://localhost:8080/debug/mem`
// and jeprof will translate the above url into sth like:
// "http://localhost:8080/debug/mem/pprof/profile?seconds=30"
"/debug/mem/pprof/profile",
get(debug_jeprof_dump_handler),
);
};
route
}
fn build_tls(config: &InnerConfig) -> Result<OpensslTlsConfig, std::io::Error> {
let cfg = OpensslTlsConfig::new()
.cert_from_file(config.query.api_tls_server_cert.as_str())
.key_from_file(config.query.api_tls_server_key.as_str());
// if Path::new(&config.query.api_tls_server_root_ca_cert).exists() {
// cfg = cfg.client_auth_required(std::fs::read(
// config.query.api_tls_server_root_ca_cert.as_str(),
// )?);
// }
Ok(cfg)
}
#[async_backtrace::framed]
async fn start_with_tls(&mut self, listening: SocketAddr) -> Result<SocketAddr, HttpError> {
info!("Http API TLS enabled");
let tls_config = Self::build_tls(&self.config)
.map_err(|e| HttpError::TlsConfigError(AnyError::new(&e)))?;
let addr = self
.shutdown_handler
.start_service(
listening,
Some(tls_config),
self.build_router(),
Some(Duration::from_millis(1000)),
)
.await?;
Ok(addr)
}
#[async_backtrace::framed]
async fn start_without_tls(&mut self, listening: SocketAddr) -> Result<SocketAddr, HttpError> {
warn!("Http API TLS not set");
let addr = self
.shutdown_handler
.start_service(
listening,
None,
self.build_router(),
Some(Duration::from_millis(1000)),
)
.await?;
Ok(addr)
}
}
#[async_trait::async_trait]
impl Server for AdminService {
#[async_backtrace::framed]
async fn shutdown(&mut self, _graceful: bool) {
// intendfully do nothing: sometimes we hope to diagnose the backtraces or metrics after
// the process got the sigterm signal, we can still leave the admin service port open until
// the process exited. it's not an user facing service, it's allowed to force shutdown.
}
#[async_backtrace::framed]
async fn start(&mut self, listening: SocketAddr) -> Result<SocketAddr, ErrorCode> {
let config = &self.config.query;
let res =
match config.api_tls_server_key.is_empty() || config.api_tls_server_cert.is_empty() {
true => self.start_without_tls(listening).await,
false => self.start_with_tls(listening).await,
};
res.map_err(|e: HttpError| match e {
HttpError::BadAddressFormat(any_err) => {
ErrorCode::BadAddressFormat(any_err.to_string())
}
le @ HttpError::ListenError { .. } => ErrorCode::CannotListenerPort(le.to_string()),
HttpError::TlsConfigError(any_err) => {
ErrorCode::TLSConfigurationFailure(any_err.to_string())
}
})
}
}