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 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 fn new(pool: &str, slot: u32, ty: CanisterRole, capacity: u32, created_at: u64) -> Self {
78        Self {
79            slot,
80            canister_type: ty,
81            capacity,
82            count: 0,
83            pool: pool.to_string(),
84            created_at,
85        }
86    }
87
88    /// Whether this shard has room for more tenants.
89    #[must_use]
90    pub const fn has_capacity(&self) -> bool {
91        self.count < self.capacity
92    }
93
94    /// Returns load in basis points (0–10_000), or `None` if capacity is 0.
95    #[must_use]
96    pub const fn load_bps(&self) -> Option<u64> {
97        if self.capacity == 0 {
98            None
99        } else {
100            Some((self.count as u64).saturating_mul(10_000) / self.capacity as u64)
101        }
102    }
103
104    #[inline]
105    const fn slot_default() -> u32 {
106        Self::UNASSIGNED_SLOT
107    }
108
109    #[must_use]
110    pub const fn has_assigned_slot(&self) -> bool {
111        self.slot != Self::UNASSIGNED_SLOT
112    }
113}
114
115impl_storable_bounded!(ShardEntry, ShardEntry::STORABLE_MAX_SIZE, false);
116
117///
118/// ShardingCore
119/// Registry + assignments
120///
121
122pub(crate) struct ShardingCore<M: Memory> {
123    registry: BTreeMap<Principal, ShardEntry, M>,
124    assignments: BTreeMap<ShardKey, Principal, M>,
125}
126
127impl<M: Memory> ShardingCore<M> {
128    pub const fn new(
129        registry: BTreeMap<Principal, ShardEntry, M>,
130        assignments: BTreeMap<ShardKey, Principal, M>,
131    ) -> Self {
132        Self {
133            registry,
134            assignments,
135        }
136    }
137
138    // ---------------------------
139    // Registry CRUD
140    // ---------------------------
141
142    pub(crate) fn insert_entry(&mut self, pid: Principal, entry: ShardEntry) {
143        self.registry.insert(pid, entry);
144    }
145
146    pub(crate) fn get_entry(&self, pid: &Principal) -> Option<ShardEntry> {
147        self.registry.get(pid)
148    }
149
150    pub(crate) fn all_entries(&self) -> Vec<(Principal, ShardEntry)> {
151        self.registry
152            .iter()
153            .map(|e| (*e.key(), e.value()))
154            .collect()
155    }
156
157    // ---------------------------
158    // Assignments CRUD
159    // ---------------------------
160
161    pub(crate) fn insert_assignment(&mut self, key: ShardKey, shard: Principal) {
162        self.assignments.insert(key, shard);
163    }
164
165    pub(crate) fn remove_assignment(&mut self, key: &ShardKey) -> Option<Principal> {
166        self.assignments.remove(key)
167    }
168
169    pub(crate) fn get_assignment(&self, key: &ShardKey) -> Option<Principal> {
170        self.assignments.get(key)
171    }
172
173    pub(crate) fn all_assignments(&self) -> Vec<(ShardKey, Principal)> {
174        self.assignments
175            .iter()
176            .map(|e| (e.key().clone(), e.value()))
177            .collect()
178    }
179}
canic_core/model/memory/sharding/mod.rs

canic_core/model/memory/sharding/
mod.rs
