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

canic_core/model/memory/
scaling.rs
