use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct ResourcesConfig {
#[serde(default)]
pub scan_threads: usize,
#[serde(default)]
pub embed_threads: usize,
#[serde(default)]
pub max_concurrent_documents: usize,
#[serde(default = "ResourcesConfig::default_embed_batch_size")]
#[schemars(range(min = 1))]
pub embed_batch_size: usize,
#[serde(default)]
pub max_footprint_mb: MaxFootprint,
#[serde(default = "ResourcesConfig::default_max_map_cache_mb")]
pub max_map_cache_mb: usize,
#[serde(default)]
pub document_models: DocumentModelProfile,
#[serde(default)]
pub onnx_provider: OnnxProvider,
}
impl ResourcesConfig {
fn default_embed_batch_size() -> usize {
32
}
fn default_max_map_cache_mb() -> usize {
256
}
pub fn effective_embed_threads(&self, deprecated_alias: usize) -> usize {
if self.embed_threads != 0 {
self.embed_threads
} else {
deprecated_alias
}
}
}
impl Default for ResourcesConfig {
fn default() -> Self {
Self {
scan_threads: 0,
embed_threads: 0,
max_concurrent_documents: 0,
embed_batch_size: Self::default_embed_batch_size(),
max_footprint_mb: MaxFootprint::default(),
max_map_cache_mb: Self::default_max_map_cache_mb(),
document_models: DocumentModelProfile::default(),
onnx_provider: OnnxProvider::default(),
}
}
}
const BYTES_PER_MB: u64 = 1024 * 1024;
const AUTO_ENFORCED_NUMERATOR: u64 = 3;
const AUTO_ENFORCED_DENOMINATOR: u64 = 4;
const AUTO_ADVISORY_DENOMINATOR: u64 = 2;
const AUTO_FLOOR_BYTES: u64 = 512 * BYTES_PER_MB;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum FootprintKeyword {
Off,
Auto,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(untagged)]
pub enum MaxFootprint {
Keyword(FootprintKeyword),
Mebibytes(usize),
}
impl Default for MaxFootprint {
fn default() -> Self {
MaxFootprint::Mebibytes(0)
}
}
impl MaxFootprint {
pub fn is_off(self) -> bool {
matches!(self, MaxFootprint::Keyword(FootprintKeyword::Off))
}
pub fn is_auto(self) -> bool {
matches!(
self,
MaxFootprint::Keyword(FootprintKeyword::Auto) | MaxFootprint::Mebibytes(0)
)
}
pub fn explicit_bytes(self) -> Option<u64> {
match self {
MaxFootprint::Mebibytes(mb) if mb > 0 => Some((mb as u64).saturating_mul(BYTES_PER_MB)),
_ => None,
}
}
pub fn auto_ceiling_bytes(limit_bytes: Option<u64>, enforced: bool) -> Option<u64> {
let limit = limit_bytes?;
let budget = if enforced {
limit / AUTO_ENFORCED_DENOMINATOR * AUTO_ENFORCED_NUMERATOR
} else {
limit / AUTO_ADVISORY_DENOMINATOR
};
Some(budget.max(AUTO_FLOOR_BYTES))
}
pub fn resolve_bytes(self) -> Option<u64> {
if self.is_off() {
return None;
}
if let Some(bytes) = self.explicit_bytes() {
return Some(bytes);
}
let reading = crate::sysres::memory_reading()?;
Self::auto_ceiling_bytes(reading.limit_bytes, reading.limit_is_enforced)
}
pub fn resolve_mb(self) -> usize {
match self.resolve_bytes() {
Some(bytes) => (bytes / BYTES_PER_MB) as usize,
None => 0,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
#[serde(rename_all = "lowercase")]
pub enum OnnxProvider {
#[default]
Cpu,
Auto,
#[serde(alias = "core_ml")]
CoreMl,
Cuda,
#[serde(alias = "tensor_rt")]
TensorRt,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
#[serde(rename_all = "snake_case")]
pub enum DocumentModelProfile {
#[default]
Full,
CodeOnly,
#[serde(rename = "none")]
None_,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_resources_config_has_expected_field_values() {
let cfg = ResourcesConfig::default();
assert_eq!(cfg.scan_threads, 0);
assert_eq!(cfg.embed_threads, 0);
assert_eq!(cfg.max_concurrent_documents, 0);
assert_eq!(cfg.embed_batch_size, 32);
assert_eq!(cfg.max_footprint_mb, MaxFootprint::Mebibytes(0));
assert_eq!(cfg.max_map_cache_mb, 256);
assert_eq!(cfg.document_models, DocumentModelProfile::Full);
}
#[test]
fn document_model_profile_defaults_to_full() {
assert_eq!(DocumentModelProfile::default(), DocumentModelProfile::Full);
}
#[test]
fn resources_toml_roundtrips_embed_batch_size_override() {
let cfg: ResourcesConfig = toml::from_str("embed_batch_size = 8\n").expect("parse [resources] body");
assert_eq!(cfg.embed_batch_size, 8);
assert_eq!(cfg.scan_threads, 0);
assert_eq!(cfg.document_models, DocumentModelProfile::Full);
}
#[test]
fn resources_empty_toml_falls_back_to_all_defaults() {
let cfg: ResourcesConfig = toml::from_str("").expect("empty [resources] body");
assert_eq!(cfg.embed_batch_size, 32);
assert_eq!(cfg.scan_threads, 0);
assert_eq!(cfg.embed_threads, 0);
assert_eq!(cfg.max_concurrent_documents, 0);
assert_eq!(cfg.max_footprint_mb, MaxFootprint::Mebibytes(0));
assert_eq!(cfg.max_map_cache_mb, 256);
assert_eq!(cfg.document_models, DocumentModelProfile::Full);
}
#[test]
fn max_map_cache_mb_accepts_zero_as_unbounded() {
let cfg: ResourcesConfig = toml::from_str("max_map_cache_mb = 0\n").expect("parse [resources] body");
assert_eq!(cfg.max_map_cache_mb, 0);
}
#[test]
fn document_model_profile_none_serializes_as_none_string() {
let profile = DocumentModelProfile::None_;
let json = serde_json::to_string(&profile).expect("serialize");
assert_eq!(json, "\"none\"");
let back: DocumentModelProfile = serde_json::from_str("\"none\"").expect("deserialize");
assert_eq!(back, DocumentModelProfile::None_);
}
#[test]
fn document_model_profile_code_only_uses_snake_case() {
let json = serde_json::to_string(&DocumentModelProfile::CodeOnly).expect("serialize");
assert_eq!(json, "\"code_only\"");
}
const MB: u64 = 1024 * 1024;
#[test]
fn max_footprint_accepts_all_three_spellings_from_toml() {
let auto: ResourcesConfig = toml::from_str("max_footprint_mb = 0\n").expect("parse 0");
assert_eq!(auto.max_footprint_mb, MaxFootprint::Mebibytes(0));
assert!(auto.max_footprint_mb.is_auto());
let explicit: ResourcesConfig = toml::from_str("max_footprint_mb = 512\n").expect("parse 512");
assert_eq!(explicit.max_footprint_mb, MaxFootprint::Mebibytes(512));
assert!(!explicit.max_footprint_mb.is_auto());
let off: ResourcesConfig = toml::from_str("max_footprint_mb = \"off\"\n").expect("parse off");
assert_eq!(off.max_footprint_mb, MaxFootprint::Keyword(FootprintKeyword::Off));
assert!(off.max_footprint_mb.is_off());
let named_auto: ResourcesConfig = toml::from_str("max_footprint_mb = \"auto\"\n").expect("parse auto");
assert_eq!(
named_auto.max_footprint_mb,
MaxFootprint::Keyword(FootprintKeyword::Auto)
);
assert!(named_auto.max_footprint_mb.is_auto());
}
#[test]
fn max_footprint_rejects_an_unknown_keyword() {
let bad = toml::from_str::<ResourcesConfig>("max_footprint_mb = \"disabled\"\n");
assert!(bad.is_err(), "only `off` and `auto` are accepted keywords");
}
#[test]
fn max_footprint_serialises_auto_back_as_the_integer_zero() {
let json = serde_json::to_string(&MaxFootprint::default()).expect("serialize");
assert_eq!(json, "0");
assert_eq!(
serde_json::to_string(&MaxFootprint::Mebibytes(256)).expect("serialize"),
"256"
);
assert_eq!(
serde_json::to_string(&MaxFootprint::Keyword(FootprintKeyword::Off)).expect("serialize"),
"\"off\""
);
}
#[test]
fn explicit_bytes_is_only_set_for_a_positive_mebibyte_count() {
assert_eq!(MaxFootprint::Mebibytes(512).explicit_bytes(), Some(512 * MB));
assert_eq!(MaxFootprint::Mebibytes(0).explicit_bytes(), None);
assert_eq!(MaxFootprint::Keyword(FootprintKeyword::Off).explicit_bytes(), None);
assert_eq!(MaxFootprint::Keyword(FootprintKeyword::Auto).explicit_bytes(), None);
}
#[test]
fn auto_ceiling_takes_three_quarters_of_an_enforced_cgroup_limit() {
assert_eq!(MaxFootprint::auto_ceiling_bytes(Some(2048 * MB), true), Some(1536 * MB));
assert_eq!(MaxFootprint::auto_ceiling_bytes(Some(8192 * MB), true), Some(6144 * MB));
}
#[test]
fn auto_ceiling_takes_half_of_a_machine_wide_limit() {
assert_eq!(
MaxFootprint::auto_ceiling_bytes(Some(16 * 1024 * MB), false),
Some(8 * 1024 * MB)
);
}
#[test]
fn auto_ceiling_floor_applies_to_a_small_container() {
assert_eq!(MaxFootprint::auto_ceiling_bytes(Some(256 * MB), true), Some(512 * MB));
assert_eq!(MaxFootprint::auto_ceiling_bytes(Some(1024 * MB), false), Some(512 * MB));
assert_eq!(MaxFootprint::auto_ceiling_bytes(Some(800 * MB), false), Some(512 * MB));
}
#[test]
fn auto_ceiling_is_absent_when_no_limit_can_be_detected() {
assert_eq!(MaxFootprint::auto_ceiling_bytes(None, true), None);
assert_eq!(MaxFootprint::auto_ceiling_bytes(None, false), None);
}
#[test]
fn auto_ceiling_never_overflows_on_a_sentinel_sized_limit() {
let ceiling = MaxFootprint::auto_ceiling_bytes(Some(u64::MAX), true).expect("a ceiling");
assert!((512 * MB..u64::MAX).contains(&ceiling));
}
#[test]
fn resolve_bytes_honours_off_and_an_explicit_ceiling_without_sampling() {
assert_eq!(MaxFootprint::Keyword(FootprintKeyword::Off).resolve_bytes(), None);
assert_eq!(MaxFootprint::Keyword(FootprintKeyword::Off).resolve_mb(), 0);
assert_eq!(MaxFootprint::Mebibytes(384).resolve_bytes(), Some(384 * MB));
assert_eq!(MaxFootprint::Mebibytes(384).resolve_mb(), 384);
}
#[test]
fn resolve_mb_auto_yields_a_ceiling_on_a_supported_platform() {
let resolved = MaxFootprint::default().resolve_mb();
if cfg!(any(target_os = "linux", target_os = "macos", windows)) {
assert!(
resolved >= 512,
"auto must clear the 512 MiB floor on a supported platform, got {resolved}"
);
} else {
assert_eq!(resolved, 0, "an unsupported platform reports no ceiling");
}
}
#[test]
fn onnx_provider_defaults_to_cpu_and_parses_every_spelling() {
let absent: ResourcesConfig = toml::from_str("scan_threads = 2\n").expect("pre-existing config");
assert_eq!(absent.onnx_provider, OnnxProvider::Cpu);
for (text, expected) in [
("cpu", OnnxProvider::Cpu),
("auto", OnnxProvider::Auto),
("coreml", OnnxProvider::CoreMl),
("cuda", OnnxProvider::Cuda),
("tensorrt", OnnxProvider::TensorRt),
] {
let cfg: ResourcesConfig =
toml::from_str(&format!("onnx_provider = \"{text}\"\n")).expect("provider spelling");
assert_eq!(cfg.onnx_provider, expected, "{text}");
}
assert!(toml::from_str::<ResourcesConfig>("onnx_provider = \"metal\"\n").is_err());
}
#[test]
fn effective_embed_threads_prefers_resources_then_deprecated_alias() {
let cfg = ResourcesConfig {
embed_threads: 4,
..ResourcesConfig::default()
};
assert_eq!(cfg.effective_embed_threads(8), 4);
let cfg = ResourcesConfig::default();
assert_eq!(cfg.effective_embed_threads(8), 8);
assert_eq!(cfg.effective_embed_threads(0), 0);
}
}