mz_cluster_client/
instances.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Types related to storage instances.
11
12use std::fmt;
13use std::str::FromStr;
14
15use anyhow::bail;
16#[cfg(any(test, feature = "proptest"))]
17use proptest::arbitrary::Arbitrary;
18#[cfg(any(test, feature = "proptest"))]
19use proptest::strategy::{BoxedStrategy, Strategy};
20use serde::{Deserialize, Serialize};
21use tracing::error;
22
23/// Identifier of a storage instance.
24#[derive(
25    Clone,
26    Copy,
27    Debug,
28    Eq,
29    PartialEq,
30    Ord,
31    PartialOrd,
32    Hash,
33    Serialize,
34    Deserialize
35)]
36pub enum StorageInstanceId {
37    /// A system storage instance.
38    System(u64),
39    /// A user storage instance.
40    User(u64),
41}
42
43impl StorageInstanceId {
44    /// Creates a new `StorageInstanceId` in the system namespace. The top 16 bits of `id` must be
45    /// 0, because this ID is packed into 48 bits of
46    /// `mz_repr::GlobalId::IntrospectionSourceIndex`.
47    pub fn system(id: u64) -> Option<Self> {
48        Self::new(id, Self::System)
49    }
50
51    /// Creates a new `StorageInstanceId` in the user namespace. The top 16 bits of `id` must be
52    /// 0, because this ID is packed into 48 bits of
53    /// `mz_repr::GlobalId::IntrospectionSourceIndex`.
54    pub fn user(id: u64) -> Option<Self> {
55        Self::new(id, Self::User)
56    }
57
58    fn new(id: u64, variant: fn(u64) -> Self) -> Option<Self> {
59        const MASK: u64 = 0xFFFF << 48;
60        const WARN_MASK: u64 = 1 << 47;
61        if MASK & id == 0 {
62            if WARN_MASK & id != 0 {
63                error!("{WARN_MASK} or more `StorageInstanceId`s allocated, we will run out soon");
64            }
65            Some(variant(id))
66        } else {
67            None
68        }
69    }
70
71    /// Returns the numeric id without its namespace.
72    pub fn inner_id(&self) -> u64 {
73        match self {
74            StorageInstanceId::System(id) | StorageInstanceId::User(id) => *id,
75        }
76    }
77
78    /// Returns whether this is a user storage instance.
79    pub fn is_user(&self) -> bool {
80        matches!(self, Self::User(_))
81    }
82
83    /// Returns whether this is a system storage instance.
84    pub fn is_system(&self) -> bool {
85        matches!(self, Self::System(_))
86    }
87}
88
89impl FromStr for StorageInstanceId {
90    type Err = anyhow::Error;
91
92    fn from_str(s: &str) -> Result<Self, Self::Err> {
93        // Validate the (single-byte, ASCII) tag before slicing so that a
94        // multi-byte leading character doesn't slice inside a UTF-8 boundary.
95        let variant = match s.chars().next() {
96            Some('s') => Self::System,
97            Some('u') => Self::User,
98            _ => bail!("couldn't parse compute instance id {}", s),
99        };
100        let val: u64 = s[1..].parse()?;
101        Ok(variant(val))
102    }
103}
104
105impl fmt::Display for StorageInstanceId {
106    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
107        match self {
108            Self::System(id) => write!(f, "s{}", id),
109            Self::User(id) => write!(f, "u{}", id),
110        }
111    }
112}
113
114#[cfg(any(test, feature = "proptest"))]
115impl Arbitrary for StorageInstanceId {
116    type Parameters = ();
117    type Strategy = BoxedStrategy<StorageInstanceId>;
118
119    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
120        // The inner id must fit in 48 bits: the top 16 are reserved because the
121        // id gets packed into `mz_repr::GlobalId::IntrospectionSourceIndex` (see
122        // `Self::new`). Only generate ids in that valid range so we never produce
123        // an instance that couldn't actually be allocated. Build the variants
124        // directly rather than via `Self::system`/`Self::user` to avoid their
125        // soft "running out of IDs" warning firing during tests.
126        (proptest::arbitrary::any::<bool>(), 0u64..(1 << 48))
127            .prop_map(|(is_system, id)| {
128                if is_system {
129                    StorageInstanceId::System(id)
130                } else {
131                    StorageInstanceId::User(id)
132                }
133            })
134            .boxed()
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use proptest::prelude::*;
141
142    use super::*;
143
144    #[mz_ore::test]
145    fn proptest_storage_instance_id_roundtrips() {
146        fn testcase(og: StorageInstanceId) {
147            let s = og.to_string();
148            let rnd: StorageInstanceId = s.parse().unwrap();
149            assert_eq!(og, rnd);
150        }
151
152        proptest!(|(id in any::<StorageInstanceId>())| {
153            testcase(id);
154        })
155    }
156
157    #[mz_ore::test]
158    fn test_storage_instance_id_from_str() {
159        assert_eq!(
160            "s5".parse::<StorageInstanceId>().unwrap(),
161            StorageInstanceId::System(5)
162        );
163        assert_eq!(
164            "u5".parse::<StorageInstanceId>().unwrap(),
165            StorageInstanceId::User(5)
166        );
167
168        // Regression test for a panic on multi-byte leading characters, where
169        // slicing off a single byte landed inside a UTF-8 char boundary (SQL-195).
170        for invalid in ["ü1", "ü", "é42", "🦀7", "", "x1", "u"] {
171            assert!(
172                invalid.parse::<StorageInstanceId>().is_err(),
173                "expected {invalid:?} to fail to parse"
174            );
175        }
176    }
177}
mz_cluster_client/instances.rs

mz_cluster_client/
instances.rs
