From e1c1aa74fe7a9e080f10c5c638fc9bc61184fcdf Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Wed, 4 Jun 2025 18:28:53 +0200
Subject: [PATCH] get benchmark to work

---
 libs/sk_ps_discovery/benches/bench.rs | 11 +++++------
 libs/sk_ps_discovery/src/tests.rs     |  1 -
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/libs/sk_ps_discovery/benches/bench.rs b/libs/sk_ps_discovery/benches/bench.rs
index 4c7b4a486d..898026d257 100644
--- a/libs/sk_ps_discovery/benches/bench.rs
+++ b/libs/sk_ps_discovery/benches/bench.rs
@@ -2,15 +2,14 @@
 
 use std::time::Instant;
 
-use criterion::Criterion;
+use criterion::{criterion_group, criterion_main, Criterion};
 use hex::FromHex;
+use pprof::criterion::{Output, PProfProfiler};
 use sk_ps_discovery::{
     AttachmentUpdate, RemoteConsistentLsnAdv, TenantShardAttachmentId, TimelineAttachmentId,
 };
 use utils::{
-    generation::Generation,
-    id::{TenantId, TenantTimelineId, TimelineId},
-    shard::ShardIndex,
+    generation::Generation, id::{NodeId, TenantId, TenantTimelineId, TimelineId}, lsn::Lsn, shard::ShardIndex
 };
 
 /// Use jemalloc and enable profiling, to mirror bin/safekeeper.rs.
@@ -39,12 +38,12 @@ fn bench_simple(c: &mut Criterion) {
     let n_tenants = 400_000;
     for t in 1..=n_tenants {
         let ps_id = NodeId(23);
-        let tenant_id = TenantId::from_hex(format!("{t:x}")).unwrap();
+        let tenant_id = TenantId::generate();
         let timeline_id = TimelineId::generate();
         let tenant_shard_attachment_id = TenantShardAttachmentId {
             tenant_id,
             shard_id: ShardIndex::unsharded(),
-            generation: Generation(0),
+            generation: Generation::Valid(0),
         };
         let timeline_attachment = TimelineAttachmentId {
             tenant_shard_attachment_id,
diff --git a/libs/sk_ps_discovery/src/tests.rs b/libs/sk_ps_discovery/src/tests.rs
index ead1067257..54a36cd4f8 100644
--- a/libs/sk_ps_discovery/src/tests.rs
+++ b/libs/sk_ps_discovery/src/tests.rs
@@ -23,7 +23,6 @@ fn basic() {
     };
 
     let ps1 = NodeId(0x100);
-    let ps2 = NodeId(0x200);
 
     // Out of order; in happy path, commit_lsn advances first, but let's test the
     // case where safekeeper doesn't know about the attachments yet first, before