risingwave_common/hash/consistent_hash/
vnode_count.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::num::NonZeroUsize;
16
17use super::vnode::VirtualNode;
18
19/// The different cases of `maybe_vnode_count` field in the protobuf message.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
21pub enum VnodeCount {
22    /// The field is a placeholder and has to be filled first before using it.
23    #[default]
24    Placeholder,
25    /// The table/fragment is a singleton, thus the value should always be interpreted as `1`.
26    Singleton,
27    /// The field is set to a specific value.
28    Set(NonZeroUsize),
29    /// The field is unset because the table/fragment is persisted as hash-distributed
30    /// in an older version.
31    CompatHash,
32}
33
34impl VnodeCount {
35    /// Creates a `VnodeCount` set to the given value.
36    pub fn set(v: impl TryInto<usize> + Copy + std::fmt::Debug) -> Self {
37        let v = (v.try_into().ok())
38            .filter(|v| (1..=VirtualNode::MAX_COUNT).contains(v))
39            .unwrap_or_else(|| panic!("invalid vnode count {v:?}"));
40
41        VnodeCount::Set(NonZeroUsize::new(v).unwrap())
42    }
43
44    /// Creates a `VnodeCount` set to the value for testing.
45    ///
46    /// Equivalent to `VnodeCount::set(VirtualNode::COUNT_FOR_TEST)`.
47    pub fn for_test() -> Self {
48        Self::set(VirtualNode::COUNT_FOR_TEST)
49    }
50
51    /// Converts from protobuf representation of `maybe_vnode_count`.
52    ///
53    /// The value will be ignored if `is_singleton` returns `true`.
54    pub fn from_protobuf(v: Option<u32>, is_singleton: impl FnOnce() -> bool) -> Self {
55        match v {
56            Some(0) => VnodeCount::Placeholder,
57            _ => {
58                if is_singleton() {
59                    if let Some(v) = v
60                        && v != 1
61                    {
62                        tracing::debug!(
63                            vnode_count = v,
64                            "singleton has vnode count set to non-1, \
65                            ignoring as it could be due to backward compatibility"
66                        );
67                    }
68                    VnodeCount::Singleton
69                } else {
70                    v.map_or(VnodeCount::CompatHash, VnodeCount::set)
71                }
72            }
73        }
74    }
75
76    /// Converts to protobuf representation for `maybe_vnode_count`.
77    pub fn to_protobuf(self) -> Option<u32> {
78        // Effectively fills the compatibility cases with values.
79        self.value_opt()
80            .map_or(Some(0) /* placeholder */, |v| Some(v as _))
81    }
82
83    /// Returns the value of the vnode count, or `None` if it's a placeholder.
84    pub fn value_opt(self) -> Option<usize> {
85        match self {
86            VnodeCount::Placeholder => None,
87            VnodeCount::Singleton => Some(1),
88            VnodeCount::Set(v) => Some(v.get()),
89            VnodeCount::CompatHash => Some(VirtualNode::COUNT_FOR_COMPAT),
90        }
91    }
92
93    /// Returns the value of the vnode count. Panics if it's a placeholder.
94    pub fn value(self) -> usize {
95        self.value_opt()
96            .expect("vnode count is a placeholder that must be filled by the meta service first")
97    }
98}
99
100/// A trait for checking whether a table/fragment is a singleton.
101pub trait IsSingleton {
102    /// Returns `true` if the table/fragment is a singleton.
103    ///
104    /// By singleton, we mean that all data read from or written to the storage belongs to
105    /// the only `SINGLETON_VNODE`. This must be consistent with the behavior of
106    /// [`TableDistribution`](crate::hash::table_distribution::TableDistribution::new).
107    /// As a result, the `vnode_count` of such table/fragment can be `1`.
108    fn is_singleton(&self) -> bool;
109}
110
111/// A trait for accessing the vnode count field with backward compatibility.
112///
113/// # `maybe_`?
114///
115/// The reason why there's a `maybe_` prefix on the protobuf field is that, a getter
116/// method with the same name as the field will be generated for `prost` structs.
117/// Directly naming it `vnode_count` will lead to the method `vnode_count()` returning
118/// `0` when the field is unset, which can be misleading sometimes.
119///
120/// Instead, we name the field as `maybe_vnode_count` and provide the method `vnode_count`
121/// through this trait, ensuring that backward compatibility is handled properly.
122pub trait VnodeCountCompat {
123    /// Get the `maybe_vnode_count` field.
124    fn vnode_count_inner(&self) -> VnodeCount;
125
126    /// Returns the vnode count if it's set. Otherwise, returns [`VirtualNode::COUNT_FOR_COMPAT`]
127    /// for distributed tables/fragments, and `1` for singleton tables/fragments, for backward
128    /// compatibility. Panics if the field is a placeholder.
129    ///
130    /// See the documentation on the field of the implementing type for more details.
131    fn vnode_count(&self) -> usize {
132        self.vnode_count_inner().value()
133    }
134}
135
136impl IsSingleton for risingwave_pb::catalog::Table {
137    fn is_singleton(&self) -> bool {
138        self.distribution_key.is_empty()
139            && self.dist_key_in_pk.is_empty()
140            && self.vnode_col_index.is_none()
141    }
142}
143impl VnodeCountCompat for risingwave_pb::catalog::Table {
144    fn vnode_count_inner(&self) -> VnodeCount {
145        VnodeCount::from_protobuf(self.maybe_vnode_count, || self.is_singleton())
146    }
147}
148
149impl IsSingleton for risingwave_pb::plan_common::StorageTableDesc {
150    fn is_singleton(&self) -> bool {
151        self.dist_key_in_pk_indices.is_empty() && self.vnode_col_idx_in_pk.is_none()
152    }
153}
154impl VnodeCountCompat for risingwave_pb::plan_common::StorageTableDesc {
155    fn vnode_count_inner(&self) -> VnodeCount {
156        VnodeCount::from_protobuf(self.maybe_vnode_count, || self.is_singleton())
157    }
158}
risingwave_common/hash/consistent_hash/vnode_count.rs

risingwave_common/hash/consistent_hash/
vnode_count.rs
