feat(pageserver): migrate reldir to sparse keyspace

Signed-off-by: Alex Chi Z <chi@neon.tech>

libs/pageserver_api/src/key.rs

@@ -47,6 +47,12 @@ pub const AUX_KEY_PREFIX: u8 = 0x62;
 /// The key prefix of ReplOrigin keys.
 pub const REPL_ORIGIN_KEY_PREFIX: u8 = 0x63;
 
+/// The key prefix of db directory keys.
+pub const DB_DIR_KEY_PREFIX: u8 = 0x64;
+
+/// The key prefix of rel direcotry keys.
+pub const REL_DIR_KEY_PREFIX: u8 = 0x65;
+
 /// Check if the key falls in the range of metadata keys.
 pub const fn is_metadata_key_slice(key: &[u8]) -> bool {
     key[0] >= METADATA_KEY_BEGIN_PREFIX && key[0] < METADATA_KEY_END_PREFIX
@@ -108,6 +114,24 @@ impl Key {
         }
     }
 
+    pub fn rel_dir_sparse_key_range() -> Range<Self> {
+        Key {
+            field1: REL_DIR_KEY_PREFIX,
+            field2: 0,
+            field3: 0,
+            field4: 0,
+            field5: 0,
+            field6: 0,
+        }..Key {
+            field1: REL_DIR_KEY_PREFIX + 1,
+            field2: 0,
+            field3: 0,
+            field4: 0,
+            field5: 0,
+            field6: 0,
+        }
+    }
+
     /// This function checks more extensively what keys we can take on the write path.
     /// If a key beginning with 00 does not have a global/default tablespace OID, it
     /// will be rejected on the write path.
@@ -438,6 +462,36 @@ pub fn rel_dir_to_key(spcnode: Oid, dbnode: Oid) -> Key {
     }
 }
 
