RustProxy/tools/checks/memory_contracts.py

#!/usr/bin/env python3
"""Static regression checks for bounded-memory/background-worker contracts.

This is intentionally Rust-toolchain independent. It does not replace runtime RSS/
heap profiling; it prevents the source-level bounds audited for this release from
silently disappearing before compiler-enabled acceptance runs.
"""
from __future__ import annotations

from pathlib import Path
from rust_source import rust_source

ROOT = Path(__file__).resolve().parents[2]


def read(path: str) -> str:
    target = ROOT / path
    return rust_source(ROOT, path) if target.suffix == ".rs" else target.read_text(encoding="utf-8")


def require(path: str, *markers: str) -> None:
    text = read(path)
    for marker in markers:
        assert marker in text, f"{path}: missing memory contract: {marker}"


def forbid(path: str, *markers: str) -> None:
    text = read(path)
    for marker in markers:
        assert marker not in text, f"{path}: forbidden memory pattern returned: {marker}"


# Transform path: hard aggregate/decode/concurrency budgets and no JSON DOM tree.
require(
    "src/transform/mod.rs",
    "DEFAULT_MEMORY_BUDGET_BYTES: usize = 128 * 1024 * 1024",
    "pub memory_budget_bytes: usize",
    "pub max_inflight: usize",
    "try_reserve_transform_memory(required_memory, settings.memory_budget_bytes)",
    "try_reserve_transform_slot(settings.max_inflight)",
    "TRANSFORM_MEMORY_IN_USE_BYTES: AtomicUsize",
    "TRANSFORM_INFLIGHT: AtomicUsize",
    "pub struct ResourceUsage",
    "pub fn resource_usage(settings: &Settings)",
    "memory_reserved_bytes: TRANSFORM_MEMORY_IN_USE_BYTES.load",
    "Limited::new(body, settings.max_body_bytes)",
    ".try_reserve_exact(read)",
    '"memory_pressure"',
    '"concurrency_pressure"',
    "aggregation_is_bounded_even_when_content_length_is_false",
    "transform_memory_reservation_obeys_runtime_budget",
    "transform_inflight_reservation_obeys_runtime_limit",
)
require("src/transform/json.rs", "IgnoredAny::deserialize", "Vec::with_capacity(body.len())")
forbid("src/transform/json.rs", "use serde_json::Value", "from_slice::<Value>")
require("src/transform/css.rs", "let mut last_char: Option<char> = None")
require("src/transform/html.rs", "enum RawElement", "fn raw_tag(tag: &str)")

# Cache background work: all previously high-retention paths have hard bounds.
require(
    "src/cache/model.rs",
    "COLD_DELETE_FALLBACKS: usize = 2",
    "COLD_CLEANUP_SLOTS: usize = 1",
    "BACKGROUND_REVALIDATION_LIMIT: usize = 8",
    "pub cold_limit_bytes: u64",
    "pub cold_limit_entries: usize",
    "pub hot_auto_trim_enabled: bool",
    "pub hot_auto_trim_after_seconds: u64",
    "pub hot_auto_trim_target_percent: u8",
)
require(
    "src/admin/mod.rs",
    'payload["transform_resources"]',
    'payload["memory_accounting"]',
    '"non_cache_rss_includes_transform": true',
)
require(
    "src/cache/cold.rs",
    "BinaryHeap::with_capacity(requested_limit.saturating_add(1))",
    "self.cold_delete_fallbacks.clone().try_acquire_owned()",
    "self.cleanup_slots.clone().try_acquire_owned()",
    "std::mem::take(&mut delete)",
)
require(
    "src/cache/tickets.rs",
    "active.len() >= BACKGROUND_REVALIDATION_LIMIT",
)
require(
    "src/cache/collapse.rs",
    ".saturating_mul(64)",
    ".clamp(64, 256)",
    "settings.proactive_refresh_max_inflight as u64",
)
require(
    "src/cache/lifecycle.rs",
    "mpsc::channel(COLD_QUEUE)",
    "cold_delete_fallbacks: Arc::new(Semaphore::new(COLD_DELETE_FALLBACKS))",
    "cleanup_slots: Arc::new(Semaphore::new(COLD_CLEANUP_SLOTS))",
    "cache.trim_hot_idle();",
    "fn trim_hot_idle(&self)",
)

# COLD bodies stay disk-backed and the long-lived index must not retain full
# response header/validator allocations for every object.
cache_model = read("src/cache/model/cache.rs")
cold_entry = cache_model.split("struct ColdEntry {", 1)[1].split("#[derive(Default)]", 1)[0]
assert "headers: HeaderMap" not in cold_entry, "COLD index retains full response headers"
assert "origin_etag: Option<HeaderValue>" not in cold_entry, "COLD index retains ETag HeaderValue"
assert "origin_last_modified: Option<HeaderValue>" not in cold_entry, "COLD index retains Last-Modified HeaderValue"
assert "body_sha256: Key" in cold_entry, "COLD digest must use the compact 32-byte representation"
assert "has_validator: bool" in cold_entry, "COLD index must retain validator presence without full values"
require(
    "src/cache/model/cold_metadata.rs",
    "tokio::fs::read(&entry.meta_path).await",
    "pairs_to_headers(&file.headers)",
    "body_sha256 != entry.body_sha256",
)

