1use std::{path::PathBuf, sync::Arc};
4
5use eredu_core::residency::CacheEvictionPolicy;
6
7use super::{CachePoolError, CachePoolLimits, CacheResidencyPool};
8
9#[derive(Debug, Clone, Default, Eq, PartialEq)]
11pub enum CacheResidencyPolicy {
12 #[default]
14 Device,
15 Paged(PagedCacheOptions),
17}
18
19#[derive(Debug, Clone, Default, Eq, PartialEq)]
21pub enum LiveCacheDiskPolicy {
22 #[default]
24 Disabled,
25 Enabled {
27 directory: PathBuf,
29 budget_bytes: u64,
31 queue_capacity: usize,
33 },
34}
35
36#[derive(Debug, Clone, Eq, PartialEq)]
38pub struct PagedCacheOptions {
39 block_size_tokens: i32,
40 device_budget_bytes: u64,
41 host_budget_bytes: u64,
42 recent_device_blocks: usize,
43 eviction_policy: CacheEvictionPolicy,
44 full_attention: bool,
45 retain_discarded_for_persistence: bool,
46 live_disk: LiveCacheDiskPolicy,
47 sample_process: bool,
48 pool: Option<Arc<CacheResidencyPool>>,
49}
50
51impl PagedCacheOptions {
52 pub fn new(
54 block_size_tokens: i32,
55 device_budget_bytes: u64,
56 host_budget_bytes: u64,
57 recent_device_blocks: usize,
58 ) -> Result<Self, CacheResidencyConfigurationError> {
59 if block_size_tokens <= 0 {
60 return Err(CacheResidencyConfigurationError::InvalidOptions(
61 "cache block size must be positive".into(),
62 ));
63 }
64 if device_budget_bytes == 0 {
65 return Err(CacheResidencyConfigurationError::InvalidOptions(
66 "paged cache device budget must be nonzero".into(),
67 ));
68 }
69 if recent_device_blocks == 0 {
70 return Err(CacheResidencyConfigurationError::InvalidOptions(
71 "paged cache must protect at least one recent device block".into(),
72 ));
73 }
74 Ok(Self {
75 block_size_tokens,
76 device_budget_bytes,
77 host_budget_bytes,
78 recent_device_blocks,
79 eviction_policy: CacheEvictionPolicy::LeastRecentlyUsed,
80 full_attention: false,
81 retain_discarded_for_persistence: false,
82 live_disk: LiveCacheDiskPolicy::Disabled,
83 sample_process: false,
84 pool: None,
85 })
86 }
87
88 pub const fn with_full_attention(mut self, enabled: bool) -> Self {
90 self.full_attention = enabled;
91 self
92 }
93
94 pub const fn with_persistence_retention(mut self, enabled: bool) -> Self {
96 self.retain_discarded_for_persistence = enabled;
97 self
98 }
99
100 pub const fn with_eviction_policy(mut self, policy: CacheEvictionPolicy) -> Self {
102 self.eviction_policy = policy;
103 self
104 }
105
106 pub fn with_live_disk(
108 mut self,
109 directory: impl Into<PathBuf>,
110 budget_bytes: u64,
111 queue_capacity: usize,
112 ) -> Result<Self, CacheResidencyConfigurationError> {
113 if budget_bytes == 0 {
114 return Err(CacheResidencyConfigurationError::InvalidOptions(
115 "live cache disk budget must be nonzero".into(),
116 ));
117 }
118 if queue_capacity == 0 {
119 return Err(CacheResidencyConfigurationError::InvalidOptions(
120 "live cache disk queue capacity must be nonzero".into(),
121 ));
122 }
123 let directory = directory.into();
124 if directory.as_os_str().is_empty() {
125 return Err(CacheResidencyConfigurationError::InvalidOptions(
126 "live cache disk directory must not be empty".into(),
127 ));
128 }
129 self.live_disk = LiveCacheDiskPolicy::Enabled {
130 directory,
131 budget_bytes,
132 queue_capacity,
133 };
134 Ok(self)
135 }
136
137 pub const fn with_process_sampling(mut self, enabled: bool) -> Self {
139 self.sample_process = enabled;
140 self
141 }
142
143 pub fn with_pool(
145 mut self,
146 pool: CacheResidencyPool,
147 ) -> Result<Self, CacheResidencyConfigurationError> {
148 if self.device_budget_bytes > pool.limits().device_bytes() {
149 return Err(CacheResidencyConfigurationError::InvalidOptions(format!(
150 "per-cache device budget {} exceeds cache pool budget {}",
151 self.device_budget_bytes,
152 pool.limits().device_bytes()
153 )));
154 }
155 if self.host_budget_bytes > pool.limits().host_bytes() {
156 return Err(CacheResidencyConfigurationError::InvalidOptions(format!(
157 "per-cache host budget {} exceeds cache pool budget {}",
158 self.host_budget_bytes,
159 pool.limits().host_bytes()
160 )));
161 }
162 if let LiveCacheDiskPolicy::Enabled { budget_bytes, .. } = &self.live_disk {
163 if *budget_bytes > pool.limits().disk_bytes() {
164 return Err(CacheResidencyConfigurationError::InvalidOptions(format!(
165 "per-cache disk budget {budget_bytes} exceeds cache pool budget {}",
166 pool.limits().disk_bytes()
167 )));
168 }
169 }
170 self.pool = Some(Arc::new(pool));
171 Ok(self)
172 }
173
174 pub const fn block_size_tokens(&self) -> i32 {
176 self.block_size_tokens
177 }
178
179 pub const fn device_budget_bytes(&self) -> u64 {
181 self.device_budget_bytes
182 }
183
184 pub const fn host_budget_bytes(&self) -> u64 {
186 self.host_budget_bytes
187 }
188
189 pub const fn recent_device_blocks(&self) -> usize {
191 self.recent_device_blocks
192 }
193
194 pub const fn eviction_policy(&self) -> CacheEvictionPolicy {
196 self.eviction_policy
197 }
198
199 pub const fn full_attention_enabled(&self) -> bool {
201 self.full_attention
202 }
203
204 pub const fn retains_discarded_for_persistence(&self) -> bool {
206 self.retain_discarded_for_persistence
207 }
208
209 pub const fn live_disk_policy(&self) -> &LiveCacheDiskPolicy {
211 &self.live_disk
212 }
213
214 pub const fn process_sampling_enabled(&self) -> bool {
216 self.sample_process
217 }
218
219 pub fn pool(&self) -> Option<&CacheResidencyPool> {
221 self.pool.as_deref()
222 }
223
224 pub fn create_pool(&self) -> Result<CacheResidencyPool, CacheResidencyConfigurationError> {
226 let disk_bytes = match self.live_disk_policy() {
227 LiveCacheDiskPolicy::Disabled => 0,
228 LiveCacheDiskPolicy::Enabled { budget_bytes, .. } => *budget_bytes,
229 };
230 let transfer_bytes = self
231 .device_budget_bytes()
232 .max(self.host_budget_bytes())
233 .saturating_mul(2)
234 .max(1);
235 Ok(CacheResidencyPool::new(CachePoolLimits::new(
236 self.device_budget_bytes(),
237 self.host_budget_bytes(),
238 transfer_bytes,
239 disk_bytes,
240 )?))
241 }
242}
243
244#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
246pub enum CacheResidencyConfigurationError {
247 #[error("invalid paged cache options: {0}")]
249 InvalidOptions(String),
250 #[error(transparent)]
252 Pool(#[from] CachePoolError),
253}
254
255#[cfg(test)]
256mod tests {
257 use super::*;
258
259 #[test]
260 fn finite_limits_and_live_disk_are_validated_without_a_backend() {
261 assert!(PagedCacheOptions::new(0, 1, 1, 1).is_err());
262 assert!(PagedCacheOptions::new(16, 0, 1, 1).is_err());
263 assert!(PagedCacheOptions::new(16, 1, 1, 0).is_err());
264 let options = PagedCacheOptions::new(16, 8, 0, 1)
265 .unwrap()
266 .with_live_disk("cache", 32, 2)
267 .unwrap();
268 assert_eq!(options.create_pool().unwrap().limits().disk_bytes(), 32);
269 }
270
271 #[test]
272 fn attached_pool_must_cover_every_enabled_residency_tier() {
273 let pool = CacheResidencyPool::new(CachePoolLimits::new(8, 4, 8, 0).unwrap());
274 let error = PagedCacheOptions::new(16, 16, 4, 1)
275 .unwrap()
276 .with_pool(pool)
277 .unwrap_err();
278 assert!(matches!(
279 error,
280 CacheResidencyConfigurationError::InvalidOptions(_)
281 ));
282 }
283}