feat(storage): introduce non-proto type for vnode bitmap

src/common/src/consistent_hash/mod.rs

@@ -0,0 +1,16 @@
+// Copyright 2022 Singularity Data
+//
+// 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.
+
+mod vnode;
+pub use vnode::*;

src/common/src/consistent_hash/vnode.rs

@@ -0,0 +1,131 @@
+// Copyright 2022 Singularity Data
+//
+// 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.
+
+/// `VirtualNode` is the logical key for consistent hash. Virtual nodes stand for the intermediate
+/// layer between data and physical nodes.
+pub type VirtualNode = u16;
+pub const VNODE_BITS: usize = 11;
+pub const VIRTUAL_NODE_COUNT: usize = 1 << VNODE_BITS;
+pub const VNODE_BITMAP_LEN: usize = 1 << (VNODE_BITS - 3);
+
+/// `VNodeBitmap` is a bitmap of vnodes for a state table, which indicates the vnodes that the state
+/// table owns.
+#[derive(Clone, Default)]
+pub struct VNodeBitmap {
+    /// Id of state table that the bitmap belongs to.
+    table_id: u32,
+    /// Bitmap that indicates vnodes that the state tables owns.
+    bitmap: Vec<u8>,
+}
+
+impl From<risingwave_pb::common::VNodeBitmap> for VNodeBitmap {
+    fn from(proto: risingwave_pb::common::VNodeBitmap) -> Self {
+        Self {
+            table_id: proto.table_id,
+            bitmap: proto.bitmap,
+        }
+    }
+}
+
+impl VNodeBitmap {
+    pub fn new(table_id: u32, bitmap: Vec<u8>) -> Self {
+        Self { table_id, bitmap }
+    }
+
+    /// Checks whether the bitmap overlaps with the given bitmaps. Two bitmaps overlap only when
+    /// they belong to the same state table and have at least one common "1" bit. The given
+    /// bitmaps are always from `SstableInfo`, which are in the form of protobuf.
+    pub fn check_overlap(&self, bitmaps: &[risingwave_pb::common::VNodeBitmap]) -> bool {
+        if bitmaps.is_empty() {
+            return true;
+        }
+        if let Ok(pos) =
+            bitmaps.binary_search_by_key(&self.table_id, |bitmap| bitmap.get_table_id())
+        {
+            let candidate_bitmap = &bitmaps[pos];
+            assert_eq!(self.bitmap.len(), VNODE_BITMAP_LEN);
+            assert_eq!(candidate_bitmap.bitmap.len(), VNODE_BITMAP_LEN);
+            for i in 0..VNODE_BITMAP_LEN as usize {
+                if (self.bitmap[i] & candidate_bitmap.get_bitmap()[i]) != 0 {
+                    return true;
+                }
+            }
+        }
+        false
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use itertools::Itertools;
+    use risingwave_pb::common::VNodeBitmap as ProstBitmap;
+
+    use super::*;
+
+    #[test]
+    fn test_check_overlap() {
+        // The bitmap is for table with id 3, and owns vnodes from 64 to 128.
+        let table_id = 3;
+        let mut bitmap = [0; VNODE_BITMAP_LEN].to_vec();
+        for byte in bitmap.iter_mut().take(16).skip(8) {
+            *byte = 0b11111111;
+        }
+        let vnode_bitmap = VNodeBitmap::new(table_id, bitmap);
+
+        // Test overlap.
+        let bitmaps_1 = (2..4)
+            .map(|table_id| {
+                let mut test_bitmap = [0; VNODE_BITMAP_LEN].to_vec();
+                test_bitmap[10] = 0b1;
+                ProstBitmap {
+                    table_id,
+                    bitmap: test_bitmap,
+                }
+            })
+            .collect_vec();
+        assert!(vnode_bitmap.check_overlap(&bitmaps_1));
+
+        // Test non-overlap with same table id.
+        let bitmaps_2 = (2..4)
+            .map(|table_id| {
+                let mut test_bitmap = [0; VNODE_BITMAP_LEN].to_vec();
+                test_bitmap[20] = 0b1;
+                ProstBitmap {
+                    table_id,
+                    bitmap: test_bitmap,
+                }
+            })
+            .collect_vec();
+        assert!(!vnode_bitmap.check_overlap(&bitmaps_2));
+
+        // Test non-overlap with different table ids and same vnodes.
+        let bitmaps_3 = (4..6)
+            .map(|table_id| {
+                let mut test_bitmap = [0; VNODE_BITMAP_LEN].to_vec();
+                for byte in test_bitmap.iter_mut().take(16).skip(8) {
+                    *byte = 0b11111111;
+                }
+                ProstBitmap {
+                    table_id,
+                    bitmap: test_bitmap,
+                }
+            })
+            .collect_vec();
+        assert!(!vnode_bitmap.check_overlap(&bitmaps_3));
+
+        // Test empty
+        assert!(vnode_bitmap.check_overlap(&[]));
+    }
+}

src/common/src/lib.rs

@@ -48,6 +48,7 @@ pub mod cache;
 pub mod catalog;
 pub mod collection;
 pub mod config;
+pub mod consistent_hash;
 pub mod hash;
 pub mod monitor;
 pub mod service;

src/storage/src/hummock/shared_buffer/mod.rs

@@ -23,9 +23,9 @@ use std::ops::{Bound, RangeBounds};
 use std::sync::Arc;
 
 use itertools::Itertools;
+use risingwave_common::consistent_hash::VNodeBitmap;
 use risingwave_hummock_sdk::is_remote_sst_id;
 use risingwave_hummock_sdk::key::user_key;
-use risingwave_pb::common::VNodeBitmap;
 use risingwave_pb::hummock::{KeyRange, SstableInfo};
 
 use self::shared_buffer_batch::SharedBufferBatch;

src/storage/src/hummock/state_store.rs

@@ -19,9 +19,9 @@ use std::sync::Arc;
 
 use bytes::Bytes;
 use itertools::Itertools;
+use risingwave_common::consistent_hash::VNodeBitmap;
 use risingwave_hummock_sdk::key::key_with_epoch;
 use risingwave_hummock_sdk::HummockEpoch;
-use risingwave_pb::common::VNodeBitmap;
 use risingwave_pb::hummock::SstableInfo;
 
 use super::iterator::{

src/storage/src/hummock/state_store_tests.rs

@@ -16,11 +16,11 @@ use std::sync::Arc;
 
 use bytes::{BufMut, Bytes, BytesMut};
 use futures::executor::block_on;
+use risingwave_common::consistent_hash::VNodeBitmap;
 use risingwave_common::hash::VNODE_BITMAP_LEN;
 use risingwave_hummock_sdk::HummockEpoch;
 use risingwave_meta::hummock::test_utils::setup_compute_env;
 use risingwave_meta::hummock::MockHummockMetaClient;
-use risingwave_pb::common::VNodeBitmap;
 use risingwave_rpc_client::HummockMetaClient;
 
 use super::HummockStorage;
@@ -222,10 +222,7 @@ async fn test_vnode_filter() {
         .get_with_vnode_set(
             &Bytes::from(&b"t\0\0\0\0"[..]),
             epoch,
-            Some(VNodeBitmap {
-                table_id: 0,
-                bitmap: [1; VNODE_BITMAP_LEN].to_vec(),
-            }),
+            Some(VNodeBitmap::new(0, [1; VNODE_BITMAP_LEN].to_vec())),
         )
         .await
         .unwrap();
@@ -235,10 +232,7 @@ async fn test_vnode_filter() {
         .get_with_vnode_set(
             &Bytes::from(&b"t\0\0\0\0"[..]),
             epoch,
-            Some(VNodeBitmap {
-                table_id: 0,
-                bitmap: [0; VNODE_BITMAP_LEN].to_vec(),
-            }),
+            Some(VNodeBitmap::new(0, [0; VNODE_BITMAP_LEN].to_vec())),
         )
         .await
         .unwrap();
@@ -248,10 +242,7 @@ async fn test_vnode_filter() {
         .get_with_vnode_set(
             &Bytes::from(&b"t\0\0\0\0"[..]),
             epoch,
-            Some(VNodeBitmap {
-                table_id: 5,
-                bitmap: [1; VNODE_BITMAP_LEN].to_vec(),
-            }),
+            Some(VNodeBitmap::new(5, [1; VNODE_BITMAP_LEN].to_vec())),
         )
         .await
         .unwrap();

src/storage/src/hummock/utils.rs

@@ -16,9 +16,8 @@ use std::cmp::Ordering;
 use std::ops::Bound::{Excluded, Included, Unbounded};
 use std::ops::RangeBounds;
 
-use risingwave_common::hash::VNODE_BITMAP_LEN;
+use risingwave_common::consistent_hash::VNodeBitmap;
 use risingwave_hummock_sdk::key::user_key;
-use risingwave_pb::common::VNodeBitmap;
 use risingwave_pb::hummock::{Level, SstableInfo};
 
 use super::{HummockError, HummockResult};
@@ -74,25 +73,6 @@ pub fn validate_table_key_range(levels: &[Level]) -> HummockResult<()> {
     Ok(())
 }
 
-pub fn bitmap_overlap(pattern: &VNodeBitmap, sst_bitmaps: &Vec<VNodeBitmap>) -> bool {
-    if sst_bitmaps.is_empty() {
-        return true;
-    }
-    if let Ok(pos) =
-        sst_bitmaps.binary_search_by_key(&pattern.get_table_id(), |bitmap| bitmap.get_table_id())
-    {
-        let text = &sst_bitmaps[pos];
-        assert_eq!(pattern.bitmap.len(), VNODE_BITMAP_LEN);
-        assert_eq!(text.bitmap.len(), VNODE_BITMAP_LEN);
-        for i in 0..VNODE_BITMAP_LEN as usize {
-            if (pattern.get_bitmap()[i] & text.get_bitmap()[i]) != 0 {
-                return true;
-            }
-        }
-    }
-    false
-}
-
 pub fn filter_single_sst<R, B>(
     info: &SstableInfo,
     key_range: &R,
@@ -108,7 +88,7 @@ where
         let table_end = user_key(table_range.right.as_slice());
         range_overlap(key_range, table_start, table_end)
     }) && vnode_set.map_or(true, |vnode_set| {
-        bitmap_overlap(vnode_set, &info.vnode_bitmaps)
+        vnode_set.check_overlap(&info.vnode_bitmaps)
     })
 }