# Rate-limit client state stores a vhost allocation once and reclaims expired
# fixed-window buckets without waiting for the global cardinality ceiling.
require(
    "src/rate_limit.rs",
    "buckets: HashMap<String, HashMap<IpAddr, Bucket>>",
    "fn sweep_expired(&mut self, now: Instant, force: bool)",
    "now.duration_since(bucket.started) < Duration::from_secs(bucket.window_seconds)",
    "state.sweep_expired(now, false)",
    "state.sweep_expired(now, true)",
    "clients.shrink_to_fit()",
    "self.buckets.shrink_to_fit()",
    "pub fn maintenance(&self)",
    "MAX_BUCKETS: usize = 65_536",
)
require(
    "src/runtime/constructors.rs",
    "app.rate_limiter.maintenance()",
    "spawn_blocking(crate::diagnostics::trim_allocator)",
)

# Compression remains response-streaming; no complete-body collection is added.
require(
    "src/compression.rs",
    "No full-body aggregation is introduced.",
    "CompressionLevel::Fastest",
    "Compression::new(service)",
)

# One JSONL buffer is reused across local file and Logstash sends (Loki-only
# sinks should not allocate the JSONL line). Both output queues remain bounded.
logging_sink = read("src/logging/runtime.rs")
assert logging_sink.count("serde_json::to_vec(&payload)") == 1
assert "if config.file.enabled || config.logstash.enabled" in logging_sink
assert "opened.write_all(&jsonl_line)" in logging_sink
assert "state.stream.write_all(&jsonl_line)" in logging_sink
assert "PER_VHOST_OPEN_FILES: usize = 64" in logging_sink

# Stored Logs file discovery must also be cardinality bounded before sort:
# rotated files and per-vhost directories may contain arbitrary historical files.
require("src/admin/file_logs.rs", "newest_log_files(std::fs::read_dir(&directory)", ", 256);")
require("src/admin/file_logs_io.rs", "BinaryHeap::with_capacity(limit)", "newest_log_files(std::fs::read_dir(parent)", "), 64);")

# A per-key invalidation generation is useful only inside the current PURGE
# publication epoch; retaining each purged URL after eviction is an unbounded
# lifetime-memory leak on high-cardinality cache invalidation workloads.
invalidation = read("src/cache/invalidation/invalidate.rs")
start = invalidation.index("InvalidationOperation::Purge =>")
barrier = invalidation.index("store.generation = store.generation.wrapping_add(1);", start)
reset = invalidation.index("store.key_generations.clear();", barrier)
assert "store.key_generations.capacity() > 256" in invalidation[barrier:reset]
matcher = invalidation.index("let affected = self.purge_matcher(&matcher);", reset)
assert start < barrier < reset < matcher
require("src/cache/tests/v022_cache_tests/part_02.rs", "many_distinct_selective_purges_do_not_retain_tombstones_forever")

# Repeated hot-path analytics labels should use borrowed lookups instead of
# allocating a String for every request in an already populated top table.
require("src/traffic.rs", "fn truncate(value: &str, max_bytes: usize) -> &str", "self.values.get_mut(value)", "self.values.insert(value.to_owned(), weight)")
forbid("src/web_server/serve.rs", "Body::from(listing.clone())")
require("src/web_server/serve.rs", "entries.retain(|entry|", "let listing_bytes = listing.len();", "Body::from(listing)")
require("src/logging/runtime.rs", "String::from_utf8_lossy(&jsonl_line[..jsonl_line.len() - 1])")

# Other long-lived/background facilities: queue/task/worker cardinality is bounded
# or owned by a cancellation-aware supervisor/TaskTracker.
require(
    "src/logging.rs",
    "HISTORY_CAPACITY: usize = 4096",
    "LIVE_STREAM_CAPACITY: usize = 4096",
    "mpsc::channel(SINK_QUEUE)",
    "self.sink.try_send(entry)",
    "pub capture_user_agent: bool",
    "pub capture_country: bool",
)
require(
    "src/traffic.rs",
    "QUEUE_CAPACITY: usize = 2048",
    "TOP_CAPACITY: usize = 32",
    "MAX_VHOSTS: usize = 64",
    "self.tx.try_send(sample)",
    "client_ips: TopTable",
)
require(
    "src/diagnostics.rs",
    "tokio::task::spawn_blocking(move || system_stats::sample(previous)).await",
    "stop.cancelled()",
    "pub fn trim_allocator() -> bool",
)
require(
    "src/runtime.rs",
    "Duration::from_secs(60)",
    "spawn_blocking(crate::diagnostics::trim_allocator)",
)
require(
    "src/upstream.rs",
    "Semaphore::new(16)",
    ".try_acquire_owned()",
    "tokio::time::timeout(DNS_TIMEOUT, lookup)",
)
require(
    "src/upstream/client.rs",
    "pool_recycle_requests",
    "pool_recycle_seconds",
    "http1_max_buf_size(keepalive.http1_max_buffer_bytes)",
    "ArcSwap<RawProxyClient>",
    "self.current.store(Arc::new(client))",
    '"recycled upstream client pool"',
)
require(
    "src/balance.rs",
    "if checks.len() >= 16 { break; }",
    "checks.abort_all()",
)
require(
    "src/stream.rs",
    "Semaphore::new(host.max_connections)",
    "mpsc::channel(64)",
    "JoinSet::new()",
)
require(
    "src/server.rs",
    "Semaphore::new(max_connections)",
    "let mut tasks = JoinSet::new()",
)
require(
    "src/runtime.rs",
    "tokio_util::task::TaskTracker::new()",
    "update_slots: Arc::new(Semaphore::new(8))",
    "inflight: Arc::new(Semaphore::new(max_inflight))",
)
require("src/acme.rs", "slot: Arc::new(Semaphore::new(1))", ".try_acquire_owned()")
require("src/docker.rs", "shutdown.cancelled()", "Duration::from_secs(5)")
require("src/update_check.rs", "shutdown.cancelled()", "tokio::time::sleep(interval)")

print("PASS: bounded-memory transform/cache contracts and background-worker limits are present.")