mz_storage_types/
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    pub fn inner_id(&self) -> u64 {
72        match self {
73            StorageInstanceId::System(id) | StorageInstanceId::User(id) => *id,
74        }
75    }
76
77    pub fn is_user(&self) -> bool {
78        matches!(self, Self::User(_))
79    }
80
81    pub fn is_system(&self) -> bool {
82        matches!(self, Self::System(_))
83    }
84}
85
86impl FromStr for StorageInstanceId {
87    type Err = anyhow::Error;
88
89    fn from_str(s: &str) -> Result<Self, Self::Err> {
90        // Validate the (single-byte, ASCII) tag before slicing so that a
91        // multi-byte leading character doesn't slice inside a UTF-8 boundary.
92        let variant = match s.chars().next() {
93            Some('s') => Self::System,
94            Some('u') => Self::User,
95            _ => bail!("couldn't parse compute instance id {}", s),
96        };
97        let val: u64 = s[1..].parse()?;
98        Ok(variant(val))
99    }
100}
101
102impl fmt::Display for StorageInstanceId {
103    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
104        match self {
105            Self::System(id) => write!(f, "s{}", id),
106            Self::User(id) => write!(f, "u{}", id),
107        }
108    }
109}
110
111#[cfg(any(test, feature = "proptest"))]
112impl Arbitrary for StorageInstanceId {
113    type Parameters = ();
114    type Strategy = BoxedStrategy<StorageInstanceId>;
115
116    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
117        // The inner id must fit in 48 bits: the top 16 are reserved because the
118        // id gets packed into `mz_repr::GlobalId::IntrospectionSourceIndex` (see
119        // `Self::new`). Only generate ids in that valid range so we never produce
120        // an instance that couldn't actually be allocated. Build the variants
121        // directly rather than via `Self::system`/`Self::user` to avoid their
122        // soft "running out of IDs" warning firing during tests.
123        (proptest::arbitrary::any::<bool>(), 0u64..(1 << 48))
124            .prop_map(|(is_system, id)| {
125                if is_system {
126                    StorageInstanceId::System(id)
127                } else {
128                    StorageInstanceId::User(id)
129                }
130            })
131            .boxed()
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use proptest::prelude::*;
138
139    use super::*;
140
141    #[mz_ore::test]
142    fn proptest_storage_instance_id_roundtrips() {
143        fn testcase(og: StorageInstanceId) {
144            let s = og.to_string();
145            let rnd: StorageInstanceId = s.parse().unwrap();
146            assert_eq!(og, rnd);
147        }
148
149        proptest!(|(id in any::<StorageInstanceId>())| {
150            testcase(id);
151        })
152    }
153
154    #[mz_ore::test]
155    fn test_storage_instance_id_from_str() {
156        assert_eq!(
157            "s5".parse::<StorageInstanceId>().unwrap(),
158            StorageInstanceId::System(5)
159        );
160        assert_eq!(
161            "u5".parse::<StorageInstanceId>().unwrap(),
162            StorageInstanceId::User(5)
163        );
164
165        // Regression test for a panic on multi-byte leading characters, where
166        // slicing off a single byte landed inside a UTF-8 char boundary (SQL-195).
167        for invalid in ["ü1", "ü", "é42", "🦀7", "", "x1", "u"] {
168            assert!(
169                invalid.parse::<StorageInstanceId>().is_err(),
170                "expected {invalid:?} to fail to parse"
171            );
172        }
173    }
174}
mz_storage_types/instances.rs

mz_storage_types/
instances.rs
