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};