Skip to main content

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}