From 3a5b44ea539b62f3b0a30468badf803cbe885cbd Mon Sep 17 00:00:00 2001
From: Yuchen Liang <yuchen@neon.tech>
Date: Tue, 1 Oct 2024 08:16:18 -0400
Subject: [PATCH] add set_io_mode option to getpage_latest_lsn

Signed-off-by: Yuchen Liang <yuchen@neon.tech>
---
 libs/pageserver_api/src/models.rs                 | 15 +++++++++++++--
 pageserver/client/src/mgmt_api.rs                 |  5 ++++-
 .../pagebench/src/cmd/getpage_latest_lsn.rs       |  8 ++++++++
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/libs/pageserver_api/src/models.rs b/libs/pageserver_api/src/models.rs
index 67da2783f5..3ec9cac2c3 100644
--- a/libs/pageserver_api/src/models.rs
+++ b/libs/pageserver_api/src/models.rs
@@ -992,8 +992,19 @@ pub mod virtual_file {
     }
 
     /// Direct IO modes for a pageserver.
-    #[derive(Debug, PartialEq, Eq, Clone, serde::Deserialize, serde::Serialize)]
-    #[serde(rename_all = "kebab-case", deny_unknown_fields)]
+    #[derive(
+        Copy,
+        Clone,
+        PartialEq,
+        Eq,
+        Hash,
+        strum_macros::EnumString,
+        strum_macros::Display,
+        serde_with::DeserializeFromStr,
+        serde_with::SerializeDisplay,
+        Debug,
+    )]
+    #[strum(serialize_all = "kebab-case")]
     #[repr(u8)]
     pub enum IoMode {
         /// Uses buffered IO.
diff --git a/pageserver/client/src/mgmt_api.rs b/pageserver/client/src/mgmt_api.rs
index 4c54f4274b..2028fa8f30 100644
--- a/pageserver/client/src/mgmt_api.rs
+++ b/pageserver/client/src/mgmt_api.rs
@@ -551,7 +551,10 @@ impl Client {
     }
 
     /// Configs io mode at runtime.
-    pub async fn put_io_mode(&self, mode: virtual_file::IoMode) -> Result<()> {
+    pub async fn put_io_mode(
+        &self,
+        mode: &pageserver_api::models::virtual_file::IoMode,
+    ) -> Result<()> {
         let uri = format!("{}/v1/io_alignment", self.mgmt_api_endpoint);
         self.request(Method::PUT, uri, mode)
             .await?
diff --git a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs
index ac4a732377..603a8b118a 100644
--- a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs
+++ b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs
@@ -63,6 +63,10 @@ pub(crate) struct Args {
     #[clap(long)]
     set_io_alignment: Option<usize>,
 
+    /// Before starting the benchmark, live-reconfigure the pageserver to use specified io mode (buffered vs. direct).
+    #[clap(long)]
+    set_io_mode: Option<pageserver_api::models::virtual_file::IoMode>,
+
     targets: Option<Vec<TenantTimelineId>>,
 }
 
@@ -133,6 +137,10 @@ async fn main_impl(
         mgmt_api_client.put_io_alignment(align).await?;
     }
 
+    if let Some(mode) = &args.set_io_mode {
+        mgmt_api_client.put_io_mode(mode).await?;
+    }
+
     // discover targets
     let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(
         &mgmt_api_client,