hacky: ability to set io_engine via mgmt_api => pagebench

This commit is contained in:
Christian Schwarz
2024-01-29 10:50:20 +00:00
parent d23ea718ee
commit 8d6ce71b29
5 changed files with 49 additions and 12 deletions

View File

@@ -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)
}
}

View File

@@ -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<Utf8PathBuf>,
#[clap(long)]
set_io_engine: Option<String>,
targets: Option<Vec<TenantTimelineId>>,
}
@@ -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<TenantTimelineId> = crate::util::cli::targets::discover(
&mgmt_api_client,

View File

@@ -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<Body>,
_cancel: CancellationToken,
) -> Result<Response<Body>, 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))
}

View File

@@ -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::*;

View File

@@ -26,23 +26,31 @@ pub enum IoEngineKind {
TokioEpollUring,
}
static IO_ENGINE: once_cell::sync::OnceCell<IoEngineKind> = once_cell::sync::OnceCell::new();
static IO_ENGINE: std::sync::RwLock<Option<IoEngineKind>> = 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::<IoEngineKind>() {
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;