Skip to main content

xet_client/cas_client/
interface.rs

1use std::sync::Arc;
2
3use bytes::Bytes;
4use xet_core_structures::merklehash::MerkleHash;
5use xet_core_structures::metadata_shard::file_structs::MDBFileInfo;
6use xet_core_structures::xorb_object::SerializedXorbObject;
7
8use super::adaptive_concurrency::ConnectionPermit;
9use super::progress_tracked_streams::ProgressCallback;
10use crate::cas_types::{
11    BatchQueryReconstructionResponse, FileChunkHashesResponse, FileRange, HttpRange, QueryReconstructionResponseV2,
12    ShardUploadEvent,
13};
14use crate::error::Result;
15
16/// Progress update delivered to the shard-upload callback.
17#[derive(Debug)]
18pub enum ShardUploadProgressType<'a> {
19    /// Incremental request-body bytes transferred to the server.
20    Transfer(u64),
21    /// Subtract previously reported transfer bytes (e.g. after a failed V2 attempt
22    /// before restarting the upload on V1).
23    DecrementTransfer(u64),
24    /// NDJSON progress event from the v2 response stream.
25    Response(&'a ShardUploadEvent),
26}
27pub type ShardUploadProgressCallback = Arc<dyn Fn(ShardUploadProgressType) + Send + Sync>;
28
29#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
30#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
31pub trait URLProvider: Send + Sync {
32    /// Retrieves the URL and the byte ranges to fetch.
33    /// For single-range (V1) blocks, the Vec has one entry.
34    /// For multi-range (V2) blocks, all ranges are included.
35    async fn retrieve_url(&self) -> Result<(String, Vec<HttpRange>)>;
36
37    /// Asks for a refresh of the URL; triggered on 403 errors.
38    async fn refresh_url(&self) -> Result<()>;
39}
40
41/// A Client to the Shard service. The shard service
42/// provides for
43/// 1. upload shard to the shard service
44/// 2. querying of file->reconstruction information
45/// 3. querying of chunk->shard information
46#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
47#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
48pub trait Client: Send + Sync {
49    async fn get_file_reconstruction_info(
50        &self,
51        file_hash: &MerkleHash,
52    ) -> Result<Option<(MDBFileInfo, Option<MerkleHash>)>>;
53
54    /// Returns reconstruction info always in V2 format.
55    /// Implementations may try V2 first and fall back to V1 + convert.
56    async fn get_reconstruction(
57        &self,
58        file_id: &MerkleHash,
59        bytes_range: Option<FileRange>,
60    ) -> Result<Option<QueryReconstructionResponseV2>>;
61
62    async fn batch_get_reconstruction(&self, file_ids: &[MerkleHash]) -> Result<BatchQueryReconstructionResponse>;
63
64    async fn acquire_download_permit(&self) -> Result<ConnectionPermit>;
65
66    /// Optional progress callback receives (delta, completed, total) in transfer bytes.
67    /// When [uncompressed_size_if_known] is [Some], the returned Bytes must have len() equal to that value.
68    async fn get_file_term_data(
69        &self,
70        url_info: Box<dyn URLProvider>,
71        download_permit: ConnectionPermit,
72        progress_callback: Option<ProgressCallback>,
73        uncompressed_size_if_known: Option<usize>,
74    ) -> Result<(Bytes, Vec<u32>)>;
75
76    async fn query_for_global_dedup_shard(&self, prefix: &str, chunk_hash: &MerkleHash) -> Result<Option<Bytes>>;
77
78    /// Acquire an upload permit.
79    async fn acquire_upload_permit(&self) -> Result<ConnectionPermit>;
80
81    /// Upload a new shard. The optional callback receives v2 NDJSON progress events.
82    async fn upload_shard(
83        &self,
84        shard_data: bytes::Bytes,
85        upload_permit: ConnectionPermit,
86        progress_callback: Option<ShardUploadProgressCallback>,
87    ) -> Result<()>;
88
89    /// Upload a new xorb. Optional progress callback receives (delta, completed, total) in transfer bytes.
90    async fn upload_xorb(
91        &self,
92        prefix: &str,
93        serialized_xorb_object: SerializedXorbObject,
94        progress_callback: Option<ProgressCallback>,
95        upload_permit: ConnectionPermit,
96    ) -> Result<u64>;
97
98    /// Compute chunk-aligned dirty windows + opaque gap [`MerkleHashSubtree`] summaries for the
99    /// given file, narrowed to `dirty_ranges`.
100    ///
101    /// `dirty_ranges` must be sorted and non-overlapping. Per-chunk hashes are never returned;
102    /// the response carries only `windows.len()` dirty windows and `windows.len() + 1` gap
103    /// subtrees, which the client merges with locally-recomputed window subtrees to obtain the
104    /// new file hash.
105    async fn get_file_chunk_hashes(
106        &self,
107        file_id: &MerkleHash,
108        dirty_ranges: Vec<FileRange>,
109    ) -> Result<FileChunkHashesResponse>;
110}