perf(storage): simplify table watermark index

src/frontend/src/catalog/system_catalog/rw_catalog/rw_hummock_version.rs

@@ -169,14 +169,14 @@ async fn read_hummock_table_watermarks(
             for (vnode, epoch, watermark) in
                 table_watermarks
                     .watermarks
-                    .into_iter()
+                    .iter()
                     .flat_map(move |(epoch, watermarks)| {
-                        watermarks.into_iter().flat_map(move |vnode_watermark| {
+                        watermarks.iter().flat_map(move |vnode_watermark| {
                             let watermark = vnode_watermark.watermark().clone();
                             let vnodes = vnode_watermark.vnode_bitmap().iter_ones().collect_vec();
                             vnodes
                                 .into_iter()
-                                .map(move |vnode| (vnode, epoch, Vec::from(watermark.as_ref())))
+                                .map(move |vnode| (vnode, *epoch, Vec::from(watermark.as_ref())))
                         })
                     })
             {

src/storage/Cargo.toml

@@ -145,5 +145,9 @@ harness = false
 name = "bench_imm_compact"
 harness = false
 
+[[bench]]
+name = "bench_table_watermarks"
+harness = false
+
 [lints]
 workspace = true

src/storage/benches/bench_table_watermarks.rs

@@ -0,0 +1,259 @@
+// Copyright 2024 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#![feature(lazy_cell)]
+
+use std::collections::hash_map::Entry;
+use std::collections::{HashMap, VecDeque};
+use std::sync::{Arc, LazyLock};
+
+use bytes::Bytes;
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use itertools::Itertools;
+use risingwave_common::buffer::{Bitmap, BitmapBuilder};
+use risingwave_common::catalog::TableId;
+use risingwave_common::hash::VirtualNode;
+use risingwave_common::util::epoch::test_epoch;
+use risingwave_hummock_sdk::table_watermark::{
+    TableWatermarks, TableWatermarksIndex, VnodeWatermark, WatermarkDirection,
+};
+use risingwave_hummock_sdk::version::HummockVersion;
+use risingwave_hummock_sdk::HummockEpoch;
+use risingwave_storage::hummock::local_version::pinned_version::PinnedVersion;
+use spin::Mutex;
+use tokio::sync::mpsc::unbounded_channel;
+
+fn vnode_bitmaps(part_count: usize) -> impl Iterator<Item = Arc<Bitmap>> {
+    static BITMAP_CACHE: LazyLock<Mutex<HashMap<usize, Vec<Arc<Bitmap>>>>> =
+        LazyLock::new(|| Mutex::new(HashMap::new()));
+    assert_eq!(VirtualNode::COUNT % part_count, 0);
+    let mut cache = BITMAP_CACHE.lock();
+    match cache.entry(part_count) {
+        Entry::Occupied(entry) => entry.get().clone().into_iter(),
+        Entry::Vacant(entry) => entry
+            .insert({
+                let part_size = VirtualNode::COUNT / part_count;
+                (0..part_count)
+                    .map(move |part_idx| {
+                        let start = part_idx * part_size;
+                        let end = part_idx * part_size + part_size;
+                        let mut bitmap = BitmapBuilder::zeroed(VirtualNode::COUNT);
+                        for i in start..end {
+                            bitmap.set(i, true);
+                        }
+                        Arc::new(bitmap.finish())
+                    })
+                    .collect()
+            })
+            .clone()
+            .into_iter(),
+    }
+}
+
+fn gen_watermark(epoch_idx: usize) -> Bytes {
+    static WATERMARK_CACHE: LazyLock<Mutex<HashMap<usize, Bytes>>> =
+        LazyLock::new(|| Mutex::new(HashMap::new()));
+    let mut cache = WATERMARK_CACHE.lock();
+    match cache.entry(epoch_idx) {
+        Entry::Occupied(entry) => entry.get().clone(),
+        Entry::Vacant(entry) => entry
+            .insert(Bytes::copy_from_slice(
+                format!("key_test_{:010}", epoch_idx).as_bytes(),
+            ))
+            .clone(),
+    }
+}
+
+fn gen_epoch_watermarks(
+    epoch_idx: usize,
+    vnode_part_count: usize,
+) -> (HummockEpoch, Vec<VnodeWatermark>) {
+    (
+        test_epoch(epoch_idx as _),
+        vnode_bitmaps(vnode_part_count)
+            .map(|bitmap| VnodeWatermark::new(bitmap, gen_watermark(epoch_idx)))
+            .collect_vec(),
+    )
+}
+
+fn gen_committed_table_watermarks(
+    old_epoch_idx: usize,
+    new_epoch_idx: usize,
+    vnode_part_count: usize,
+) -> TableWatermarks {
+    assert!(old_epoch_idx <= new_epoch_idx);
+    TableWatermarks {
+        watermarks: (old_epoch_idx..=new_epoch_idx)
+            .map(|epoch_idx| {
+                let (epoch, watermarks) = gen_epoch_watermarks(epoch_idx, vnode_part_count);
+                (epoch, watermarks.into())
+            })
+            .collect(),
+        direction: WatermarkDirection::Ascending,
+    }
+}
+
+fn gen_version(
+    old_epoch_idx: usize,
+    new_epoch_idx: usize,
+    table_count: usize,
+    vnode_part_count: usize,
+) -> HummockVersion {
+    let table_watermarks = Arc::new(gen_committed_table_watermarks(
+        old_epoch_idx,
+        new_epoch_idx,
+        vnode_part_count,
+    ));
+    // let table_watermarks =
+    //     gen_committed_table_watermarks(old_epoch_idx, new_epoch_idx, vnode_part_count);
+    HummockVersion {
+        id: new_epoch_idx as _,
+        max_committed_epoch: test_epoch(new_epoch_idx as _),
+        safe_epoch: test_epoch(old_epoch_idx as _),
+        table_watermarks: (0..table_count)
+            .map(|table_id| (TableId::new(table_id as _), table_watermarks.clone()))
+            .collect(),
+        ..Default::default()
+    }
+}
+
+fn bench_table_watermarks(c: &mut Criterion) {
+    let version_count = 500;
+    let epoch_count = 1000;
+    let table_count = 1;
+    let vnode_part_count = 16;
+    let pre_generated_versions: VecDeque<_> = (1..version_count + 1)
+        .map(|epoch_idx| {
+            gen_version(
+                epoch_idx,
+                epoch_idx + epoch_count,
+                table_count,
+                vnode_part_count,
+            )
+        })
+        .collect();
+    c.bench_function("new pinned version", |b| {
+        b.iter_batched(
+            || pre_generated_versions.clone(),
+            |mut versions| {
+                let mut pinned_version =
+                    PinnedVersion::new(versions.pop_front().unwrap(), unbounded_channel().0);
+                while let Some(version) = versions.pop_front() {
+                    pinned_version = pinned_version.new_pin_version(version);
+                }
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    let safe_epoch_idx: usize = 9500;
+    let committed_epoch_idx: usize = 10000;
+    let staging_epoch_count: usize = 500;
+    let vnode_part_count = 16;
+
+    let mut table_watermarks = TableWatermarksIndex::new_committed(
+        gen_committed_table_watermarks(safe_epoch_idx, committed_epoch_idx, vnode_part_count)
+            .into(),
+        test_epoch(committed_epoch_idx as u64),
+    );
+    for i in 0..staging_epoch_count {
+        let (epoch, watermarks) =
+            gen_epoch_watermarks(committed_epoch_idx + i + 1, vnode_part_count);
+        table_watermarks.add_epoch_watermark(
+            epoch,
+            watermarks.into(),
+            WatermarkDirection::Ascending,
+        );
+    }
+    let table_watermarks = table_watermarks;
+    let committed_watermarks = gen_committed_table_watermarks(
+        safe_epoch_idx + 1,
+        committed_epoch_idx + 1,
+        vnode_part_count,
+    );
+    let batch_size = 100;
+    c.bench_function("apply committed watermark", |b| {
+        b.iter_batched(
+            || {
+                (0..batch_size)
+                    .map(|_| (table_watermarks.clone(), committed_watermarks.clone()))
+                    .collect_vec()
+            },
+            |list| {
+                for (mut table_watermarks, committed_watermarks) in list {
+                    table_watermarks.apply_committed_watermarks(
+                        committed_watermarks.into(),
+                        test_epoch((committed_epoch_idx + 1) as _),
+                    );
+                }
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    // Code for the original table watermark index
+    // let mut table_watermarks =
+    //     gen_committed_table_watermarks(safe_epoch_idx, committed_epoch_idx, vnode_part_count)
+    //         .build_index(test_epoch(committed_epoch_idx as u64));
+    // for i in 0..staging_epoch_count {
+    //     let (epoch, watermarks) =
+    //         gen_epoch_watermarks(committed_epoch_idx + i + 1, vnode_part_count);
+    //     table_watermarks.add_epoch_watermark(epoch, &watermarks, WatermarkDirection::Ascending);
+    // }
+    // let table_watermarks = table_watermarks;
+    // let committed_watermarks = gen_committed_table_watermarks(
+    //     safe_epoch_idx + 1,
+    //     committed_epoch_idx + 1,
+    //     vnode_part_count,
+    // )
+    // .build_index(test_epoch((committed_epoch_idx + 1) as _));
+    // let batch_size = 100;
+    // c.bench_function("apply committed watermark", |b| {
+    //     b.iter_batched(
+    //         || {
+    //             (0..batch_size)
+    //                 .map(|_| (table_watermarks.clone(), committed_watermarks.clone()))
+    //                 .collect_vec()
+    //         },
+    //         |list| {
+    //             for (mut table_watermarks, committed_watermarks) in list {
+    //                 table_watermarks.apply_committed_watermarks(&committed_watermarks);
+    //             }
+    //         },
+    //         BatchSize::SmallInput,
+    //     )
+    // });
+
+    c.bench_function("read latest watermark", |b| {
+        b.iter(|| {
+            for i in 0..VirtualNode::COUNT {
+                let _ = table_watermarks.latest_watermark(VirtualNode::from_index(i));
+            }
+        })
+    });
+
+    c.bench_function("read committed watermark", |b| {
+        b.iter(|| {
+            for i in 0..VirtualNode::COUNT {
+                let _ = table_watermarks.read_watermark(
+                    VirtualNode::from_index(i),
+                    test_epoch(committed_epoch_idx as u64),
+                );
+            }
+        })
+    });
+}
+
+criterion_group!(benches, bench_table_watermarks);
+criterion_main!(benches);

src/storage/hummock_sdk/src/compaction_group/hummock_version_ext.rs

@@ -15,6 +15,7 @@
 use std::cmp::Ordering;
 use std::collections::hash_map::Entry;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
+use std::sync::Arc;
 
 use itertools::Itertools;
 use risingwave_common::catalog::TableId;
@@ -33,7 +34,7 @@ use super::StateTableId;
 use crate::compaction_group::StaticCompactionGroupId;
 use crate::key_range::KeyRangeCommon;
 use crate::prost_key_range::KeyRangeExt;
-use crate::table_watermark::{TableWatermarks, TableWatermarksIndex, VnodeWatermark};
+use crate::table_watermark::{TableWatermarks, VnodeWatermark};
 use crate::version::{HummockVersion, HummockVersionDelta};
 use crate::{can_concat, CompactionGroupId, HummockSstableId, HummockSstableObjectId};
 
@@ -195,18 +196,6 @@ impl HummockVersion {
             .unwrap_or(0)
     }
 
-    pub fn build_table_watermarks_index(&self) -> HashMap<TableId, TableWatermarksIndex> {
-        self.table_watermarks
-            .iter()
-            .map(|(table_id, table_watermarks)| {
-                (
-                    *table_id,
-                    table_watermarks.build_index(self.max_committed_epoch),
-                )
-            })
-            .collect()
-    }
-
     pub fn safe_epoch_table_watermarks(
         &self,
         existing_table_ids: &[u32],
@@ -600,25 +589,34 @@ impl HummockVersion {
         for table_id in &version_delta.removed_table_ids {
             let _ = self.table_watermarks.remove(table_id);
         }
+
+        let mut modified_table_watermarks: HashMap<TableId, TableWatermarks> = HashMap::new();
+
         for (table_id, table_watermarks) in &version_delta.new_table_watermarks {
-            match self.table_watermarks.entry(*table_id) {
-                Entry::Occupied(mut entry) => {
-                    entry.get_mut().apply_new_table_watermarks(table_watermarks);
-                }
-                Entry::Vacant(entry) => {
-                    entry.insert(table_watermarks.clone());
-                }
+            if let Some(current_table_watermarks) = self.table_watermarks.get(table_id) {
+                let mut current_table_watermarks = (**current_table_watermarks).clone();
+                current_table_watermarks.apply_new_table_watermarks(table_watermarks);
+                modified_table_watermarks.insert(*table_id, current_table_watermarks);
+            } else {
+                modified_table_watermarks.insert(*table_id, table_watermarks.clone());
             }
         }
         if version_delta.safe_epoch != self.safe_epoch {
             assert!(version_delta.safe_epoch > self.safe_epoch);
-            self.table_watermarks
-                .values_mut()
-                .for_each(|table_watermarks| {
-                    table_watermarks.clear_stale_epoch_watermark(version_delta.safe_epoch)
-                });
+            for (table_id, table_watermarks) in &self.table_watermarks {
+                let table_watermarks = modified_table_watermarks
+                    .entry(*table_id)
+                    .or_insert_with(|| (**table_watermarks).clone());
+                table_watermarks.clear_stale_epoch_watermark(version_delta.safe_epoch);
+            }
             self.safe_epoch = version_delta.safe_epoch;
         }
+
+        for (table_id, table_watermarks) in modified_table_watermarks {
+            self.table_watermarks
+                .insert(table_id, Arc::new(table_watermarks));
+        }
+
         sst_split_info
     }