canic-core 0.26.9

pub mod lifecycle;
pub mod registry;

use crate::{
    cdk::{
        structures::{BTreeMap, DefaultMemoryImpl, Memory, memory::VirtualMemory},
        types::{BoundedString64, BoundedString128},
    },
    storage::{
        prelude::*,
        stable::{
            memory::placement::{SHARDING_ASSIGNMENT_ID, SHARDING_REGISTRY_ID},
            sharding::registry::ShardingRegistry,
        },
    },
};
use std::cell::RefCell;

//
// SHARDING CORE
//

eager_static! {
    static SHARDING_CORE: RefCell<ShardingCore<VirtualMemory<DefaultMemoryImpl>>> = RefCell::new(
        ShardingCore::new(
            BTreeMap::init(ic_memory!(ShardingRegistry, SHARDING_REGISTRY_ID)),
            BTreeMap::init(ic_memory!(ShardingRegistry, SHARDING_ASSIGNMENT_ID)),
        )
    );
}

///
/// ShardKey
/// Composite key: (pool, partition_key) → shard
///

#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
pub struct ShardKey {
    pub pool: BoundedString64,
    pub partition_key: BoundedString128,
}

impl ShardKey {
    pub const STORABLE_MAX_SIZE: u32 = 192;

    pub(crate) fn try_new(pool: &str, partition_key: &str) -> Result<Self, String> {
        Ok(Self {
            pool: pool.try_into()?,
            partition_key: partition_key.try_into()?,
        })
    }
}

impl_storable_bounded!(ShardKey, ShardKey::STORABLE_MAX_SIZE, false);

///
/// ShardEntryRecord
/// (bare-bones; policy like has_capacity is higher-level)
///

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct ShardEntryRecord {
    /// Logical slot index within the pool (assigned deterministically).
    #[serde(default = "ShardEntryRecord::slot_default")]
    pub slot: u32,
    pub capacity: u32,
    pub count: u32,
    pub pool: BoundedString64,
    pub canister_role: CanisterRole,
    pub created_at: u64,
}

impl ShardEntryRecord {
    pub const STORABLE_MAX_SIZE: u32 = 240;
    pub const UNASSIGNED_SLOT: u32 = u32::MAX;

    pub(crate) fn try_new(
        pool: &str,
        slot: u32,
        role: CanisterRole,
        capacity: u32,
        created_at: u64,
    ) -> Result<Self, String> {
        let pool = BoundedString64::try_new(pool).map_err(|err| format!("pool name: {err}"))?;

        Ok(Self {
            slot,
            canister_role: role,
            capacity,
            count: 0,
            pool,
            created_at,
        })
    }

    const fn slot_default() -> u32 {
        Self::UNASSIGNED_SLOT
    }

    #[must_use]
    pub const fn has_assigned_slot(&self) -> bool {
        self.slot != Self::UNASSIGNED_SLOT
    }
}

impl_storable_bounded!(ShardEntryRecord, ShardEntryRecord::STORABLE_MAX_SIZE, false);

///
/// ShardingCore
/// Registry + assignments
///

pub struct ShardingCore<M: Memory> {
    registry: BTreeMap<Principal, ShardEntryRecord, M>,
    assignments: BTreeMap<ShardKey, Principal, M>,
}

impl<M: Memory> ShardingCore<M> {
    pub const fn new(
        registry: BTreeMap<Principal, ShardEntryRecord, M>,
        assignments: BTreeMap<ShardKey, Principal, M>,
    ) -> Self {
        Self {
            registry,
            assignments,
        }
    }

    // ---------------------------
    // Registry CRUD
    // ---------------------------

    pub fn insert_entry(&mut self, pid: Principal, entry: ShardEntryRecord) {
        self.registry.insert(pid, entry);
    }

    pub fn get_entry(&self, pid: &Principal) -> Option<ShardEntryRecord> {
        self.registry.get(pid)
    }

    pub fn all_entries(&self) -> Vec<(Principal, ShardEntryRecord)> {
        self.registry
            .iter()
            .map(|e| (*e.key(), e.value()))
            .collect()
    }

    // ---------------------------
    // Assignments CRUD
    // ---------------------------

    pub fn insert_assignment(&mut self, key: ShardKey, shard: Principal) {
        self.assignments.insert(key, shard);
    }

    #[expect(dead_code)] // Used by future rebalance / eviction workflows
    pub fn remove_assignment(&mut self, key: &ShardKey) -> Option<Principal> {
        self.assignments.remove(key)
    }

    pub fn get_assignment(&self, key: &ShardKey) -> Option<Principal> {
        self.assignments.get(key)
    }

    pub fn all_assignments(&self) -> Vec<(ShardKey, Principal)> {
        self.assignments
            .iter()
            .map(|e| (e.key().clone(), e.value()))
            .collect()
    }
}