Skip to main content

kvbm_engine/offload/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Offload Engine for asynchronous block transfers between storage tiers.
5//!
6//! The offload engine provides a policy-based, cancellable pipeline for moving
7//! blocks from higher-performance tiers (G1/G2) to lower-cost tiers (G3/G4).
8//!
9//! # Architecture
10//!
11//! ```text
12//! ┌─────────────────────────────────────────────────────────────────┐
13//! │                        OffloadEngine                            │
14//! │                                                                 │
15//! │  ┌───────────────┐    ┌───────────────┐    ┌───────────────┐    │
16//! │  │G1→G2 Pipeline │────│ G2→G3 Pipeline│    │ G2→G4 Pipeline│    │
17//! │  └───────────────┘    └───────────────┘    └───────────────┘    │
18//! │         │                     │                     │           │
19//! │         └─────────auto_chain──┘                     │           │
20//! │                                                                 │
21//! └─────────────────────────────────────────────────────────────────┘
22//!
23//! Pipeline stages:
24//! ┌─────────────┐    ┌────────────────┐    ┌──────────────────┐
25//! │   Policy    │───▶│     Batch      │───▶│    Transfer      │
26//! │  Evaluator  │    │   Collector    │    │    Executor      │
27//! └─────────────┘    └────────────────┘    └──────────────────┘
28//!       │                   │                      │
29//!       ▼                   ▼                      ▼
30//!   cancel check       cancel check          wait for in-flight
31//! ```
32//!
33//! # Features
34//!
35//! - **Policy-based filtering**: Blocks pass through configurable policies
36//!   (presence checks, LFU thresholds) before transfer
37//! - **Batched transfers**: Blocks are accumulated into batches for efficient
38//!   bulk transfers
39//! - **Cancellation**: Clean cancellation with confirmation that all blocks
40//!   are released and no outstanding operations remain
41//! - **Pipeline chaining**: G1→G2 completions can automatically feed G2→G3
42//!
43//! See also: [Developer Guide](../../docs/offload-developer.md) for implementation
44//! details and extension rules.
45//!
46//! # Example
47//!
48//! ```ignore
49//! use kvbm::v2::distributed::offload::{
50//!     OffloadEngine, PipelineBuilder, PresenceFilter, PresenceAndLFUFilter,
51//! };
52//!
53//! // Build engine with pipelines
54//! let engine = OffloadEngine::builder(leader.clone())
55//!     .with_registry(registry.clone())
56//!     .with_g2_manager(g2_manager.clone())
57//!     .with_g3_manager(g3_manager.clone())
58//!     .with_g2_to_g3_pipeline(
59//!         PipelineBuilder::<G2, G3>::new()
60//!             .policy(Arc::new(PresenceAndLFUFilter::with_default_threshold(registry.clone())))
61//!             .batch_size(64)
62//!             .build()
63//!     )
64//!     .build()?;
65//!
66//! // Enqueue blocks for offload
67//! let handle = engine.enqueue_g2_to_g3(blocks)?;
68//!
69//! // Wait for completion or cancel
70//! tokio::select! {
71//!     result = handle.wait() => {
72//!         println!("Completed: {:?}", result?.completed_blocks);
73//!     }
74//!     _ = shutdown_signal => {
75//!         handle.cancel().wait().await;
76//!         println!("Cancelled");
77//!     }
78//! }
79//! ```
80//!
81//! See also: [Developer Guide](../../docs/offload-developer.md)
82
83/// Helper macro to create an NVTX range when the nvtx feature is enabled.
84/// The range automatically ends when the returned guard is dropped.
85macro_rules! nvtx_range {
86    ($name:expr) => {{
87        #[cfg(feature = "nvtx")]
88        let _range = nvtx::range!($name);
89        #[cfg(not(feature = "nvtx"))]
90        let _range = ();
91        _range
92    }};
93}
94
95mod batch;
96mod cancel;
97mod engine;
98mod handle;
99mod pending;
100mod pipeline;
101mod policy;
102mod queue;
103mod settlement;
104mod source;
105
106#[cfg(test)]
107mod cancel_tests;
108
109// Re-export public API
110pub use cancel::{CancelConfirmation, CancelState, CancellationToken};
111pub use engine::{OffloadEngine, OffloadEngineBuilder};
112pub use handle::{
113    TransferHandle, TransferId, TransferProgressCounts, TransferProgressCursor,
114    TransferProgressDelta, TransferResult, TransferStatus,
115};
116pub use pending::{PendingGuard, PendingTracker};
117pub use pipeline::{
118    ObjectPipeline, ObjectPipelineBuilder, ObjectPipelineConfig, Pipeline, PipelineBuilder,
119    PipelineConfig, ResolvedBatch, ResolvedBlock, upgrade_batch,
120};
121pub use policy::{
122    AllOfPolicy, AnyOfPolicy, BoxFuture, EvalContext, ObjectLockPresenceFilter,
123    ObjectPresenceFilter, OffloadPolicy, PassAllPolicy, PolicyBatchFuture, PolicyFuture,
124    PresenceAndLFUFilter, PresenceChecker, PresenceFilter, S3PresenceChecker, async_batch_result,
125    async_result, create_policy_from_config, sync_batch_result, sync_result,
126};
127pub use queue::CancellableQueue;
128pub use settlement::{
129    PipelineFailure, PipelineFailureKind, PipelineLane, SettlementError, SettlementTarget,
130    SettlementToken,
131};
132pub use source::{ExternalBlock, SourceBlock, SourceBlocks};
133
134// Re-export batch config for advanced users
135pub use batch::{BatchConfig, TimingTrace};