canic-core 0.56.0

Canic — a canister orchestration and management toolkit for the Internet Computer
Documentation
use super::ShardingWorkflow;
use crate::{
    cdk::types::Principal,
    ops::{
        placement::sharding::mapper::ShardPlacementPolicyInputMapper,
        storage::{children::CanisterChildrenOps, placement::sharding::ShardingRegistryOps},
    },
    view::placement::sharding::ShardPlacement,
};
use std::collections::BTreeSet;

impl ShardingWorkflow {
    pub(super) fn pool_entry_views(pool: &str) -> Vec<(Principal, ShardPlacement)> {
        let direct_children = Self::direct_child_pid_set();
        ShardingRegistryOps::entries_for_pool(pool)
            .iter()
            .filter(|(pid, _)| direct_children.is_empty() || direct_children.contains(pid))
            .map(|(pid, entry)| {
                ShardPlacementPolicyInputMapper::record_to_policy_input(*pid, entry)
            })
            .collect()
    }

    pub(super) fn routable_active_set(active: &BTreeSet<Principal>) -> BTreeSet<Principal> {
        let direct_children = Self::direct_child_pid_set();
        if direct_children.is_empty() {
            return active.clone();
        }

        active.intersection(&direct_children).copied().collect()
    }

    fn direct_child_pid_set() -> BTreeSet<Principal> {
        CanisterChildrenOps::data()
            .entries
            .into_iter()
            .map(|(pid, _)| pid)
            .collect()
    }

    pub(super) fn free_slots(max_shards: u32, entries: &[(Principal, ShardPlacement)]) -> Vec<u32> {
        let mut occupied = BTreeSet::new();
        for (_, entry) in entries {
            if entry.slot != ShardPlacement::UNASSIGNED_SLOT {
                occupied.insert(entry.slot);
            }
        }

        (0..max_shards)
            .filter(|slot| !occupied.contains(slot))
            .collect()
    }
}