Add script to create safekeeper timelines

pageserver/src/repository.rs

@@ -28,9 +28,8 @@ impl Key {
     /// As long as Neon does not support tablespace (because of lack of access to local file system),
     /// we can assume that only some predefined namespace OIDs are used which can fit in u16
     pub fn to_i128(&self) -> i128 {
-        assert!(self.field1 < 0xF);
-        assert!(self.field2 < 0xFFFF);
-        (((self.field1 & 0xF) as i128) << 120)
+        assert!(self.field2 < 0xFFFF || self.field2 == 0xFFFFFFFF || self.field2 == 0x22222222);
+        (((self.field1 & 0xf) as i128) << 120)
             | (((self.field2 & 0xFFFF) as i128) << 104)
             | ((self.field3 as i128) << 72)
             | ((self.field4 as i128) << 40)
@@ -150,8 +149,8 @@ impl Key {
         field6: u32::MIN,
     };
     pub const MAX: Key = Key {
-        field1: 0xF - 1,
-        field2: 0xFFFF - 1,
+        field1: u8::MAX,
+        field2: u32::MAX,
         field3: u32::MAX,
         field4: u32::MAX,
         field5: u8::MAX,

pageserver/src/tenant.rs

@@ -3692,7 +3692,7 @@ mod tests {
     use tokio_util::sync::CancellationToken;
 
     static TEST_KEY: Lazy<Key> =
-        Lazy::new(|| Key::from_slice(&hex!("010000000033333333444444445500000001")));
+        Lazy::new(|| Key::from_slice(&hex!("112222222233333333444444445500000001")));
 
     #[tokio::test]
     async fn test_basic() -> anyhow::Result<()> {
@@ -3788,9 +3788,9 @@ mod tests {
         let writer = tline.writer().await;
 
         #[allow(non_snake_case)]
-        let TEST_KEY_A: Key = Key::from_hex("110000000033333333444444445500000001").unwrap();
+        let TEST_KEY_A: Key = Key::from_hex("112222222233333333444444445500000001").unwrap();
         #[allow(non_snake_case)]
-        let TEST_KEY_B: Key = Key::from_hex("110000000033333333444444445500000002").unwrap();
+        let TEST_KEY_B: Key = Key::from_hex("112222222233333333444444445500000002").unwrap();
 
         // Insert a value on the timeline
         writer
@@ -4374,7 +4374,7 @@ mod tests {
 
         let mut keyspace = KeySpaceAccum::new();
 
-        let mut test_key = Key::from_hex("010000000033333333444444445500000000").unwrap();
+        let mut test_key = Key::from_hex("012222222233333333444444445500000000").unwrap();
         let mut blknum = 0;
         for _ in 0..50 {
             for _ in 0..10000 {
@@ -4420,7 +4420,7 @@ mod tests {
 
         const NUM_KEYS: usize = 1000;
 
-        let mut test_key = Key::from_hex("010000000033333333444444445500000000").unwrap();
+        let mut test_key = Key::from_hex("012222222233333333444444445500000000").unwrap();
 
         let mut keyspace = KeySpaceAccum::new();
 
@@ -4501,7 +4501,7 @@ mod tests {
 
         const NUM_KEYS: usize = 1000;
 
-        let mut test_key = Key::from_hex("010000000033333333444444445500000000").unwrap();
+        let mut test_key = Key::from_hex("012222222233333333444444445500000000").unwrap();
 
         let mut keyspace = KeySpaceAccum::new();
 
@@ -4592,7 +4592,7 @@ mod tests {
         const NUM_KEYS: usize = 100;
         const NUM_TLINES: usize = 50;
 
-        let mut test_key = Key::from_hex("010000000033333333444444445500000000").unwrap();
+        let mut test_key = Key::from_hex("012222222233333333444444445500000000").unwrap();
         // Track page mutation lsns across different timelines.
         let mut updated = [[Lsn(0); NUM_KEYS]; NUM_TLINES];
 

pageserver/src/tenant/storage_layer/inmemory_layer.rs

@@ -345,19 +345,14 @@ impl InMemoryLayer {
 
         let cursor = inner.file.block_cursor();
 
-        // Sort the keys because delta layer writer expects them sorted.
-        //
-        // NOTE: this sort can take up significant time if the layer has millions of
-        //       keys. To speed up all the comparisons we convert the key to i128 and
-        //       keep the value as a reference.
-        let mut keys: Vec<_> = inner.index.iter().map(|(k, m)| (k.to_i128(), m)).collect();
-        keys.sort_unstable_by_key(|k| k.0);
+        let mut keys: Vec<(&Key, &VecMap<Lsn, u64>)> = inner.index.iter().collect();
+        keys.sort_by_key(|k| k.0);
 
         let ctx = RequestContextBuilder::extend(ctx)
             .page_content_kind(PageContentKind::InMemoryLayer)
             .build();
         for (key, vec_map) in keys.iter() {
-            let key = Key::from_i128(*key);
+            let key = **key;
             // Write all page versions
             for (lsn, pos) in vec_map.as_slice() {
                 cursor.read_blob_into_buf(*pos, &mut buf, &ctx).await?;

pageserver/src/tenant/storage_layer/layer.rs

@@ -337,16 +337,18 @@ enum ResidentOrWantedEvicted {
 }
 
 impl ResidentOrWantedEvicted {
-    fn get_and_upgrade(&mut self) -> Option<(Arc<DownloadedLayer>, bool)> {
+    /// If `Some` is returned, the ResidentOrWantedEvicted has been upgraded back from
+    /// `ResidentOrWantedEvicted::WantedEvicted` to `ResidentOrWantedEvicted::Resident`.
+    fn get_and_upgrade(&mut self) -> Option<Arc<DownloadedLayer>> {
         match self {
-            ResidentOrWantedEvicted::Resident(strong) => Some((strong.clone(), false)),
+            ResidentOrWantedEvicted::Resident(strong) => Some(strong.clone()),
             ResidentOrWantedEvicted::WantedEvicted(weak, _) => match weak.upgrade() {
                 Some(strong) => {
                     LAYER_IMPL_METRICS.inc_raced_wanted_evicted_accesses();
 
                     *self = ResidentOrWantedEvicted::Resident(strong.clone());
 
-                    Some((strong, true))
+                    Some(strong)
                 }
                 None => None,
             },
@@ -635,16 +637,14 @@ impl LayerInner {
 
                 // check if we really need to be downloaded; could have been already downloaded by a
                 // cancelled previous attempt.
+                //
+                // FIXME: what if it's a directory? that is currently needs_download == true
                 let needs_download = self
                     .needs_download()
                     .await
                     .map_err(DownloadError::PreStatFailed)?;
 
                 let permit = if let Some(reason) = needs_download {
-                    if let NeedsDownload::NotFile(ft) = reason {
-                        return Err(DownloadError::NotFile(ft));
-                    }
-
                     // only reset this after we've decided we really need to download. otherwise it'd
                     // be impossible to mark cancelled downloads for eviction, like one could imagine
                     // we would like to do for prefetching which was not needed.
@@ -694,7 +694,7 @@ impl LayerInner {
                 // below paths anymore essentially limiting the max loop iterations to 2.
                 let (value, init_permit) = download(init_permit).await?;
                 let mut guard = self.inner.set(value, init_permit);
-                let (strong, _upgraded) = guard
+                let strong = guard
                     .get_and_upgrade()
                     .expect("init creates strong reference, we held the init permit");
                 return Ok(strong);
@@ -703,17 +703,11 @@ impl LayerInner {
             let (weak, permit) = {
                 let mut locked = self.inner.get_or_init(download).await?;
 
-                if let Some((strong, upgraded)) = locked.get_and_upgrade() {
-                    if upgraded {
-                        // when upgraded back, the Arc<DownloadedLayer> is still available, but
-                        // previously a `evict_and_wait` was received.
-                        self.wanted_evicted.store(false, Ordering::Relaxed);
+                if let Some(strong) = locked.get_and_upgrade() {
+                    self.wanted_evicted.store(false, Ordering::Relaxed);
 
-                        // error out any `evict_and_wait`
-                        drop(self.status.send(Status::Downloaded));
-                        LAYER_IMPL_METRICS
-                            .inc_eviction_cancelled(EvictionCancelled::UpgradedBackOnAccess);
-                    }
+                    // error out any `evict_and_wait`
+                    drop(self.status.send(Status::Downloaded));
 
                     return Ok(strong);
                 } else {
@@ -889,7 +883,7 @@ impl LayerInner {
     fn is_file_present_and_good_size(&self, m: &std::fs::Metadata) -> Result<(), NeedsDownload> {
         // in future, this should include sha2-256 validation of the file.
         if !m.is_file() {
-            Err(NeedsDownload::NotFile(m.file_type()))
+            Err(NeedsDownload::NotFile)
         } else if m.len() != self.desc.file_size {
             Err(NeedsDownload::WrongSize {
                 actual: m.len(),
@@ -1088,8 +1082,6 @@ enum DownloadError {
     ContextAndConfigReallyDeniesDownloads,
     #[error("downloading is really required but not allowed by this method")]
     DownloadRequired,
-    #[error("layer path exists, but it is not a file: {0:?}")]
-    NotFile(std::fs::FileType),
     /// Why no error here? Because it will be reported by page_service. We should had also done
     /// retries already.
     #[error("downloading evicted layer file failed")]
@@ -1105,7 +1097,7 @@ enum DownloadError {
 #[derive(Debug, PartialEq)]
 pub(crate) enum NeedsDownload {
     NotFound,
-    NotFile(std::fs::FileType),
+    NotFile,
     WrongSize { actual: u64, expected: u64 },
 }
 
@@ -1113,7 +1105,7 @@ impl std::fmt::Display for NeedsDownload {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             NeedsDownload::NotFound => write!(f, "file was not found"),
-            NeedsDownload::NotFile(ft) => write!(f, "path is not a file; {ft:?}"),
+            NeedsDownload::NotFile => write!(f, "path is not a file"),
             NeedsDownload::WrongSize { actual, expected } => {
                 write!(f, "file size mismatch {actual} vs. {expected}")
             }
@@ -1509,8 +1501,6 @@ enum EvictionCancelled {
     AlreadyReinitialized,
     /// Not evicted because of a pending reinitialization
     LostToDownload,
-    /// After eviction, there was a new layer access which cancelled the eviction.
-    UpgradedBackOnAccess,
 }
 
 impl EvictionCancelled {
@@ -1523,7 +1513,6 @@ impl EvictionCancelled {
             EvictionCancelled::RemoveFailed => "remove_failed",
             EvictionCancelled::AlreadyReinitialized => "already_reinitialized",
             EvictionCancelled::LostToDownload => "lost_to_download",
-            EvictionCancelled::UpgradedBackOnAccess => "upgraded_back_on_access",
         }
     }
 }

safekeeper/Cargo.toml

@@ -50,3 +50,7 @@ storage_broker.workspace = true
 utils.workspace = true
 
 workspace_hack.workspace = true
+
+[[bench]]
+name = "sk_capacity"
+harness = false

safekeeper/benches/sk_capacity.rs

@@ -0,0 +1,64 @@
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use clap::Parser;
+use safekeeper_api::models::TimelineCreateRequest;
+use utils::id::TenantTimelineId;
+use utils::lsn::Lsn;
+
+const ABOUT: &str = r#"
+Creates many random timelines on the safekeeper.
+
+For example,
+cargo build -r -p safekeeper && target/release/safekeeper
+cargo bench --bench sk_capacity -- -n 1000 --http-addr=http://127.0.0.1:7676
+"#;
+
+#[derive(Parser, Debug)]
+#[clap(author, version, about = ABOUT)]
+struct Args {
+    /// Number of timelines to create
+    #[clap(short = 'n', long, value_parser, default_value_t = 1)]
+    num_timelines: u64,
+    /// HTTP safekeeper address
+    #[clap(long)]
+    http_addr: String,
+    // Fake value to satisfy `cargo bench` passing it.
+    #[clap(long)]
+    bench: bool,
+}
+
+async fn create_timeline(args: &Args) -> Result<(), Box<dyn std::error::Error>> {
+    let client = reqwest::Client::new();
+
+    let ttid = TenantTimelineId::generate();
+
+    let request = TimelineCreateRequest {
+        tenant_id: ttid.tenant_id,
+        timeline_id: ttid.timeline_id,
+        peer_ids: None,
+        pg_version: 160000,
+        system_id: None,
+        wal_seg_size: None,
+        commit_lsn: Lsn(21623024),
+        local_start_lsn: None,
+    };
+
+    // Send request to /v1/tenant/timeline
+    let url = format!("{}/v1/tenant/timeline", args.http_addr);
+    let res = client.post(url).json(&request).send().await?;
+
+    println!("Response: {:?}", res.status());
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args = Args::parse();
+
+    for i in 0..args.num_timelines {
+        create_timeline(&args).await?;
+    }
+    Ok(())
+}