clusterflux-public/crates/disasmer-core/src/limits.rs

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::TaskId;

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum LimitKind {
    ApiCall,
    Spawn,
    LogBytes,
    MetadataBytes,
    DebugReadBytes,
    UiEvent,
    RendezvousAttempt,
    ArtifactDownloadBytes,
    HostedFuel,
    HostedMemoryBytes,
    HostedWallClockMs,
    HostedStateBytes,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceLimits {
    pub limits: BTreeMap<LimitKind, u64>,
}

impl ResourceLimits {
    pub fn community_tier_defaults() -> Self {
        Self {
            limits: BTreeMap::from([
                (LimitKind::ApiCall, 10_000),
                (LimitKind::Spawn, 64),
                (LimitKind::LogBytes, 256 * 1024),
                (LimitKind::MetadataBytes, 512 * 1024),
                (LimitKind::DebugReadBytes, 128 * 1024),
                (LimitKind::UiEvent, 1_000),
                (LimitKind::RendezvousAttempt, 128),
                (LimitKind::ArtifactDownloadBytes, 256 * 1024 * 1024),
                (LimitKind::HostedFuel, 10_000_000),
                (LimitKind::HostedMemoryBytes, 64 * 1024 * 1024),
                (LimitKind::HostedWallClockMs, 30_000),
                (LimitKind::HostedStateBytes, 1024 * 1024),
            ]),
        }
    }

    pub fn limit(&self, kind: &LimitKind) -> u64 {
        *self.limits.get(kind).unwrap_or(&0)
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceMeter {
    used: BTreeMap<LimitKind, u64>,
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum LimitError {
    #[error(
        "resource limit exceeded for {kind:?}: requested {requested}, used {used}, limit {limit}"
    )]
    Exceeded {
        kind: LimitKind,
        requested: u64,
        used: u64,
        limit: u64,
    },
    #[error("task argument is too large: {size} bytes exceeds {limit} bytes")]
    LargeTaskArgument { size: u64, limit: u64 },
}

impl ResourceMeter {
    pub fn can_charge(
        &self,
        limits: &ResourceLimits,
        kind: LimitKind,
        amount: u64,
    ) -> Result<(), LimitError> {
        let used = self.used.get(&kind).copied().unwrap_or(0);
        let limit = limits.limit(&kind);
        if used.saturating_add(amount) > limit {
            return Err(LimitError::Exceeded {
                kind,
                requested: amount,
                used,
                limit,
            });
        }
        Ok(())
    }

    pub fn charge(
        &mut self,
        limits: &ResourceLimits,
        kind: LimitKind,
        amount: u64,
    ) -> Result<(), LimitError> {
        self.can_charge(limits, kind.clone(), amount)?;
        let used = self.used.get(&kind).copied().unwrap_or(0);
        self.used.insert(kind, used + amount);
        Ok(())
    }

    pub fn used(&self, kind: &LimitKind) -> u64 {
        self.used.get(kind).copied().unwrap_or(0)
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LogRecord {
    pub task: TaskId,
    pub bytes: Vec<u8>,
    pub truncated: bool,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LogBuffer {
    max_bytes: usize,
    used_bytes: usize,
    records: Vec<LogRecord>,
    backpressured: bool,
}

impl LogBuffer {
    pub fn new(max_bytes: usize) -> Self {
        Self {
            max_bytes,
            used_bytes: 0,
            records: Vec::new(),
            backpressured: false,
        }
    }

    pub fn push(&mut self, task: TaskId, bytes: impl AsRef<[u8]>) {
        let bytes = bytes.as_ref();
        let remaining = self.max_bytes.saturating_sub(self.used_bytes);
        let truncated = bytes.len() > remaining;
        let stored = bytes[..bytes.len().min(remaining)].to_vec();
        self.used_bytes += stored.len();
        if truncated {
            self.backpressured = true;
        }
        self.records.push(LogRecord {
            task,
            bytes: stored,
            truncated,
        });
    }

    pub fn records(&self) -> &[LogRecord] {
        &self.records
    }

    pub fn backpressured(&self) -> bool {
        self.backpressured
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum LargeArgumentPolicy {
    Allow,
    Warn,
    Reject,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TaskArgumentBudget {
    pub max_inline_bytes: u64,
    pub policy: LargeArgumentPolicy,
}

impl TaskArgumentBudget {
    pub fn validate(&self, size: u64) -> Result<Option<String>, LimitError> {
        if size <= self.max_inline_bytes {
            return Ok(None);
        }

        match self.policy {
            LargeArgumentPolicy::Allow => Ok(None),
            LargeArgumentPolicy::Warn => Ok(Some(format!(
                "task argument is {size} bytes; prefer SourceSnapshot, Blob, Artifact, or VFS handles"
            ))),
            LargeArgumentPolicy::Reject => Err(LimitError::LargeTaskArgument {
                size,
                limit: self.max_inline_bytes,
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn resource_meter_rejects_usage_before_work_starts() {
        let limits = ResourceLimits {
            limits: BTreeMap::from([(LimitKind::Spawn, 1)]),
        };
        let mut meter = ResourceMeter::default();

        meter.charge(&limits, LimitKind::Spawn, 1).unwrap();
        let error = meter.charge(&limits, LimitKind::Spawn, 1).unwrap_err();

        assert!(matches!(error, LimitError::Exceeded { .. }));
    }

    #[test]
    fn resource_meter_can_check_limits_without_consuming() {
        let limits = ResourceLimits {
            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 4)]),
        };
        let mut meter = ResourceMeter::default();

        meter
            .can_charge(&limits, LimitKind::ArtifactDownloadBytes, 4)
            .unwrap();
        assert_eq!(meter.used(&LimitKind::ArtifactDownloadBytes), 0);
        meter
            .charge(&limits, LimitKind::ArtifactDownloadBytes, 3)
            .unwrap();
        assert!(matches!(
            meter.can_charge(&limits, LimitKind::ArtifactDownloadBytes, 2),
            Err(LimitError::Exceeded { .. })
        ));
    }

    #[test]
    fn log_buffer_caps_backpressures_and_keeps_task_association() {
        let mut logs = LogBuffer::new(4);
        logs.push(TaskId::from("task-a"), b"abcdef");

        assert!(logs.backpressured());
        assert_eq!(logs.records()[0].task, TaskId::from("task-a"));
        assert_eq!(logs.records()[0].bytes, b"abcd");
        assert!(logs.records()[0].truncated);
    }

    #[test]
    fn large_task_arguments_are_rejected_or_warned() {
        let reject = TaskArgumentBudget {
            max_inline_bytes: 4,
            policy: LargeArgumentPolicy::Reject,
        };
        assert!(reject.validate(5).is_err());

        let warn = TaskArgumentBudget {
            max_inline_bytes: 4,
            policy: LargeArgumentPolicy::Warn,
        };
        assert!(warn.validate(5).unwrap().unwrap().contains("Artifact"));
    }
}