+#[inline(always)]
+pub fn rel_tag_sparse_key(spcnode: Oid, dbnode: Oid, relnode: Oid, forknum: u8) -> Key {
+    Key {
+        field1: REL_DIR_KEY_PREFIX,
+        field2: spcnode,
+        field3: dbnode,
+        field4: relnode,
+        field5: forknum,
+        field6: 1,
+    }
+}
+
+pub fn rel_tag_sparse_key_range(spcnode: Oid, dbnode: Oid) -> Range<Key> {
+    Key {
+        field1: REL_DIR_KEY_PREFIX,
+        field2: spcnode,
+        field3: dbnode,
+        field4: 0,
+        field5: 0,
+        field6: 0,
+    }..Key {
+        field1: REL_DIR_KEY_PREFIX,
+        field2: spcnode,
+        field3: dbnode,
+        field4: u32::MAX,
+        field5: u8::MAX,
+        field6: u32::MAX,
+    } // it's fine to exclude the last key b/c we only use field6 == 1
+}
+
 #[inline(always)]
 pub fn rel_block_to_key(rel: RelTag, blknum: BlockNumber) -> Key {
     Key {

pageserver/src/pgdatadir_mapping.rs

@@ -22,13 +22,14 @@ use anyhow::{ensure, Context};
 use bytes::{Buf, Bytes, BytesMut};
 use enum_map::Enum;
 use itertools::Itertools;
-use pageserver_api::key::Key;
 use pageserver_api::key::{
     dbdir_key_range, rel_block_to_key, rel_dir_to_key, rel_key_range, rel_size_to_key,
-    relmap_file_key, repl_origin_key, repl_origin_key_range, slru_block_to_key, slru_dir_to_key,
-    slru_segment_key_range, slru_segment_size_to_key, twophase_file_key, twophase_key_range,
-    CompactKey, AUX_FILES_KEY, CHECKPOINT_KEY, CONTROLFILE_KEY, DBDIR_KEY, TWOPHASEDIR_KEY,
+    rel_tag_sparse_key_range, relmap_file_key, repl_origin_key, repl_origin_key_range,
+    slru_block_to_key, slru_dir_to_key, slru_segment_key_range, slru_segment_size_to_key,
+    twophase_file_key, twophase_key_range, CompactKey, AUX_FILES_KEY, CHECKPOINT_KEY,
+    CONTROLFILE_KEY, DBDIR_KEY, TWOPHASEDIR_KEY,
 };
+use pageserver_api::key::{rel_tag_sparse_key, Key};
 use pageserver_api::keyspace::SparseKeySpace;
 use pageserver_api::record::NeonWalRecord;
 use pageserver_api::reltag::{BlockNumber, RelTag, SlruKind};
@@ -55,6 +56,8 @@ pub const MAX_AUX_FILE_DELTAS: usize = 1024;
 /// Max number of aux-file-related delta layers. The compaction will create a new image layer once this threshold is reached.
 pub const MAX_AUX_FILE_V2_DELTAS: usize = 16;
 
+pub const REL_STORE_V2: bool = true;
+
 #[derive(Debug)]
 pub enum LsnForTimestamp {
     /// Found commits both before and after the given timestamp
@@ -483,12 +486,20 @@ impl Timeline {
         if !dbdirs.contains_key(&(tag.spcnode, tag.dbnode)) {
             return Ok(false);
         }
-        // fetch directory listing
-        let key = rel_dir_to_key(tag.spcnode, tag.dbnode);
-        let buf = version.get(self, key, ctx).await?;
 
-        let dir = RelDirectory::des(&buf)?;
-        Ok(dir.rels.contains(&(tag.relnode, tag.forknum)))
+        if REL_STORE_V2 {
+            // fetch directory listing
+            let key = rel_tag_sparse_key(tag.spcnode, tag.dbnode, tag.relnode, tag.forknum);
+            let buf = version.get(self, key, ctx).await;
+            Ok(buf.is_ok()) // TODO: sparse keyspace needs a different get function
+        } else {
+            // fetch directory listing
+            let key = rel_dir_to_key(tag.spcnode, tag.dbnode);
+            let buf = version.get(self, key, ctx).await?;
+
+            let dir = RelDirectory::des(&buf)?;
+            Ok(dir.rels.contains(&(tag.relnode, tag.forknum)))
+        }
     }
 
     /// Get a list of all existing relations in given tablespace and database.
@@ -506,20 +517,41 @@ impl Timeline {
         version: Version<'_>,
         ctx: &RequestContext,
     ) -> Result<HashSet<RelTag>, PageReconstructError> {
-        // fetch directory listing
-        let key = rel_dir_to_key(spcnode, dbnode);
-        let buf = version.get(self, key, ctx).await?;
+        if REL_STORE_V2 {
+            // scan directory listing
+            let key_range = rel_tag_sparse_key_range(spcnode, dbnode);
+            let results = self
+                .get_vectored(KeySpace::single(key_range), version.get_lsn(), ctx)
+                .await?;
+            let mut rels = HashSet::new();
+            for (key, _) in results {
+                assert_eq!(key.field6, 1);
+                assert_eq!(key.field2, spcnode);
+                assert_eq!(key.field3, dbnode);
+                rels.insert(RelTag {
+                    spcnode,
+                    dbnode,
+                    relnode: key.field4,
+                    forknum: key.field5,
+                });
+            }
+            Ok(rels)
+        } else {
+            // fetch directory listing
+            let key = rel_dir_to_key(spcnode, dbnode);
+            let buf = version.get(self, key, ctx).await?;
 
-        let dir = RelDirectory::des(&buf)?;
-        let rels: HashSet<RelTag> =
-            HashSet::from_iter(dir.rels.iter().map(|(relnode, forknum)| RelTag {
-                spcnode,
-                dbnode,
-                relnode: *relnode,
-                forknum: *forknum,
-            }));
+            let dir = RelDirectory::des(&buf)?;
+            let rels: HashSet<RelTag> =
+                HashSet::from_iter(dir.rels.iter().map(|(relnode, forknum)| RelTag {
+                    spcnode,
+                    dbnode,
+                    relnode: *relnode,
+                    forknum: *forknum,
+                }));
 
-        Ok(rels)
+            Ok(rels)
+        }
     }
 
     /// Get the whole SLRU segment
@@ -1116,7 +1148,11 @@ impl Timeline {
 
         let dense_keyspace = result.to_keyspace();
         let sparse_keyspace = SparseKeySpace(KeySpace {
-            ranges: vec![Key::metadata_aux_key_range(), repl_origin_key_range()],
+            ranges: vec![
+                Key::metadata_aux_key_range(),
+                repl_origin_key_range(),
+                Key::rel_dir_sparse_key_range(),
+            ],
         });
 
         if cfg!(debug_assertions) {
@@ -1608,7 +1644,7 @@ impl DatadirModification<'_> {
             let buf = DbDirectory::ser(&dbdir)?;
             self.put(DBDIR_KEY, Value::Image(buf.into()));
         }
-        if r.is_none() {
+        if REL_STORE_V2 && r.is_none() {
             // Create RelDirectory
             let buf = RelDirectory::ser(&RelDirectory {
                 rels: HashSet::new(),
@@ -1730,8 +1766,8 @@ impl DatadirModification<'_> {
         // tablespace.  Create the reldir entry for it if so.
         let mut dbdir = DbDirectory::des(&self.get(DBDIR_KEY, ctx).await.context("read db")?)
             .context("deserialize db")?;
-        let rel_dir_key = rel_dir_to_key(rel.spcnode, rel.dbnode);
-        let mut rel_dir =
+
+        if REL_STORE_V2 {
             if let hash_map::Entry::Vacant(e) = dbdir.dbdirs.entry((rel.spcnode, rel.dbnode)) {
                 // Didn't exist. Update dbdir
                 e.insert(false);
@@ -1739,30 +1775,49 @@ impl DatadirModification<'_> {
                 self.pending_directory_entries
                     .push((DirectoryKind::Db, dbdir.dbdirs.len()));
                 self.put(DBDIR_KEY, Value::Image(buf.into()));
+            }
+            let rel_dir_key = rel_tag_sparse_key(rel.spcnode, rel.dbnode, rel.relnode, rel.forknum);
+            // check if the rel_dir_key exists
+            if self.get(rel_dir_key, ctx).await.is_ok() {
+                return Err(RelationError::AlreadyExists);
+            }
+            self.put(rel_dir_key, Value::Image(Bytes::from_static(b"1")));
+            // TODO: update directory_entries_count, it seems to be a metrics
+        } else {
+            let rel_dir_key = rel_dir_to_key(rel.spcnode, rel.dbnode);
+            let mut rel_dir =
+                if let hash_map::Entry::Vacant(e) = dbdir.dbdirs.entry((rel.spcnode, rel.dbnode)) {
+                    // Didn't exist. Update dbdir
+                    e.insert(false);
+                    let buf = DbDirectory::ser(&dbdir).context("serialize db")?;
+                    self.pending_directory_entries
+                        .push((DirectoryKind::Db, dbdir.dbdirs.len()));
+                    self.put(DBDIR_KEY, Value::Image(buf.into()));
 
-                // and create the RelDirectory
-                RelDirectory::default()
-            } else {
-                // reldir already exists, fetch it
-                RelDirectory::des(&self.get(rel_dir_key, ctx).await.context("read db")?)
-                    .context("deserialize db")?
-            };
+                    // and create the RelDirectory
+                    RelDirectory::default()
+                } else {
+                    // reldir already exists, fetch it
+                    RelDirectory::des(&self.get(rel_dir_key, ctx).await.context("read db")?)
+                        .context("deserialize db")?
+                };
 
-        // Add the new relation to the rel directory entry, and write it back
-        if !rel_dir.rels.insert((rel.relnode, rel.forknum)) {
-            return Err(RelationError::AlreadyExists);
+            // Add the new relation to the rel directory entry, and write it back
+            if !rel_dir.rels.insert((rel.relnode, rel.forknum)) {
+                return Err(RelationError::AlreadyExists);
+            }
+
+            self.pending_directory_entries
+                .push((DirectoryKind::Rel, rel_dir.rels.len()));
+
+            self.put(
+                rel_dir_key,
+                Value::Image(Bytes::from(
+                    RelDirectory::ser(&rel_dir).context("serialize")?,
+                )),
+            );
         }
 
-        self.pending_directory_entries
-            .push((DirectoryKind::Rel, rel_dir.rels.len()));
-
-        self.put(
-            rel_dir_key,
-            Value::Image(Bytes::from(
-                RelDirectory::ser(&rel_dir).context("serialize")?,
-            )),
-        );
-
         // Put size
         let size_key = rel_size_to_key(rel);
         let buf = nblocks.to_le_bytes();
@@ -1841,33 +1896,57 @@ impl DatadirModification<'_> {
         drop_relations: HashMap<(u32, u32), Vec<RelTag>>,
         ctx: &RequestContext,
     ) -> anyhow::Result<()> {
-        for ((spc_node, db_node), rel_tags) in drop_relations {
-            let dir_key = rel_dir_to_key(spc_node, db_node);
-            let buf = self.get(dir_key, ctx).await?;
-            let mut dir = RelDirectory::des(&buf)?;
+        if REL_STORE_V2 {
+            for ((spc_node, db_node), rel_tags) in drop_relations {
+                for rel_tag in rel_tags {
+                    let key =
+                        rel_tag_sparse_key(spc_node, db_node, rel_tag.relnode, rel_tag.forknum);
+                    if self.get(key, ctx).await.is_ok() {
+                        // remove the relation key
+                        self.delete(key..key.next());
 
-            let mut dirty = false;
-            for rel_tag in rel_tags {
-                if dir.rels.remove(&(rel_tag.relnode, rel_tag.forknum)) {
-                    dirty = true;
+                        // update logical size
+                        let size_key = rel_size_to_key(rel_tag);
+                        let old_size = self.get(size_key, ctx).await?.get_u32_le();
+                        self.pending_nblocks -= old_size as i64;
 
-                    // update logical size
-                    let size_key = rel_size_to_key(rel_tag);
-                    let old_size = self.get(size_key, ctx).await?.get_u32_le();
-                    self.pending_nblocks -= old_size as i64;
+                        // Remove entry from relation size cache
+                        self.tline.remove_cached_rel_size(&rel_tag);
 
-                    // Remove entry from relation size cache
-                    self.tline.remove_cached_rel_size(&rel_tag);
-
-                    // Delete size entry, as well as all blocks
-                    self.delete(rel_key_range(rel_tag));
+                        // Delete size entry, as well as all blocks
+                        self.delete(rel_key_range(rel_tag));
+                    }
                 }
             }
+        } else {
+            for ((spc_node, db_node), rel_tags) in drop_relations {
+                let dir_key = rel_dir_to_key(spc_node, db_node);
+                let buf = self.get(dir_key, ctx).await?;
+                let mut dir = RelDirectory::des(&buf)?;
 
-            if dirty {
-                self.put(dir_key, Value::Image(Bytes::from(RelDirectory::ser(&dir)?)));
-                self.pending_directory_entries
-                    .push((DirectoryKind::Rel, dir.rels.len()));
+                let mut dirty = false;
+                for rel_tag in rel_tags {
+                    if dir.rels.remove(&(rel_tag.relnode, rel_tag.forknum)) {
+                        dirty = true;
+
+                        // update logical size
+                        let size_key = rel_size_to_key(rel_tag);
+                        let old_size = self.get(size_key, ctx).await?.get_u32_le();
+                        self.pending_nblocks -= old_size as i64;
+
+                        // Remove entry from relation size cache
+                        self.tline.remove_cached_rel_size(&rel_tag);
+
+                        // Delete size entry, as well as all blocks
+                        self.delete(rel_key_range(rel_tag));
+                    }
+                }
+
+                if dirty {
+                    self.put(dir_key, Value::Image(Bytes::from(RelDirectory::ser(&dir)?)));
+                    self.pending_directory_entries
+                        .push((DirectoryKind::Rel, dir.rels.len()));
+                }
             }
         }