canic_core/model/memory/sharding/
mod.rs

1mod registry;
2
3pub use registry::*;
4
5use crate::{
6    cdk::{
7        structures::{BTreeMap, DefaultMemoryImpl, Memory, memory::VirtualMemory},
8        types::Principal,
9    },
10    eager_static, ic_memory,
11    ids::CanisterRole,
12    memory::impl_storable_bounded,
13    model::memory::id::sharding::{SHARDING_ASSIGNMENT_ID, SHARDING_REGISTRY_ID},
14    types::{BoundedString64, BoundedString128},
15};
16use candid::CandidType;
17use serde::{Deserialize, Serialize};
18use std::cell::RefCell;
19
20//
21// SHARDING CORE
22//
23
24eager_static! {
25    static SHARDING_CORE: RefCell<ShardingCore<VirtualMemory<DefaultMemoryImpl>>> = RefCell::new(
26        ShardingCore::new(
27            BTreeMap::init(ic_memory!(ShardingRegistry, SHARDING_REGISTRY_ID)),
28            BTreeMap::init(ic_memory!(ShardingRegistry, SHARDING_ASSIGNMENT_ID)),
29        )
30    );
31}
32
33///
34/// ShardKey
35/// Composite key: (pool, tenant) → shard
36///
37
38#[derive(CandidType, Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
39pub struct ShardKey {
40    pub pool: BoundedString64,
41    pub tenant: BoundedString128,
42}
43
44impl ShardKey {
45    pub const STORABLE_MAX_SIZE: u32 = 192;
46
47    pub(crate) fn try_new(pool: &str, tenant: &str) -> Result<Self, String> {
48        Ok(Self {
49            pool: pool.try_into()?,
50            tenant: tenant.try_into()?,
51        })
52    }
53}
54
55impl_storable_bounded!(ShardKey, ShardKey::STORABLE_MAX_SIZE, false);
56
57///
58/// ShardEntry
59/// (bare-bones; policy like has_capacity is higher-level)
60///
61
62#[derive(CandidType, Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
63pub struct ShardEntry {
64    /// Logical slot index within the pool (assigned deterministically).
65    #[serde(default = "ShardEntry::slot_default")]
66    pub slot: u32,
67    pub capacity: u32,
68    pub count: u32,
69    pub pool: BoundedString64,
70    pub canister_role: CanisterRole,
71    pub created_at: u64,
72}
73
74impl ShardEntry {
75    pub const STORABLE_MAX_SIZE: u32 = 240;
76    pub const UNASSIGNED_SLOT: u32 = u32::MAX;
77
78    pub(crate) fn try_new(
79        pool: &str,
80        slot: u32,
81        role: CanisterRole,
82        capacity: u32,
83        created_at: u64,
84    ) -> Result<Self, String> {
85        let pool = BoundedString64::try_new(pool).map_err(|err| format!("pool name: {err}"))?;
86
87        Ok(Self {
88            slot,
89            canister_role: role,
90            capacity,
91            count: 0,
92            pool,
93            created_at,
94        })
95    }
96
97    /// Whether this shard has room for more tenants.
98    #[must_use]
99    pub const fn has_capacity(&self) -> bool {
100        self.count < self.capacity
101    }
102
103    /// Returns load in basis points (0–10_000), or `None` if capacity is 0.
104    #[must_use]
105    pub const fn load_bps(&self) -> Option<u64> {
106        if self.capacity == 0 {
107            None
108        } else {
109            Some((self.count as u64).saturating_mul(10_000) / self.capacity as u64)
110        }
111    }
112
113    const fn slot_default() -> u32 {
114        Self::UNASSIGNED_SLOT
115    }
116
117    #[must_use]
118    pub const fn has_assigned_slot(&self) -> bool {
119        self.slot != Self::UNASSIGNED_SLOT
120    }
121}
122
123impl_storable_bounded!(ShardEntry, ShardEntry::STORABLE_MAX_SIZE, false);
124
125///
126/// ShardingCore
127/// Registry + assignments
128///
129
130pub struct ShardingCore<M: Memory> {
131    registry: BTreeMap<Principal, ShardEntry, M>,
132    assignments: BTreeMap<ShardKey, Principal, M>,
133}
134
135impl<M: Memory> ShardingCore<M> {
136    pub const fn new(
137        registry: BTreeMap<Principal, ShardEntry, M>,
138        assignments: BTreeMap<ShardKey, Principal, M>,
139    ) -> Self {
140        Self {
141            registry,
142            assignments,
143        }
144    }
145
146    // ---------------------------
147    // Registry CRUD
148    // ---------------------------
149
150    pub(crate) fn insert_entry(&mut self, pid: Principal, entry: ShardEntry) {
151        self.registry.insert(pid, entry);
152    }
153
154    pub(crate) fn get_entry(&self, pid: &Principal) -> Option<ShardEntry> {
155        self.registry.get(pid)
156    }
157
158    pub(crate) fn all_entries(&self) -> Vec<(Principal, ShardEntry)> {
159        self.registry
160            .iter()
161            .map(|e| (*e.key(), e.value()))
162            .collect()
163    }
164
165    // ---------------------------
166    // Assignments CRUD
167    // ---------------------------
168
169    pub(crate) fn insert_assignment(&mut self, key: ShardKey, shard: Principal) {
170        self.assignments.insert(key, shard);
171    }
172
173    pub(crate) fn remove_assignment(&mut self, key: &ShardKey) -> Option<Principal> {
174        self.assignments.remove(key)
175    }
176
177    pub(crate) fn get_assignment(&self, key: &ShardKey) -> Option<Principal> {
178        self.assignments.get(key)
179    }
180
181    pub(crate) fn all_assignments(&self) -> Vec<(ShardKey, Principal)> {
182        self.assignments
183            .iter()
184            .map(|e| (e.key().clone(), e.value()))
185            .collect()
186    }
187}
canic_core/model/memory/sharding/mod.rs

canic_core/model/memory/sharding/
mod.rs
