pagebench: add a 'getpage@lsn' benchmark

libs/pageserver_api/src/shard.rs

@@ -81,6 +81,10 @@ impl TenantShardId {
     pub fn is_zero(&self) -> bool {
         self.shard_number == ShardNumber(0)
     }
+
+    pub fn is_unsharded(&self) -> bool {
+        self.shard_number == ShardNumber(0) && self.shard_count == ShardCount(0)
+    }
 }
 
 /// Formatting helper

pageserver/pagebench/src/basebackup.rs

@@ -1,12 +1,12 @@
 use anyhow::Context;
 use pageserver_client::page_service::BasebackupRequest;
+use utils::id::TenantId;
 use utils::lsn::Lsn;
 
 use rand::prelude::*;
 use tokio::sync::Barrier;
 use tokio::task::JoinSet;
 use tracing::{debug, info, instrument};
-use utils::id::TenantShardId;
 use utils::logging;
 
 use std::cell::RefCell;
@@ -189,12 +189,16 @@ async fn main_impl(
     if args.targets.is_some() {
         timelines = args.targets.clone().unwrap();
     } else {
-        let tenants: Vec<TenantShardId> = mgmt_api_client
-            .list_tenants()
-            .await?
-            .into_iter()
-            .map(|ti| ti.id)
-            .collect();
+        let mut tenants: Vec<TenantId> = Vec::new();
+        for ti in mgmt_api_client.list_tenants().await? {
+            if !ti.id.is_unsharded() {
+                anyhow::bail!(
+                    "only unsharded tenants are supported at this time: {}",
+                    ti.id
+                );
+            }
+            tenants.push(ti.id.tenant_id)
+        }
         let mut js = JoinSet::new();
         for tenant_id in tenants {
             js.spawn({
@@ -211,7 +215,7 @@ async fn main_impl(
             let (tenant_id, tl_infos) = res.unwrap();
             for tl in tl_infos {
                 timelines.push(TenantTimelineId {
-                    tenant_shard_id: tenant_id,
+                    tenant_id,
                     timeline_id: tl.timeline_id,
                 });
             }
@@ -225,7 +229,7 @@ async fn main_impl(
         js.spawn({
             let timeline = *timeline;
             let info = mgmt_api_client
-                .timeline_info(timeline.tenant_shard_id, timeline.timeline_id)
+                .timeline_info(timeline.tenant_id, timeline.timeline_id)
                 .await
                 .unwrap();
             async move {
@@ -357,19 +361,18 @@ async fn client(
 ) {
     start_work_barrier.wait().await;
 
-    let client =
-        pageserver_client::page_service::Client::new(crate::util::connstring::connstring(
-            &args.page_service_host_port,
-            args.pageserver_jwt.as_deref(),
-        ))
-        .await
-        .unwrap();
+    let client = pageserver_client::page_service::Client::new(crate::util::connstring::connstring(
+        &args.page_service_host_port,
+        args.pageserver_jwt.as_deref(),
+    ))
+    .await
+    .unwrap();
 
     while let Some(Work { lsn, gzip }) = work.recv().await {
         let start = Instant::now();
         let copy_out_stream = client
             .basebackup(&BasebackupRequest {
-                tenant_id: timeline.tenant_shard_id,
+                tenant_id: timeline.tenant_id,
                 timeline_id: timeline.timeline_id,
                 lsn,
                 gzip,

pageserver/pagebench/src/getpage_latest_lsn.rs

@@ -3,15 +3,14 @@ use futures::future::join_all;
 use pageserver::pgdatadir_mapping::key_to_rel_block;
 use pageserver::repository;
 use pageserver_api::key::is_rel_block_key;
-use pageserver_api::shard::TenantShardId;
 use pageserver_client::page_service::RelTagBlockNo;
+use utils::id::TenantId;
 use utils::lsn::Lsn;
 
 use rand::prelude::*;
 use tokio::sync::Barrier;
 use tokio::task::JoinSet;
 use tracing::{info, instrument};
-use utils::id::TenantShardId;
 use utils::logging;
 
 use std::cell::RefCell;
@@ -202,12 +201,16 @@ async fn main_impl(
     if args.targets.is_some() {
         timelines = args.targets.clone().unwrap();
     } else {
-        let tenants: Vec<TenantShardId> = mgmt_api_client
-            .list_tenants()
-            .await?
-            .into_iter()
-            .map(|ti| ti.id)
-            .collect();
+        let mut tenants: Vec<TenantId> = Vec::new();
+        for ti in mgmt_api_client.list_tenants().await? {
+            if !ti.id.is_unsharded() {
+                anyhow::bail!(
+                    "only unsharded tenants are supported at this time: {}",
+                    ti.id
+                );
+            }
+            tenants.push(ti.id.tenant_id)
+        }
         let mut js = JoinSet::new();
         for tenant_id in tenants {
             js.spawn({
@@ -224,7 +227,7 @@ async fn main_impl(
             let (tenant_id, tl_infos) = res.unwrap();
             for tl in tl_infos {
                 timelines.push(TenantTimelineId {
-                    tenant_shard_id: tenant_id,
+                    tenant_id,
                     timeline_id: tl.timeline_id,
                 });
             }
@@ -240,7 +243,7 @@ async fn main_impl(
             let timeline = *timeline;
             async move {
                 let partitioning = mgmt_api_client
-                    .keyspace(timeline.tenant_shard_id, timeline.timeline_id)
+                    .keyspace(timeline.tenant_id, timeline.timeline_id)
                     .await?;
                 let lsn = partitioning.at_lsn;
 
@@ -443,7 +446,7 @@ async fn client(
         .await
         .unwrap();
     let mut client = client
-        .pagestream(timeline.tenant_shard_id, timeline.timeline_id)
+        .pagestream(timeline.tenant_id, timeline.timeline_id)
         .await
         .unwrap();
 

pageserver/pagebench/src/util/tenant_timeline_id.rs

@@ -1,12 +1,11 @@
 use std::str::FromStr;
 
 use anyhow::Context;
-use pageserver_api::shard::TenantShardId;
-use utils::id::TimelineId;
+use utils::id::{TenantId, TimelineId};
 
 #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
 pub(crate) struct TenantTimelineId {
-    pub(crate) tenant_shard_id: TenantShardId,
+    pub(crate) tenant_id: TenantId,
     pub(crate) timeline_id: TimelineId,
 }
 
@@ -17,12 +16,12 @@ impl FromStr for TenantTimelineId {
         let (tenant_id, timeline_id) = s
             .split_once("/")
             .context("tenant and timeline id must be separated by `/`")?;
-        let tenant_id = TenantShardId::from_str(&tenant_id)
+        let tenant_id = TenantId::from_str(&tenant_id)
             .with_context(|| format!("invalid tenant id: {tenant_id:?}"))?;
         let timeline_id = TimelineId::from_str(&timeline_id)
             .with_context(|| format!("invalid timeline id: {timeline_id:?}"))?;
         Ok(Self {
-            tenant_shard_id: tenant_id,
+            tenant_id,
             timeline_id,
         })
     }
@@ -30,6 +29,6 @@ impl FromStr for TenantTimelineId {
 
 impl std::fmt::Display for TenantTimelineId {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}/{}", self.tenant_shard_id, self.timeline_id)
+        write!(f, "{}/{}", self.tenant_id, self.timeline_id)
     }
 }