kvbm_engine/offload/policy.rs
1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Policy trait and built-in implementations for offload filtering.
5//!
6//! Policies determine which blocks should be offloaded. They are evaluated
7//! as filters - blocks that fail any filter are removed from the transfer.
8//!
9//! # Performance Optimization
10//!
11//! This module uses `Either<Ready, BoxFuture>` instead of `#[async_trait]` to
12//! avoid heap allocations for synchronous policies. Policies that only perform
13//! local, synchronous operations (like `PresenceFilter`, `PassAllPolicy`) return
14//! `Either::Left(ready(...))` which requires zero heap allocation. Policies that
15//! need actual async operations return `Either::Right(Box::pin(...))`.
16//!
17//! # Built-in Policies
18//!
19//! - `PresenceFilter<Src, Dst>`: Skip blocks already present in destination tier
20//! - `PresenceAndLFUFilter<Src, Dst>`: Presence check + LFU count threshold
21//! - `PassAllPolicy`: No filtering (pass all blocks)
22//! - `AllOfPolicy`: Composite AND policy
23//! - `AnyOfPolicy`: Composite OR policy
24
25use std::future::{Future, Ready, ready};
26use std::marker::PhantomData;
27use std::pin::Pin;
28use std::sync::Arc;
29
30use anyhow::Result;
31use futures::future::Either;
32use kvbm_config::{PolicyType, TierOffloadConfig};
33
34use crate::{BlockId, SequenceHash};
35use kvbm_logical::blocks::{BlockMetadata, BlockRegistry, ImmutableBlock};
36
37use super::pending::{PendingCheck, PendingTracker};
38use crate::object::{ObjectBlockOps, ObjectLockManager};
39
40/// Boxed future type for async policy evaluation.
41pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
42
43/// Future type for single-block policy evaluation.
44///
45/// - `Left(Ready<...>)`: Synchronous result, zero heap allocation
46/// - `Right(BoxFuture<...>)`: Async result, requires heap allocation
47pub type PolicyFuture<'a> = Either<Ready<Result<bool>>, BoxFuture<'a, Result<bool>>>;
48
49/// Future type for batch policy evaluation.
50///
51/// - `Left(Ready<...>)`: Synchronous result, zero heap allocation
52/// - `Right(BoxFuture<...>)`: Async result, requires heap allocation
53pub type PolicyBatchFuture<'a> = Either<Ready<Result<Vec<bool>>>, BoxFuture<'a, Result<Vec<bool>>>>;
54
55/// Create a synchronous policy result (zero allocation).
56#[inline]
57pub fn sync_result(result: Result<bool>) -> PolicyFuture<'static> {
58 Either::Left(ready(result))
59}
60
61/// Create a synchronous batch policy result (zero allocation).
62#[inline]
63pub fn sync_batch_result(result: Result<Vec<bool>>) -> PolicyBatchFuture<'static> {
64 Either::Left(ready(result))
65}
66
67/// Create an async policy result (boxes the future).
68#[inline]
69pub fn async_result<'a, F>(future: F) -> PolicyFuture<'a>
70where
71 F: Future<Output = Result<bool>> + Send + 'a,
72{
73 Either::Right(Box::pin(future))
74}
75
76/// Create an async batch policy result (boxes the future).
77#[inline]
78pub fn async_batch_result<'a, F>(future: F) -> PolicyBatchFuture<'a>
79where
80 F: Future<Output = Result<Vec<bool>>> + Send + 'a,
81{
82 Either::Right(Box::pin(future))
83}
84
85// ============================================================================
86// Presence Checker Trait
87// ============================================================================
88
89/// Async presence checker for object storage or other external destinations.
90///
91/// This trait abstracts presence checking for destinations that require async
92/// operations (like S3, caching services). Unlike `BlockRegistry::check_presence`
93/// which is synchronous, this is designed for remote/external destinations.
94///
95/// # Implementations
96///
97/// - `S3PresenceChecker`: Wraps `ObjectBlockOps::has_blocks()` for S3/object storage
98/// - Future: `CachedPresenceChecker` - local bloom filter / LRU cache layer
99/// - Future: `DistributedCacheChecker` - remote caching service
100pub trait PresenceChecker: Send + Sync {
101 /// Check if blocks exist at the destination.
102 ///
103 /// Returns a vector of (SequenceHash, exists: bool) pairs.
104 fn check_presence(
105 &self,
106 keys: Vec<SequenceHash>,
107 ) -> BoxFuture<'static, Vec<(SequenceHash, bool)>>;
108}
109
110/// S3/Object storage presence checker.
111///
112/// Wraps `ObjectBlockOps::has_blocks()` and converts `Option<usize>` → `bool`.
113/// This is the default presence checker for G2→G4 (object storage) pipelines.
114///
115/// # Example
116/// ```ignore
117/// let object_ops: Arc<dyn ObjectBlockOps> = ...;
118/// let checker = S3PresenceChecker::new(object_ops);
119/// let results = checker.check_presence(keys).await;
120/// ```
121pub struct S3PresenceChecker {
122 object_ops: Arc<dyn ObjectBlockOps>,
123}
124
125impl S3PresenceChecker {
126 /// Create a new S3 presence checker wrapping the given object operations.
127 pub fn new(object_ops: Arc<dyn ObjectBlockOps>) -> Self {
128 Self { object_ops }
129 }
130}
131
132impl PresenceChecker for S3PresenceChecker {
133 fn check_presence(
134 &self,
135 keys: Vec<SequenceHash>,
136 ) -> BoxFuture<'static, Vec<(SequenceHash, bool)>> {
137 let future = self.object_ops.has_blocks(keys);
138 Box::pin(async move {
139 let results = future.await;
140 // Convert Option<usize> (size) → bool (exists)
141 results
142 .into_iter()
143 .map(|(hash, size_opt)| (hash, size_opt.is_some()))
144 .collect()
145 })
146 }
147}
148
149// ============================================================================
150// Evaluation Context
151// ============================================================================
152
153/// Context provided to policies for block evaluation.
154#[derive(Debug)]
155pub struct EvalContext<T: BlockMetadata> {
156 /// Block ID
157 pub block_id: BlockId,
158 /// Sequence hash for this block
159 pub sequence_hash: SequenceHash,
160 /// Optional strong reference to the block.
161 /// - Some: Strong blocks (held during evaluation)
162 /// - None: Weak blocks (deferred upgrade)
163 pub block: Option<ImmutableBlock<T>>,
164}
165
166impl<T: BlockMetadata> EvalContext<T> {
167 /// Create a new evaluation context from a strong block reference.
168 pub fn new(block: ImmutableBlock<T>) -> Self {
169 Self {
170 block_id: block.block_id(),
171 sequence_hash: block.sequence_hash(),
172 block: Some(block),
173 }
174 }
175
176 /// Create a context for weak block evaluation (deferred upgrade).
177 ///
178 /// Used when evaluating weak blocks - we have the metadata
179 /// but defer the actual upgrade until just before transfer.
180 pub fn from_weak(block_id: BlockId, sequence_hash: SequenceHash) -> Self {
181 Self {
182 block_id,
183 sequence_hash,
184 block: None,
185 }
186 }
187
188 /// Create a context for external block evaluation.
189 ///
190 /// Used when evaluating external blocks (e.g., G1 from vLLM) - we have
191 /// the block_id and sequence_hash but no ImmutableBlock reference.
192 pub fn from_external(block_id: BlockId, sequence_hash: SequenceHash) -> Self {
193 Self {
194 block_id,
195 sequence_hash,
196 block: None,
197 }
198 }
199}
200
201/// Trait for offload policies that filter blocks.
202///
203/// Policies are evaluated as a chain - a block must pass ALL policies to proceed.
204/// Each policy receives an `EvalContext` with block information and returns
205/// `Ok(true)` to pass or `Ok(false)` to filter out.
206///
207/// # Performance
208///
209/// This trait uses `Either<Ready, BoxFuture>` instead of `#[async_trait]` to
210/// avoid heap allocations for synchronous policies. Implement using:
211/// - `sync_result(Ok(true))` for synchronous policies (zero allocation)
212/// - `async_result(async { ... })` for async policies (boxes the future)
213///
214/// # Batch Evaluation
215///
216/// The `evaluate_batch` method provides a default implementation that calls
217/// `evaluate` for each block. Override for efficiency when the policy can
218/// benefit from batching (e.g., batch registry lookups).
219pub trait OffloadPolicy<T: BlockMetadata>: Send + Sync {
220 /// Unique name for this policy (for logging/debugging).
221 fn name(&self) -> &str;
222
223 /// Evaluate whether a block should be offloaded.
224 ///
225 /// Returns:
226 /// - `Ok(true)`: Block passes this filter, continue to next policy
227 /// - `Ok(false)`: Block filtered out, remove from transfer
228 /// - `Err(_)`: Fatal error, fail the entire transfer
229 fn evaluate<'a>(&'a self, ctx: &'a EvalContext<T>) -> PolicyFuture<'a>;
230
231 /// Batch evaluate multiple blocks.
232 ///
233 /// Default implementation calls `evaluate` for each block.
234 /// Override for efficiency when batching is beneficial.
235 fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<T>]) -> PolicyBatchFuture<'a> {
236 // Default: sequential evaluation
237 let contexts_clone: Vec<_> = contexts.iter().collect();
238 async_batch_result(async move {
239 let mut results = Vec::with_capacity(contexts_clone.len());
240 for ctx in contexts_clone {
241 // This calls the sync or async evaluate
242 let result = match self.evaluate(ctx) {
243 Either::Left(ready) => ready.await,
244 Either::Right(boxed) => boxed.await,
245 };
246 results.push(result?);
247 }
248 Ok(results)
249 })
250 }
251}
252
253/// G1→G2 filter: skip blocks already present in destination tier.
254///
255/// Uses `BlockRegistry::check_presence` to determine if a block exists
256/// in the destination tier without acquiring a full block reference.
257/// This is efficient because it only checks the registry metadata.
258///
259/// # Duplicate Prevention
260///
261/// When a `PendingTracker` is configured, this filter also checks for blocks
262/// that are currently in-flight through the pipeline. This prevents duplicate
263/// transfers when overlapping sequences are enqueued at roughly the same time.
264///
265/// # Performance
266///
267/// This policy is fully synchronous and returns `Either::Left(Ready)`,
268/// avoiding any heap allocation per evaluation.
269///
270/// # Example
271/// ```ignore
272/// let tracker = Arc::new(PendingTracker::new());
273/// let filter = PresenceFilter::<G1, G2>::new(registry.clone())
274/// .with_pending_tracker(tracker);
275/// // Blocks already in G2 OR in-flight will be filtered out
276/// ```
277pub struct PresenceFilter<Src: BlockMetadata, Dst: BlockMetadata> {
278 registry: Arc<BlockRegistry>,
279 /// Optional tracker for pending (in-flight) transfers.
280 /// When set, blocks that are already being transferred will be filtered out.
281 pending_tracker: Option<Arc<PendingTracker>>,
282 _marker: PhantomData<(Src, Dst)>,
283}
284
285impl<Src: BlockMetadata, Dst: BlockMetadata> PresenceFilter<Src, Dst> {
286 /// Create a new presence filter without pending tracking.
287 pub fn new(registry: Arc<BlockRegistry>) -> Self {
288 Self {
289 registry,
290 pending_tracker: None,
291 _marker: PhantomData,
292 }
293 }
294
295 /// Add a pending tracker for duplicate prevention.
296 ///
297 /// When set, blocks that are currently in-flight (passed policy but not
298 /// yet registered in destination) will be filtered out.
299 pub fn with_pending_tracker(mut self, tracker: Arc<PendingTracker>) -> Self {
300 self.pending_tracker = Some(tracker);
301 self
302 }
303
304 /// Get a reference to the pending tracker if configured.
305 pub fn pending_tracker(&self) -> Option<&Arc<PendingTracker>> {
306 self.pending_tracker.as_ref()
307 }
308}
309
310impl<Src: BlockMetadata, Dst: BlockMetadata> OffloadPolicy<Src> for PresenceFilter<Src, Dst> {
311 fn name(&self) -> &str {
312 "PresenceFilter"
313 }
314
315 fn evaluate<'a>(&'a self, ctx: &'a EvalContext<Src>) -> PolicyFuture<'a> {
316 // Purely synchronous - uses Left(Ready), zero heap allocation
317
318 // 1. Check if already present in destination registry
319 let presence = self.registry.check_presence::<Dst>(&[ctx.sequence_hash]);
320 if presence[0].1 {
321 return sync_result(Ok(false)); // Already transferred
322 }
323
324 // 2. Check if currently in-flight (pending transfer)
325 if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
326 return sync_result(Ok(false)); // Already being transferred
327 }
328
329 sync_result(Ok(true)) // Not present, not pending - pass
330 }
331
332 fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<Src>]) -> PolicyBatchFuture<'a> {
333 if contexts.is_empty() {
334 return sync_batch_result(Ok(Vec::new()));
335 }
336
337 // Batch lookup for efficiency - still synchronous
338 let hashes: Vec<SequenceHash> = contexts.iter().map(|ctx| ctx.sequence_hash).collect();
339 let presence = self.registry.check_presence::<Dst>(&hashes);
340
341 // Build results checking both registry presence and pending status
342 let results: Vec<bool> = presence
343 .into_iter()
344 .map(|(hash, present)| {
345 if present {
346 return false;
347 }
348 if self.pending_tracker.is_hash_pending(&hash) {
349 return false;
350 }
351 true
352 })
353 .collect();
354
355 sync_batch_result(Ok(results))
356 }
357}
358
359/// G2→G3 filter: presence check + LFU count threshold.
360///
361/// Combines two filter conditions:
362/// 1. Skip blocks already present in destination tier
363/// 2. Only offload blocks with LFU count above threshold
364///
365/// The LFU threshold ensures we only offload "hot" blocks that have been
366/// accessed frequently, avoiding wasted transfers for rarely-used blocks.
367///
368/// # Duplicate Prevention
369///
370/// When a `PendingTracker` is configured, this filter also checks for blocks
371/// that are currently in-flight through the pipeline.
372///
373/// # Performance
374///
375/// This policy is fully synchronous and returns `Either::Left(Ready)`,
376/// avoiding any heap allocation per evaluation.
377///
378/// # Example
379/// ```ignore
380/// // Only offload blocks with LFU count > 8 that aren't in G3 or in-flight
381/// let tracker = Arc::new(PendingTracker::new());
382/// let filter = PresenceAndLFUFilter::<G2, G3>::new(registry.clone(), 8)
383/// .with_pending_tracker(tracker);
384/// ```
385pub struct PresenceAndLFUFilter<Src: BlockMetadata, Dst: BlockMetadata> {
386 registry: Arc<BlockRegistry>,
387 min_lfu_count: u32,
388 /// Optional tracker for pending (in-flight) transfers.
389 pending_tracker: Option<Arc<PendingTracker>>,
390 _marker: PhantomData<(Src, Dst)>,
391}
392
393impl<Src: BlockMetadata, Dst: BlockMetadata> PresenceAndLFUFilter<Src, Dst> {
394 /// Create a new presence + LFU filter with specified threshold.
395 pub fn new(registry: Arc<BlockRegistry>, min_lfu_count: u32) -> Self {
396 Self {
397 registry,
398 min_lfu_count,
399 pending_tracker: None,
400 _marker: PhantomData,
401 }
402 }
403
404 /// Create with default threshold of 8.
405 pub fn with_default_threshold(registry: Arc<BlockRegistry>) -> Self {
406 Self::new(registry, 8)
407 }
408
409 /// Add a pending tracker for duplicate prevention.
410 pub fn with_pending_tracker(mut self, tracker: Arc<PendingTracker>) -> Self {
411 self.pending_tracker = Some(tracker);
412 self
413 }
414}
415
416impl<Src: BlockMetadata, Dst: BlockMetadata> OffloadPolicy<Src> for PresenceAndLFUFilter<Src, Dst> {
417 fn name(&self) -> &str {
418 "PresenceAndLFUFilter"
419 }
420
421 fn evaluate<'a>(&'a self, ctx: &'a EvalContext<Src>) -> PolicyFuture<'a> {
422 // 1. Skip if already in Dst
423 let presence = self.registry.check_presence::<Dst>(&[ctx.sequence_hash]);
424 if presence[0].1 {
425 return sync_result(Ok(false));
426 }
427
428 // 2. Skip if currently pending transfer
429 if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
430 return sync_result(Ok(false));
431 }
432
433 // 3. Check LFU count > threshold
434 if let Some(tracker) = self.registry.frequency_tracker() {
435 // Convert SequenceHash to u128 for the tracker
436 let count = tracker.count(ctx.sequence_hash.as_u128());
437 return sync_result(Ok(count > self.min_lfu_count));
438 }
439
440 // No frequency tracker = pass all (conservative default)
441 sync_result(Ok(true))
442 }
443
444 fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<Src>]) -> PolicyBatchFuture<'a> {
445 if contexts.is_empty() {
446 return sync_batch_result(Ok(Vec::new()));
447 }
448
449 // Batch presence lookup
450 let hashes: Vec<SequenceHash> = contexts.iter().map(|ctx| ctx.sequence_hash).collect();
451 let presence = self.registry.check_presence::<Dst>(&hashes);
452
453 // Get trackers once
454 let freq_tracker = self.registry.frequency_tracker();
455 let min_lfu = self.min_lfu_count;
456
457 let results: Vec<bool> = presence
458 .into_iter()
459 .zip(contexts.iter())
460 .map(|((hash, present), ctx)| {
461 // Skip if present in Dst
462 if present {
463 return false;
464 }
465
466 // Skip if currently pending
467 if self.pending_tracker.is_hash_pending(&hash) {
468 return false;
469 }
470
471 // Check LFU count
472 if let Some(ref t) = freq_tracker {
473 let count = t.count(ctx.sequence_hash.as_u128());
474 count > min_lfu
475 } else {
476 true // No tracker = pass
477 }
478 })
479 .collect();
480
481 sync_batch_result(Ok(results))
482 }
483}
484
485/// G2→G4 filter: async presence check for object storage destinations.
486///
487/// Unlike `PresenceFilter` which checks local `BlockRegistry` synchronously,
488/// this filter queries object storage (S3, etc.) asynchronously via a
489/// `PresenceChecker` implementation.
490///
491/// # Duplicate Prevention
492///
493/// When a `PendingTracker` is configured, this filter also checks for blocks
494/// that are currently in-flight through the pipeline before querying object storage.
495///
496/// # Performance
497///
498/// This policy returns `Either::Right(BoxFuture)` since it requires async I/O.
499/// The pending tracker check is done synchronously first to avoid unnecessary
500/// object storage queries.
501///
502/// # Example
503/// ```ignore
504/// let object_ops: Arc<dyn ObjectBlockOps> = ...;
505/// let checker = Arc::new(S3PresenceChecker::new(object_ops));
506/// let tracker = Arc::new(PendingTracker::new());
507/// let filter = ObjectPresenceFilter::<G2>::new(checker)
508/// .with_pending_tracker(tracker);
509/// // Blocks already in object storage OR in-flight will be filtered out
510/// ```
511pub struct ObjectPresenceFilter<Src: BlockMetadata> {
512 presence_checker: Arc<dyn PresenceChecker>,
513 /// Optional tracker for pending (in-flight) transfers.
514 pending_tracker: Option<Arc<PendingTracker>>,
515 _marker: PhantomData<Src>,
516}
517
518impl<Src: BlockMetadata> ObjectPresenceFilter<Src> {
519 /// Create a new object presence filter.
520 pub fn new(presence_checker: Arc<dyn PresenceChecker>) -> Self {
521 Self {
522 presence_checker,
523 pending_tracker: None,
524 _marker: PhantomData,
525 }
526 }
527
528 /// Add a pending tracker for duplicate prevention.
529 ///
530 /// When set, blocks that are currently in-flight (passed policy but not
531 /// yet stored in object storage) will be filtered out.
532 pub fn with_pending_tracker(mut self, tracker: Arc<PendingTracker>) -> Self {
533 self.pending_tracker = Some(tracker);
534 self
535 }
536
537 /// Get a reference to the pending tracker if configured.
538 pub fn pending_tracker(&self) -> Option<&Arc<PendingTracker>> {
539 self.pending_tracker.as_ref()
540 }
541}
542
543impl<Src: BlockMetadata> OffloadPolicy<Src> for ObjectPresenceFilter<Src> {
544 fn name(&self) -> &str {
545 "ObjectPresenceFilter"
546 }
547
548 fn evaluate<'a>(&'a self, ctx: &'a EvalContext<Src>) -> PolicyFuture<'a> {
549 // 1. Synchronous check: skip if currently pending
550 if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
551 return sync_result(Ok(false)); // Already being transferred
552 }
553
554 // 2. Async check: query object storage for presence
555 let checker = self.presence_checker.clone();
556 let hash = ctx.sequence_hash;
557
558 async_result(async move {
559 let results = checker.check_presence(vec![hash]).await;
560 // If present in object storage, filter out
561 let exists = results
562 .into_iter()
563 .next()
564 .map(|(_, exists)| exists)
565 .unwrap_or(false);
566 Ok(!exists) // Pass if NOT present
567 })
568 }
569
570 fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<Src>]) -> PolicyBatchFuture<'a> {
571 if contexts.is_empty() {
572 return sync_batch_result(Ok(Vec::new()));
573 }
574
575 // Collect hashes, filtering out pending ones first (sync)
576 let mut pending_status: Vec<bool> = Vec::with_capacity(contexts.len());
577 let mut hashes_to_check: Vec<SequenceHash> = Vec::new();
578 let mut hash_indices: Vec<usize> = Vec::new();
579
580 for (i, ctx) in contexts.iter().enumerate() {
581 if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
582 pending_status.push(true); // Mark as pending (will be filtered)
583 } else {
584 pending_status.push(false);
585 hashes_to_check.push(ctx.sequence_hash);
586 hash_indices.push(i);
587 }
588 }
589
590 // If all are pending, return immediately
591 if hashes_to_check.is_empty() {
592 return sync_batch_result(Ok(vec![false; contexts.len()]));
593 }
594
595 let checker = self.presence_checker.clone();
596 let num_contexts = contexts.len();
597
598 async_batch_result(async move {
599 // Query object storage for non-pending hashes
600 let presence_results = checker.check_presence(hashes_to_check).await;
601
602 // Build final results
603 let mut results = vec![false; num_contexts]; // Default: filtered out
604
605 // Map presence results back to original indices
606 for (check_idx, original_idx) in hash_indices.into_iter().enumerate() {
607 if let Some((_, exists)) = presence_results.get(check_idx) {
608 // Pass if NOT present in object storage
609 results[original_idx] = !*exists;
610 }
611 }
612
613 Ok(results)
614 })
615 }
616}
617
618/// G2→G4 filter with distributed locking: check meta, acquire lock, track acquired locks.
619///
620/// This filter implements the full locking protocol for object storage offloads:
621/// 1. Check if `.meta` file exists (block already offloaded) - skip if yes
622/// 2. Check if currently pending (in-flight transfer) - skip if yes
623/// 3. Try to acquire `.lock` file with conditional PUT
624/// - If lock doesn't exist, create it atomically
625/// - If lock exists and expired, overwrite it
626/// - If lock exists and valid (owned by another instance), skip
627/// 4. If we own the lock (either just acquired or already owned), pass the block
628///
629/// # Lock Management
630///
631/// Locks acquired during policy evaluation are tracked and must be:
632/// - Released after successful transfer (via `ObjectTransferExecutor`)
633/// - Released on error/cancellation (via guard or explicit cleanup)
634///
635/// # Duplicate Prevention
636///
637/// When a `PendingTracker` is configured, blocks currently in-flight are filtered
638/// out before checking object storage, avoiding redundant network calls.
639///
640/// # Example
641/// ```ignore
642/// let lock_manager = Arc::new(S3LockManager::new(s3_client, instance_id));
643/// let tracker = Arc::new(PendingTracker::new());
644/// let filter = ObjectLockPresenceFilter::<G2>::new(lock_manager)
645/// .with_pending_tracker(tracker);
646/// // Blocks already offloaded, in-flight, or locked by others will be filtered out
647/// ```
648pub struct ObjectLockPresenceFilter<Src: BlockMetadata> {
649 lock_manager: Arc<dyn ObjectLockManager>,
650 /// Optional tracker for pending (in-flight) transfers.
651 pending_tracker: Option<Arc<PendingTracker>>,
652 _marker: PhantomData<Src>,
653}
654
655impl<Src: BlockMetadata> ObjectLockPresenceFilter<Src> {
656 /// Create a new object lock presence filter.
657 pub fn new(lock_manager: Arc<dyn ObjectLockManager>) -> Self {
658 Self {
659 lock_manager,
660 pending_tracker: None,
661 _marker: PhantomData,
662 }
663 }
664
665 /// Add a pending tracker for duplicate prevention.
666 ///
667 /// When set, blocks that are currently in-flight (passed policy but not
668 /// yet stored in object storage) will be filtered out.
669 pub fn with_pending_tracker(mut self, tracker: Arc<PendingTracker>) -> Self {
670 self.pending_tracker = Some(tracker);
671 self
672 }
673
674 /// Get a reference to the pending tracker if configured.
675 pub fn pending_tracker(&self) -> Option<&Arc<PendingTracker>> {
676 self.pending_tracker.as_ref()
677 }
678
679 /// Get a reference to the lock manager.
680 pub fn lock_manager(&self) -> &Arc<dyn ObjectLockManager> {
681 &self.lock_manager
682 }
683}
684
685impl<Src: BlockMetadata> OffloadPolicy<Src> for ObjectLockPresenceFilter<Src> {
686 fn name(&self) -> &str {
687 "ObjectLockPresenceFilter"
688 }
689
690 fn evaluate<'a>(&'a self, ctx: &'a EvalContext<Src>) -> PolicyFuture<'a> {
691 // 1. Synchronous check: skip if currently pending
692 if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
693 return sync_result(Ok(false)); // Already being transferred
694 }
695
696 // 2. Async checks: meta presence, then lock acquisition
697 let lock_manager = self.lock_manager.clone();
698 let hash = ctx.sequence_hash;
699
700 async_result(async move {
701 // Check if meta file exists (already offloaded)
702 match lock_manager.has_meta(hash).await {
703 Ok(true) => {
704 tracing::debug!(?hash, "Block already offloaded (meta exists)");
705 return Ok(false); // Already offloaded, skip
706 }
707 Ok(false) => {
708 // Continue to lock acquisition
709 }
710 Err(e) => {
711 tracing::warn!(?hash, error = %e, "Error checking meta file");
712 return Ok(false); // Error, skip to be safe
713 }
714 }
715
716 // Try to acquire lock
717 match lock_manager.try_acquire_lock(hash).await {
718 Ok(true) => {
719 tracing::debug!(?hash, "Lock acquired");
720 Ok(true) // Pass - we own the lock
721 }
722 Ok(false) => {
723 tracing::debug!(?hash, "Lock held by another instance");
724 Ok(false) // Skip - another instance owns the lock
725 }
726 Err(e) => {
727 tracing::warn!(?hash, error = %e, "Error acquiring lock");
728 Ok(false) // Error, skip to be safe
729 }
730 }
731 })
732 }
733
734 fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<Src>]) -> PolicyBatchFuture<'a> {
735 if contexts.is_empty() {
736 return sync_batch_result(Ok(Vec::new()));
737 }
738
739 // Filter out pending blocks first (sync)
740 let mut pending_mask: Vec<bool> = Vec::with_capacity(contexts.len());
741 let mut to_check: Vec<(usize, SequenceHash)> = Vec::new();
742
743 for (i, ctx) in contexts.iter().enumerate() {
744 if self.pending_tracker.is_hash_pending(&ctx.sequence_hash) {
745 pending_mask.push(true);
746 } else {
747 pending_mask.push(false);
748 to_check.push((i, ctx.sequence_hash));
749 }
750 }
751
752 // If all are pending, return immediately
753 if to_check.is_empty() {
754 return sync_batch_result(Ok(vec![false; contexts.len()]));
755 }
756
757 let lock_manager = self.lock_manager.clone();
758 let num_contexts = contexts.len();
759
760 async_batch_result(async move {
761 let mut results = vec![false; num_contexts]; // Default: filtered out
762
763 // Process each non-pending block
764 for (original_idx, hash) in to_check {
765 // Check meta first
766 let has_meta = match lock_manager.has_meta(hash).await {
767 Ok(has) => has,
768 Err(e) => {
769 tracing::warn!(?hash, error = %e, "Error checking meta file");
770 continue; // Skip this block
771 }
772 };
773
774 if has_meta {
775 tracing::debug!(?hash, "Block already offloaded (meta exists)");
776 continue; // Already offloaded
777 }
778
779 // Try to acquire lock
780 match lock_manager.try_acquire_lock(hash).await {
781 Ok(true) => {
782 tracing::debug!(?hash, "Lock acquired");
783 results[original_idx] = true; // Pass
784 }
785 Ok(false) => {
786 tracing::debug!(?hash, "Lock held by another instance");
787 // Skip - another instance owns the lock
788 }
789 Err(e) => {
790 tracing::warn!(?hash, error = %e, "Error acquiring lock");
791 // Skip on error
792 }
793 }
794 }
795
796 Ok(results)
797 })
798 }
799}
800
801/// Composite policy that requires ALL sub-policies to pass (AND logic).
802pub struct AllOfPolicy<T: BlockMetadata> {
803 policies: Vec<Arc<dyn OffloadPolicy<T>>>,
804}
805
806impl<T: BlockMetadata> AllOfPolicy<T> {
807 /// Create a new AND composite policy.
808 pub fn new(policies: Vec<Arc<dyn OffloadPolicy<T>>>) -> Self {
809 Self { policies }
810 }
811
812 /// Add a policy to the composite.
813 pub fn with(mut self, policy: Arc<dyn OffloadPolicy<T>>) -> Self {
814 self.policies.push(policy);
815 self
816 }
817}
818
819impl<T: BlockMetadata> OffloadPolicy<T> for AllOfPolicy<T> {
820 fn name(&self) -> &str {
821 "AllOfPolicy"
822 }
823
824 fn evaluate<'a>(&'a self, ctx: &'a EvalContext<T>) -> PolicyFuture<'a> {
825 // Must use async because sub-policies might be async
826 let policies = &self.policies;
827 async_result(async move {
828 for policy in policies {
829 let result = match policy.evaluate(ctx) {
830 Either::Left(ready) => ready.await,
831 Either::Right(boxed) => boxed.await,
832 };
833 if !result? {
834 return Ok(false);
835 }
836 }
837 Ok(true)
838 })
839 }
840}
841
842/// Composite policy that requires ANY sub-policy to pass (OR logic).
843pub struct AnyOfPolicy<T: BlockMetadata> {
844 policies: Vec<Arc<dyn OffloadPolicy<T>>>,
845}
846
847impl<T: BlockMetadata> AnyOfPolicy<T> {
848 /// Create a new OR composite policy.
849 pub fn new(policies: Vec<Arc<dyn OffloadPolicy<T>>>) -> Self {
850 Self { policies }
851 }
852
853 /// Add a policy to the composite.
854 pub fn with(mut self, policy: Arc<dyn OffloadPolicy<T>>) -> Self {
855 self.policies.push(policy);
856 self
857 }
858}
859
860impl<T: BlockMetadata> OffloadPolicy<T> for AnyOfPolicy<T> {
861 fn name(&self) -> &str {
862 "AnyOfPolicy"
863 }
864
865 fn evaluate<'a>(&'a self, ctx: &'a EvalContext<T>) -> PolicyFuture<'a> {
866 if self.policies.is_empty() {
867 return sync_result(Ok(true)); // No policies = pass
868 }
869
870 // Must use async because sub-policies might be async
871 let policies = &self.policies;
872 async_result(async move {
873 for policy in policies {
874 let result = match policy.evaluate(ctx) {
875 Either::Left(ready) => ready.await,
876 Either::Right(boxed) => boxed.await,
877 };
878 if result? {
879 return Ok(true);
880 }
881 }
882 Ok(false)
883 })
884 }
885}
886
887/// A pass-all policy (no filtering).
888///
889/// # Performance
890///
891/// This policy is fully synchronous and returns `Either::Left(Ready)`,
892/// avoiding any heap allocation per evaluation.
893pub struct PassAllPolicy<T: BlockMetadata> {
894 _marker: PhantomData<T>,
895}
896
897impl<T: BlockMetadata> PassAllPolicy<T> {
898 /// Create a new pass-all policy.
899 pub fn new() -> Self {
900 Self {
901 _marker: PhantomData,
902 }
903 }
904}
905
906impl<T: BlockMetadata> Default for PassAllPolicy<T> {
907 fn default() -> Self {
908 Self::new()
909 }
910}
911
912impl<T: BlockMetadata> OffloadPolicy<T> for PassAllPolicy<T> {
913 fn name(&self) -> &str {
914 "PassAllPolicy"
915 }
916
917 fn evaluate<'a>(&'a self, _ctx: &'a EvalContext<T>) -> PolicyFuture<'a> {
918 // Zero allocation - just returns ready(Ok(true))
919 sync_result(Ok(true))
920 }
921
922 fn evaluate_batch<'a>(&'a self, contexts: &'a [EvalContext<T>]) -> PolicyBatchFuture<'a> {
923 sync_batch_result(Ok(vec![true; contexts.len()]))
924 }
925}
926
927/// Create a composite policy from tier configuration.
928///
929/// Policies are applied in order with AND logic - blocks must pass all policies.
930/// Returns `PassAllPolicy` if no policies are configured.
931///
932/// When a `pending_tracker` is provided, it is automatically wired into
933/// `Presence` and `PresenceLfu` policies to enable duplicate prevention
934/// for blocks currently in-flight through the pipeline.
935///
936/// # Example
937///
938/// ```ignore
939/// use kvbm_config::offload::TierOffloadConfig;
940///
941/// let tracker = Arc::new(PendingTracker::new());
942/// let config = TierOffloadConfig {
943/// policies: vec![PolicyType::Presence, PolicyType::PresenceLfu],
944/// presence_lfu: PresenceLfuFilterConfig { min_lfu_count: 8 },
945/// ..Default::default()
946/// };
947///
948/// // Pending tracker is automatically wired into presence-based policies
949/// let policy = create_policy_from_config::<G2, G3>(&config, registry.clone(), Some(tracker));
950/// ```
951pub fn create_policy_from_config<Src, Dst>(
952 config: &TierOffloadConfig,
953 registry: Arc<BlockRegistry>,
954 pending_tracker: Option<Arc<PendingTracker>>,
955) -> Arc<dyn OffloadPolicy<Src>>
956where
957 Src: BlockMetadata + 'static,
958 Dst: BlockMetadata + 'static,
959{
960 if config.policies.is_empty() {
961 return Arc::new(PassAllPolicy::<Src>::new());
962 }
963
964 let policies: Vec<Arc<dyn OffloadPolicy<Src>>> = config
965 .policies
966 .iter()
967 .map(|policy_type| -> Arc<dyn OffloadPolicy<Src>> {
968 match policy_type {
969 PolicyType::PassAll => Arc::new(PassAllPolicy::<Src>::new()),
970 PolicyType::Presence => {
971 let mut filter = PresenceFilter::<Src, Dst>::new(registry.clone());
972 if let Some(tracker) = &pending_tracker {
973 filter = filter.with_pending_tracker(tracker.clone());
974 }
975 Arc::new(filter)
976 }
977 PolicyType::PresenceLfu => {
978 let mut filter = PresenceAndLFUFilter::<Src, Dst>::new(
979 registry.clone(),
980 config.presence_lfu.min_lfu_count,
981 );
982 if let Some(tracker) = &pending_tracker {
983 filter = filter.with_pending_tracker(tracker.clone());
984 }
985 Arc::new(filter)
986 }
987 }
988 })
989 .collect();
990
991 if policies.len() == 1 {
992 policies.into_iter().next().unwrap()
993 } else {
994 Arc::new(AllOfPolicy::new(policies))
995 }
996}
997
998#[cfg(test)]
999mod tests {
1000 use super::*;
1001
1002 // Note: Full tests require BlockRegistry infrastructure which needs
1003 // tokio runtime and complex setup. Basic API tests here.
1004
1005 #[test]
1006 fn test_pass_all_policy() {
1007 let _policy: PassAllPolicy<()> = PassAllPolicy::new();
1008 // Would test evaluate with proper setup
1009 }
1010
1011 #[test]
1012 fn test_all_of_policy_creation() {
1013 let policies: Vec<Arc<dyn OffloadPolicy<()>>> = vec![Arc::new(PassAllPolicy::new())];
1014 let composite = AllOfPolicy::new(policies);
1015 assert_eq!(composite.name(), "AllOfPolicy");
1016 }
1017
1018 #[test]
1019 fn test_any_of_policy_creation() {
1020 let policies: Vec<Arc<dyn OffloadPolicy<()>>> = vec![Arc::new(PassAllPolicy::new())];
1021 let composite = AnyOfPolicy::new(policies);
1022 assert_eq!(composite.name(), "AnyOfPolicy");
1023 }
1024
1025 #[tokio::test]
1026 async fn test_sync_result_zero_alloc() {
1027 // Verify sync_result returns Left variant
1028 let future = sync_result(Ok(true));
1029 assert!(matches!(future, Either::Left(_)));
1030
1031 let result = match future {
1032 Either::Left(ready) => ready.await,
1033 Either::Right(_) => unreachable!(),
1034 };
1035 assert!(result.unwrap());
1036 }
1037
1038 #[tokio::test]
1039 async fn test_async_result_boxes() {
1040 // Verify async_result returns Right variant
1041 let future = async_result(async { Ok(false) });
1042 assert!(matches!(future, Either::Right(_)));
1043
1044 let result = match future {
1045 Either::Left(_) => unreachable!(),
1046 Either::Right(boxed) => boxed.await,
1047 };
1048 assert!(!result.unwrap());
1049 }
1050
1051 #[test]
1052 fn test_pending_tracker_wiring() {
1053 use super::PendingTracker;
1054
1055 // Verify pending_tracker can be set on PresenceFilter
1056 let tracker = Arc::new(PendingTracker::new());
1057 let registry = Arc::new(BlockRegistry::new());
1058
1059 let filter: PresenceFilter<(), ()> =
1060 PresenceFilter::new(registry).with_pending_tracker(tracker.clone());
1061
1062 // Verify we can get the tracker back
1063 assert!(filter.pending_tracker().is_some());
1064 assert!(Arc::ptr_eq(filter.pending_tracker().unwrap(), &tracker));
1065 }
1066
1067 #[test]
1068 fn test_pending_tracker_wiring_lfu() {
1069 use super::PendingTracker;
1070
1071 // Verify pending_tracker can be set on PresenceAndLFUFilter
1072 let tracker = Arc::new(PendingTracker::new());
1073 let registry = Arc::new(BlockRegistry::new());
1074
1075 let filter: PresenceAndLFUFilter<(), ()> =
1076 PresenceAndLFUFilter::new(registry, 8).with_pending_tracker(tracker);
1077
1078 // Filter was successfully created with pending tracker
1079 assert_eq!(filter.name(), "PresenceAndLFUFilter");
1080 }
1081}