canic/model/memory/
scaling.rs

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

canic/model/memory/
scaling.rs
