canic_core/model/memory/sharding/
mod.rs

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

canic_core/model/memory/sharding/
mod.rs
