From 8d6ce71b297c5fe5dc4b94f46155f0d2abd9800e Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Mon, 29 Jan 2024 10:50:20 +0000 Subject: [PATCH] hacky: ability to set io_engine via mgmt_api => pagebench --- pageserver/client/src/mgmt_api.rs | 9 +++++ .../pagebench/src/cmd/getpage_latest_lsn.rs | 6 ++++ pageserver/src/http/routes.rs | 11 +++++++ pageserver/src/virtual_file.rs | 2 +- pageserver/src/virtual_file/io_engine.rs | 33 ++++++++++++------- 5 files changed, 49 insertions(+), 12 deletions(-) diff --git a/pageserver/client/src/mgmt_api.rs b/pageserver/client/src/mgmt_api.rs index 077c3909e1..c7a6079e2c 100644 --- a/pageserver/client/src/mgmt_api.rs +++ b/pageserver/client/src/mgmt_api.rs @@ -275,4 +275,13 @@ impl Client { .await .map_err(Error::ReceiveBody) } + + pub async fn set_io_engine(&self, engine_str: &str) -> Result<()> { + let uri = format!("{}/v1/set_io_engine", self.mgmt_api_endpoint); + self.request(Method::PUT, uri, engine_str) + .await? + .json() + .await + .map_err(Error::ReceiveBody) + } } diff --git a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs index 400b5476b7..e54e3f6e35 100644 --- a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs +++ b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs @@ -51,6 +51,8 @@ pub(crate) struct Args { /// It doesn't get invalidated if the keyspace changes under the hood, e.g., due to new ingested data or compaction. #[clap(long)] keyspace_cache: Option, + #[clap(long)] + set_io_engine: Option, targets: Option>, } @@ -103,6 +105,10 @@ async fn main_impl( args.pageserver_jwt.as_deref(), )); + if let Some(engine_str) = &args.set_io_engine { + mgmt_api_client.set_io_engine(engine_str).await?; + } + // discover targets let timelines: Vec = crate::util::cli::targets::discover( &mgmt_api_client, diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index aa56806246..e820dd6dce 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -2052,5 +2052,16 @@ pub fn make_router( "/v1/tenant/:tenant_shard_id/timeline/:timeline_id/keyspace", |r| testing_api_handler("read out the keyspace", r, timeline_collect_keyspace), ) + .put("/v1/set_io_engine", |r| { + async fn set_io_engine_handler( + mut r: Request, + _cancel: CancellationToken, + ) -> Result, ApiError> { + let kind: crate::virtual_file::IoEngineKind = json_request(&mut r).await?; + crate::virtual_file::io_engine::set(kind); + json_response(StatusCode::OK, ()) + } + api_handler(r, set_io_engine_handler) + }) .any(handler_404)) } diff --git a/pageserver/src/virtual_file.rs b/pageserver/src/virtual_file.rs index d200a4ba5e..04435f8731 100644 --- a/pageserver/src/virtual_file.rs +++ b/pageserver/src/virtual_file.rs @@ -28,7 +28,7 @@ use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; use tokio::time::Instant; use utils::fs_ext; -mod io_engine; +pub(crate) mod io_engine; mod open_options; pub use io_engine::IoEngineKind; pub(crate) use open_options::*; diff --git a/pageserver/src/virtual_file/io_engine.rs b/pageserver/src/virtual_file/io_engine.rs index f7b46fe653..b55106938f 100644 --- a/pageserver/src/virtual_file/io_engine.rs +++ b/pageserver/src/virtual_file/io_engine.rs @@ -26,23 +26,31 @@ pub enum IoEngineKind { TokioEpollUring, } -static IO_ENGINE: once_cell::sync::OnceCell = once_cell::sync::OnceCell::new(); +static IO_ENGINE: std::sync::RwLock> = std::sync::RwLock::new(None); + +pub(crate) fn set(engine: IoEngineKind) { + let mut guard = IO_ENGINE.write().unwrap(); + *guard = Some(engine); + let metric = &crate::metrics::virtual_file_io_engine::KIND; + metric.reset(); + metric.with_label_values(&[&format!("{engine}")]).set(1); + drop(guard); +} #[cfg(not(test))] pub(super) fn init(engine: IoEngineKind) { - if IO_ENGINE.set(engine).is_err() { - panic!("called twice"); - } - crate::metrics::virtual_file_io_engine::KIND - .with_label_values(&[&format!("{engine}")]) - .set(1); + set(engine); } -pub(super) fn get() -> &'static IoEngineKind { +pub(super) fn get() -> IoEngineKind { #[cfg(test)] { let env_var_name = "NEON_PAGESERVER_UNIT_TEST_VIRTUAL_FILE_IOENGINE"; - IO_ENGINE.get_or_init(|| match std::env::var(env_var_name) { + let guard = IO_ENGINE.read().unwrap(); + if let Some(v) = guard.is_some() { + return v; + } + *guard = Some(match std::env::var(env_var_name) { Ok(v) => match v.parse::() { Ok(engine_kind) => engine_kind, Err(e) => { @@ -57,10 +65,13 @@ pub(super) fn get() -> &'static IoEngineKind { Err(std::env::VarError::NotUnicode(_)) => { panic!("env var {env_var_name} is not unicode"); } - }) + }); } #[cfg(not(test))] - IO_ENGINE.get().unwrap() + IO_ENGINE + .read() + .unwrap() + .expect("should have called set() or init() before") } use std::os::unix::prelude::FileExt;