Public dry run dryrun-1714a9eedd5b

Source commit: 1714a9eedd5b1e30c6c9aceb7a999f2f695ba83c

Public tree identity: sha256:e5daffad23fa7c7f1822adccd3c3b4ee0bc7f1a45e0d321eb9a6fb8282b269db

crates/disasmer-core/src/limits.rs

@@ -0,0 +1,255 @@
+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"));
+    }
+}