mz_storage_types/
instances.rs

Help
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;
16use mz_proto::{RustType, TryFromProtoError};
17use proptest::prelude::{Arbitrary, Strategy};
18use proptest::strategy::BoxedStrategy;
19use serde::{Deserialize, Serialize};
20use tracing::error;
21
22include!(concat!(env!("OUT_DIR"), "/mz_storage_types.instances.rs"));
23
24/// Identifier of a storage instance.
25#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
26pub enum StorageInstanceId {
27    /// A system storage instance.
28    System(u64),
29    /// A user storage instance.
30    User(u64),
31}
32
33impl StorageInstanceId {
34    /// Creates a new `StorageInstanceId` in the system namespace. The top 16 bits of `id` must be
35    /// 0, because this ID is packed into 48 bits of
36    /// [`mz_repr::GlobalId::IntrospectionSourceIndex`].
37    pub fn system(id: u64) -> Option<Self> {
38        Self::new(id, Self::System)
39    }
40
41    /// Creates a new `StorageInstanceId` in the user namespace. The top 16 bits of `id` must be
42    /// 0, because this ID is packed into 48 bits of
43    /// [`mz_repr::GlobalId::IntrospectionSourceIndex`].
44    pub fn user(id: u64) -> Option<Self> {
45        Self::new(id, Self::User)
46    }
47
48    fn new(id: u64, variant: fn(u64) -> Self) -> Option<Self> {
49        const MASK: u64 = 0xFFFF << 48;
50        const WARN_MASK: u64 = 1 << 47;
51        if MASK & id == 0 {
52            if WARN_MASK & id != 0 {
53                error!("{WARN_MASK} or more `StorageInstanceId`s allocated, we will run out soon");
54            }
55            Some(variant(id))
56        } else {
57            None
58        }
59    }
60
61    pub fn inner_id(&self) -> u64 {
62        match self {
63            StorageInstanceId::System(id) | StorageInstanceId::User(id) => *id,
64        }
65    }
66
67    pub fn is_user(&self) -> bool {
68        matches!(self, Self::User(_))
69    }
70
71    pub fn is_system(&self) -> bool {
72        matches!(self, Self::System(_))
73    }
74}
75
76impl FromStr for StorageInstanceId {
77    type Err = anyhow::Error;
78
79    fn from_str(s: &str) -> Result<Self, Self::Err> {
80        if s.len() < 2 {
81            bail!("couldn't parse compute instance id {}", s);
82        }
83        let val: u64 = s[1..].parse()?;
84        match s.chars().next().unwrap() {
85            's' => Ok(Self::System(val)),
86            'u' => Ok(Self::User(val)),
87            _ => bail!("couldn't parse compute instance id {}", s),
88        }
89    }
90}
91
92impl fmt::Display for StorageInstanceId {
93    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
94        match self {
95            Self::System(id) => write!(f, "s{}", id),
96            Self::User(id) => write!(f, "u{}", id),
97        }
98    }
99}
100
101impl RustType<ProtoStorageInstanceId> for StorageInstanceId {
102    fn into_proto(&self) -> ProtoStorageInstanceId {
103        use proto_storage_instance_id::Kind::*;
104        ProtoStorageInstanceId {
105            kind: Some(match self {
106                StorageInstanceId::System(x) => System(*x),
107                StorageInstanceId::User(x) => User(*x),
108            }),
109        }
110    }
111
112    fn from_proto(proto: ProtoStorageInstanceId) -> Result<Self, TryFromProtoError> {
113        use proto_storage_instance_id::Kind::*;
114        match proto.kind {
115            Some(System(x)) => StorageInstanceId::system(x).ok_or_else(|| {
116                TryFromProtoError::InvalidPersistState(format!(
117                    "{x} is not a valid StorageInstanceId"
118                ))
119            }),
120            Some(User(x)) => StorageInstanceId::user(x).ok_or_else(|| {
121                TryFromProtoError::InvalidPersistState(format!(
122                    "{x} is not a valid StorageInstanceId"
123                ))
124            }),
125            None => Err(TryFromProtoError::missing_field(
126                "ProtoStorageInstanceId::kind",
127            )),
128        }
129    }
130}
131
132impl Arbitrary for StorageInstanceId {
133    type Parameters = ();
134
135    fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
136        const UPPER_BOUND: u64 = 1 << 47;
137        (0..2, 0..UPPER_BOUND)
138            .prop_map(|(variant, id)| match variant {
139                0 => StorageInstanceId::System(id),
140                1 => StorageInstanceId::User(id),
141                _ => unreachable!(),
142            })
143            .boxed()
144    }
145
146    type Strategy = BoxedStrategy<StorageInstanceId>;
147}
mz_storage_types/instances.rs

mz_storage_types/
instances.rs
