Public release release-6d1a121b0a8a

Source commit: 6d1a121b0a8a636aac95998fe5d15c24c78eb7d0

Public tree identity: sha256:2bc22807f5b838286678b65d3f550b8ae4714ae673622041d4c2516a7c5b7926

crates/clusterflux-coordinator/src/service/quota.rs

@@ -0,0 +1,484 @@
+use std::collections::BTreeMap;
+
+use clusterflux_core::{LimitError, LimitKind, ProjectId, ResourceLimits, ResourceMeter, TenantId};
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct CoordinatorQuotaConfiguration {
+    pub limits: ResourceLimits,
+    pub window_seconds: BTreeMap<LimitKind, u64>,
+    pub policy_label: Option<String>,
+    pub max_projects_per_tenant: usize,
+    pub max_nodes_per_tenant: usize,
+    pub max_active_processes_per_tenant: usize,
+}
+
+pub(super) struct CoordinatorQuotaStatus {
+    pub(super) policy_label: Option<String>,
+    pub(super) limits: ResourceLimits,
+    pub(super) window_seconds: BTreeMap<LimitKind, u64>,
+    pub(super) usage: BTreeMap<LimitKind, u64>,
+    pub(super) window_started_epoch_seconds: BTreeMap<LimitKind, u64>,
+}
+
+impl CoordinatorQuotaConfiguration {
+    pub fn new(
+        limits: ResourceLimits,
+        window_seconds: impl IntoIterator<Item = (LimitKind, u64)>,
+    ) -> Result<Self, String> {
+        let window_seconds = window_seconds.into_iter().collect::<BTreeMap<_, _>>();
+        if window_seconds.values().any(|seconds| *seconds == 0) {
+            return Err("quota windows must be at least one second".to_owned());
+        }
+        Ok(Self {
+            limits,
+            window_seconds,
+            policy_label: None,
+            max_projects_per_tenant: usize::MAX,
+            max_nodes_per_tenant: usize::MAX,
+            max_active_processes_per_tenant: usize::MAX,
+        })
+    }
+
+    pub fn with_policy_label(mut self, label: impl Into<String>) -> Self {
+        let label = label.into();
+        self.policy_label = (!label.trim().is_empty()).then_some(label);
+        self
+    }
+
+    pub fn with_admission_limits(
+        mut self,
+        max_projects_per_tenant: usize,
+        max_nodes_per_tenant: usize,
+        max_active_processes_per_tenant: usize,
+    ) -> Self {
+        self.max_projects_per_tenant = max_projects_per_tenant;
+        self.max_nodes_per_tenant = max_nodes_per_tenant;
+        self.max_active_processes_per_tenant = max_active_processes_per_tenant;
+        self
+    }
+
+    pub fn unlimited() -> Self {
+        Self {
+            limits: ResourceLimits::unlimited(),
+            window_seconds: LimitKind::ALL
+                .into_iter()
+                .map(|kind| (kind, u64::MAX))
+                .collect(),
+            policy_label: None,
+            max_projects_per_tenant: usize::MAX,
+            max_nodes_per_tenant: usize::MAX,
+            max_active_processes_per_tenant: usize::MAX,
+        }
+    }
+
+    pub fn window_seconds(&self, kind: LimitKind) -> u64 {
+        self.window_seconds
+            .get(&kind)
+            .copied()
+            .unwrap_or(u64::MAX)
+            .max(1)
+    }
+}
+
+impl Default for CoordinatorQuotaConfiguration {
+    fn default() -> Self {
+        Self::unlimited()
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+struct ProjectQuotaScope {
+    tenant: TenantId,
+    project: ProjectId,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+struct MeterKey {
+    scope: ProjectQuotaScope,
+    kind: LimitKind,
+    window: u64,
+}
+
+#[derive(Clone, Debug)]
+pub(super) struct CoordinatorQuota {
+    configuration: CoordinatorQuotaConfiguration,
+    meters: BTreeMap<MeterKey, ResourceMeter>,
+}
+
+impl Default for CoordinatorQuota {
+    fn default() -> Self {
+        Self::new(CoordinatorQuotaConfiguration::default())
+    }
+}
+
+impl CoordinatorQuota {
+    pub(super) fn new(configuration: CoordinatorQuotaConfiguration) -> Self {
+        Self {
+            configuration,
+            meters: BTreeMap::new(),
+        }
+    }
+
+    pub(super) fn ensure_project_admission(
+        &self,
+        tenant: &TenantId,
+        current: usize,
+    ) -> Result<(), super::CoordinatorServiceError> {
+        let maximum = self.configuration.max_projects_per_tenant;
+        if current >= maximum {
+            return Err(super::CoordinatorServiceError::Protocol(format!(
+                "admission.project_limit: tenant {tenant} already has the maximum of {maximum} projects"
+            )));
+        }
+        Ok(())
+    }
+
+    pub(super) fn ensure_node_admission(
+        &self,
+        tenant: &TenantId,
+        current: usize,
+    ) -> Result<(), super::CoordinatorServiceError> {
+        let maximum = self.configuration.max_nodes_per_tenant;
+        if current >= maximum {
+            return Err(super::CoordinatorServiceError::Protocol(format!(
+                "admission.node_limit: tenant {tenant} already has the maximum of {maximum} nodes"
+            )));
+        }
+        Ok(())
+    }
+
+    pub(super) fn ensure_process_admission(
+        &self,
+        tenant: &TenantId,
+        current: usize,
+    ) -> Result<(), super::CoordinatorServiceError> {
+        let maximum = self.configuration.max_active_processes_per_tenant;
+        if current >= maximum {
+            return Err(super::CoordinatorServiceError::Protocol(format!(
+                "admission.active_process_limit: tenant {tenant} already has the maximum of {maximum} active processes"
+            )));
+        }
+        Ok(())
+    }
+
+    fn key(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        kind: LimitKind,
+        now_epoch_seconds: u64,
+    ) -> MeterKey {
+        MeterKey {
+            scope: ProjectQuotaScope {
+                tenant: tenant.clone(),
+                project: project.clone(),
+            },
+            kind,
+            window: now_epoch_seconds / self.configuration.window_seconds(kind),
+        }
+    }
+
+    fn meter(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        kind: LimitKind,
+        now_epoch_seconds: u64,
+    ) -> Option<&ResourceMeter> {
+        self.meters
+            .get(&self.key(tenant, project, kind, now_epoch_seconds))
+    }
+
+    fn meter_mut(
+        &mut self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        kind: LimitKind,
+        now_epoch_seconds: u64,
+    ) -> &mut ResourceMeter {
+        let key = self.key(tenant, project, kind, now_epoch_seconds);
+        self.meters.retain(|existing, _| {
+            existing.scope != key.scope || existing.kind != kind || existing.window == key.window
+        });
+        self.meters.entry(key).or_default()
+    }
+
+    fn can_charge(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        kind: LimitKind,
+        amount: u64,
+        now_epoch_seconds: u64,
+    ) -> Result<(), LimitError> {
+        self.meter(tenant, project, kind, now_epoch_seconds)
+            .cloned()
+            .unwrap_or_default()
+            .can_charge(&self.configuration.limits, kind, amount)
+    }
+
+    fn charge(
+        &mut self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        kind: LimitKind,
+        amount: u64,
+        now_epoch_seconds: u64,
+    ) -> Result<u64, LimitError> {
+        let limits = self.configuration.limits.clone();
+        let meter = self.meter_mut(tenant, project, kind, now_epoch_seconds);
+        meter.charge(&limits, kind, amount)?;
+        Ok(meter.used(&kind))
+    }
+
+    fn used(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        kind: LimitKind,
+        now_epoch_seconds: u64,
+    ) -> u64 {
+        self.meter(tenant, project, kind, now_epoch_seconds)
+            .map_or(0, |meter| meter.used(&kind))
+    }
+
+    pub(super) fn can_charge_workflow_spawn(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> Result<(), LimitError> {
+        self.can_charge(tenant, project, LimitKind::Spawn, 1, now_epoch_seconds)
+    }
+
+    pub(super) fn charge_api_call(
+        &mut self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> Result<u64, LimitError> {
+        self.charge(tenant, project, LimitKind::ApiCall, 1, now_epoch_seconds)
+    }
+
+    pub(super) fn can_charge_log_bytes(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        bytes: u64,
+        now_epoch_seconds: u64,
+    ) -> Result<(), LimitError> {
+        self.can_charge(
+            tenant,
+            project,
+            LimitKind::LogBytes,
+            bytes,
+            now_epoch_seconds,
+        )
+    }
+
+    pub(super) fn charge_log_bytes(
+        &mut self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        bytes: u64,
+        now_epoch_seconds: u64,
+    ) -> Result<u64, LimitError> {
+        self.charge(
+            tenant,
+            project,
+            LimitKind::LogBytes,
+            bytes,
+            now_epoch_seconds,
+        )
+    }
+
+    pub(super) fn charge_workflow_spawn(
+        &mut self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> Result<u64, LimitError> {
+        self.charge(tenant, project, LimitKind::Spawn, 1, now_epoch_seconds)
+    }
+
+    #[cfg(test)]
+    pub(super) fn used_workflow_spawns(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> u64 {
+        self.used(tenant, project, LimitKind::Spawn, now_epoch_seconds)
+    }
+
+    #[cfg(test)]
+    pub(super) fn used_api_calls(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> u64 {
+        self.used(tenant, project, LimitKind::ApiCall, now_epoch_seconds)
+    }
+
+    #[cfg(test)]
+    pub(super) fn used_log_bytes(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> u64 {
+        self.used(tenant, project, LimitKind::LogBytes, now_epoch_seconds)
+    }
+
+    #[cfg(test)]
+    pub(super) fn active_meter_count(&self) -> usize {
+        self.meters.len()
+    }
+
+    pub(super) fn charge_rendezvous_attempt(
+        &mut self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> Result<u64, LimitError> {
+        self.charge(
+            tenant,
+            project,
+            LimitKind::RendezvousAttempt,
+            1,
+            now_epoch_seconds,
+        )
+    }
+
+    pub(super) fn can_charge_download(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        bytes: u64,
+        now_epoch_seconds: u64,
+    ) -> Result<(), LimitError> {
+        self.can_charge(
+            tenant,
+            project,
+            LimitKind::ArtifactDownloadBytes,
+            bytes,
+            now_epoch_seconds,
+        )
+    }
+
+    pub(super) fn charge_download(
+        &mut self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        bytes: u64,
+        now_epoch_seconds: u64,
+    ) -> Result<u64, LimitError> {
+        self.charge(
+            tenant,
+            project,
+            LimitKind::ArtifactDownloadBytes,
+            bytes,
+            now_epoch_seconds,
+        )
+    }
+
+    pub(super) fn download_limit(&self) -> u64 {
+        self.configuration
+            .limits
+            .limit(&LimitKind::ArtifactDownloadBytes)
+    }
+
+    pub(super) fn used_download_bytes(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> u64 {
+        self.used(
+            tenant,
+            project,
+            LimitKind::ArtifactDownloadBytes,
+            now_epoch_seconds,
+        )
+    }
+
+    pub(super) fn charge_debug_read(
+        &mut self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        bytes: u64,
+        now_epoch_seconds: u64,
+    ) -> Result<u64, LimitError> {
+        self.charge(
+            tenant,
+            project,
+            LimitKind::DebugReadBytes,
+            bytes,
+            now_epoch_seconds,
+        )
+    }
+
+    pub(super) fn used_debug_read_bytes(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> u64 {
+        self.used(
+            tenant,
+            project,
+            LimitKind::DebugReadBytes,
+            now_epoch_seconds,
+        )
+    }
+
+    pub(super) fn project_status(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        now_epoch_seconds: u64,
+    ) -> CoordinatorQuotaStatus {
+        let mut usage = BTreeMap::new();
+        let mut window_starts = BTreeMap::new();
+        for kind in LimitKind::ALL {
+            let seconds = self.configuration.window_seconds(kind);
+            usage.insert(kind, self.used(tenant, project, kind, now_epoch_seconds));
+            window_starts.insert(kind, (now_epoch_seconds / seconds).saturating_mul(seconds));
+        }
+        CoordinatorQuotaStatus {
+            policy_label: self.configuration.policy_label.clone(),
+            limits: self.configuration.limits.clone(),
+            window_seconds: self.configuration.window_seconds.clone(),
+            usage,
+            window_started_epoch_seconds: window_starts,
+        }
+    }
+
+    #[cfg(test)]
+    pub(super) fn set_workflow_limits(&mut self, limits: ResourceLimits) {
+        self.configuration
+            .limits
+            .limits
+            .insert(LimitKind::Spawn, limits.limit(&LimitKind::Spawn));
+        self.meters.clear();
+    }
+
+    #[cfg(test)]
+    pub(super) fn set_download_limits(&mut self, limits: ResourceLimits) {
+        self.configuration.limits.limits.insert(
+            LimitKind::ArtifactDownloadBytes,
+            limits.limit(&LimitKind::ArtifactDownloadBytes),
+        );
+        self.meters.clear();
+    }
+
+    #[cfg(test)]
+    pub(super) fn set_rendezvous_limits(&mut self, limits: ResourceLimits) {
+        self.configuration.limits.limits.insert(
+            LimitKind::RendezvousAttempt,
+            limits.limit(&LimitKind::RendezvousAttempt),
+        );
+        self.meters.clear();
+    }
+}