canic_core/model/memory/
scaling.rs

1use crate::{
2    cdk::structures::{BTreeMap, DefaultMemoryImpl, memory::VirtualMemory},
3    eager_static, ic_memory,
4    ids::CanisterRole,
5    memory::impl_storable_bounded,
6    model::memory::id::scaling::SCALING_REGISTRY_ID,
7    types::BoundedString64,
8};
9use candid::{CandidType, Principal};
10use serde::{Deserialize, Serialize};
11use std::cell::RefCell;
12
13//
14// SCALING REGISTRY
15//
16
17eager_static! {
18    static SCALING_REGISTRY: RefCell<
19        BTreeMap<Principal, WorkerEntry, VirtualMemory<DefaultMemoryImpl>>
20    > = RefCell::new(
21        BTreeMap::init(ic_memory!(ScalingRegistry, SCALING_REGISTRY_ID)),
22    );
23}
24
25///
26/// WorkerEntry
27///
28
29#[derive(CandidType, Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
30pub struct WorkerEntry {
31    pub pool: BoundedString64,       // which scale pool this belongs to
32    pub canister_role: CanisterRole, // canister role
33    pub created_at_secs: u64,        // timestamp
34}
35
36impl WorkerEntry {
37    pub const STORABLE_MAX_SIZE: u32 = 160;
38
39    pub(crate) fn try_new(
40        pool: &str,
41        canister_role: CanisterRole,
42        created_at_secs: u64,
43    ) -> Result<Self, String> {
44        let pool = BoundedString64::try_new(pool).map_err(|err| format!("pool name: {err}"))?;
45
46        Ok(Self {
47            pool,
48            canister_role,
49            created_at_secs,
50        })
51    }
52}
53
54impl_storable_bounded!(WorkerEntry, WorkerEntry::STORABLE_MAX_SIZE, false);
55
56///
57/// ScalingRegistryView
58///
59
60pub type ScalingRegistryView = Vec<(Principal, WorkerEntry)>;
61
62///
63/// ScalingRegistry
64/// Registry of active scaling workers
65///
66
67#[derive(Clone, Copy, Debug, Default)]
68pub struct ScalingRegistry;
69
70impl ScalingRegistry {
71    /// Insert or update a worker entry
72    pub(crate) fn insert(pid: Principal, entry: WorkerEntry) {
73        SCALING_REGISTRY.with_borrow_mut(|map| {
74            map.insert(pid, entry);
75        });
76    }
77
78    /// Lookup all workers in a given pool
79    #[must_use]
80    pub(crate) fn find_by_pool(pool: &str) -> Vec<(Principal, WorkerEntry)> {
81        SCALING_REGISTRY.with_borrow(|map| {
82            map.iter()
83                .filter(|e| e.value().pool.as_ref() == pool)
84                .map(|e| (*e.key(), e.value()))
85                .collect()
86        })
87    }
88
89    /// Export full registry
90    #[must_use]
91    pub(crate) fn export() -> Vec<(Principal, WorkerEntry)> {
92        SCALING_REGISTRY.with_borrow(|map| map.iter().map(|e| (*e.key(), e.value())).collect())
93    }
94}
canic_core/model/memory/scaling.rs

canic_core/model/memory/
scaling.rs
