canic_core/model/memory/sharding/
mod.rs

1mod registry;
2
3pub(crate) use registry::ShardingRegistry;
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::{BoundedString32, 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: BoundedString32,
41    pub tenant: BoundedString128,
42}
43
44impl ShardKey {
45    pub const STORABLE_MAX_SIZE: u32 = 160;
46
47    #[must_use]
48    pub(crate) fn new(pool: &str, tenant: &str) -> Self {
49        Self {
50            pool: pool.try_into().unwrap(),
51            tenant: tenant.try_into().unwrap(),
52        }
53    }
54}
55
56impl_storable_bounded!(ShardKey, ShardKey::STORABLE_MAX_SIZE, false);
57
58///
59/// ShardEntry
60/// (bare-bones; policy like has_capacity is higher-level)
61///
62
63#[derive(CandidType, Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
64pub struct ShardEntry {
65    /// Logical slot index within the pool (assigned deterministically).
66    #[serde(default = "ShardEntry::slot_default")]
67    pub slot: u32,
68    pub capacity: u32,
69    pub count: u32,
70    pub pool: String,
71    pub canister_type: CanisterRole,
72    pub created_at: u64,
73}
74
75impl ShardEntry {
76    pub const STORABLE_MAX_SIZE: u32 = 208;
77    pub const UNASSIGNED_SLOT: u32 = u32::MAX;
78
79    #[must_use]
80    pub(crate) fn new(
81        pool: &str,
82        slot: u32,
83        ty: CanisterRole,
84        capacity: u32,
85        created_at: u64,
86    ) -> Self {
87        Self {
88            slot,
89            canister_type: ty,
90            capacity,
91            count: 0,
92            pool: pool.to_string(),
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    #[inline]
114    const fn slot_default() -> u32 {
115        Self::UNASSIGNED_SLOT
116    }
117
118    #[must_use]
119    pub const fn has_assigned_slot(&self) -> bool {
120        self.slot != Self::UNASSIGNED_SLOT
121    }
122}
123
124impl_storable_bounded!(ShardEntry, ShardEntry::STORABLE_MAX_SIZE, false);
125
126///
127/// ShardingCore
128/// Registry + assignments
129///
130
131pub(crate) struct ShardingCore<M: Memory> {
132    registry: BTreeMap<Principal, ShardEntry, M>,
133    assignments: BTreeMap<ShardKey, Principal, M>,
134}
135
136impl<M: Memory> ShardingCore<M> {
137    pub const fn new(
138        registry: BTreeMap<Principal, ShardEntry, M>,
139        assignments: BTreeMap<ShardKey, Principal, M>,
140    ) -> Self {
141        Self {
142            registry,
143            assignments,
144        }
145    }
146
147    // ---------------------------
148    // Registry CRUD
149    // ---------------------------
150
151    pub(crate) fn insert_entry(&mut self, pid: Principal, entry: ShardEntry) {
152        self.registry.insert(pid, entry);
153    }
154
155    pub(crate) fn get_entry(&self, pid: &Principal) -> Option<ShardEntry> {
156        self.registry.get(pid)
157    }
158
159    pub(crate) fn all_entries(&self) -> Vec<(Principal, ShardEntry)> {
160        self.registry
161            .iter()
162            .map(|e| (*e.key(), e.value()))
163            .collect()
164    }
165
166    // ---------------------------
167    // Assignments CRUD
168    // ---------------------------
169
170    pub(crate) fn insert_assignment(&mut self, key: ShardKey, shard: Principal) {
171        self.assignments.insert(key, shard);
172    }
173
174    pub(crate) fn remove_assignment(&mut self, key: &ShardKey) -> Option<Principal> {
175        self.assignments.remove(key)
176    }
177
178    pub(crate) fn get_assignment(&self, key: &ShardKey) -> Option<Principal> {
179        self.assignments.get(key)
180    }
181
182    pub(crate) fn all_assignments(&self) -> Vec<(ShardKey, Principal)> {
183        self.assignments
184            .iter()
185            .map(|e| (e.key().clone(), e.value()))
186            .collect()
187    }
188}
canic_core/model/memory/sharding/mod.rs

canic_core/model/memory/sharding/
mod.rs
