1use std::sync::Arc;
2use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
3
4use anyhow::anyhow;
5use bytes::Bytes;
6use futures::TryStreamExt;
7use http::HeaderValue;
8use http::header::{CONTENT_LENGTH, HeaderMap, RANGE};
9use reqwest::{Body, Response, StatusCode, Url};
10use reqwest_middleware::ClientWithMiddleware;
11use tracing::{event, info, instrument};
12use xet_core_structures::merklehash::MerkleHash;
13use xet_core_structures::metadata_shard::file_structs::{FileDataSequenceEntry, FileDataSequenceHeader, MDBFileInfo};
14use xet_core_structures::xorb_object::SerializedXorbObject;
15use xet_runtime::core::XetContext;
16
17use super::adaptive_concurrency::{
18 AdaptiveConcurrencyController, ConnectionPermit, download_controller, upload_controller,
19};
20use super::auth::AuthConfig;
21use super::interface::URLProvider;
22use super::progress_tracked_streams::{
23 DownloadProgressStream, ProgressCallback, StreamProgressReporter, UploadProgressStream,
24};
25use super::retry_wrapper::{RetryWrapper, RetryableReqwestError};
26use super::{Client, INFORMATION_LOG_LEVEL};
27use crate::cas_types::{
28 BatchQueryReconstructionResponse, FileChunkHashesResponse, FileRange, HttpRange, Key, QueryReconstructionResponse,
29 QueryReconstructionResponseV2, UploadShardResponse, UploadShardResponseType, UploadXorbResponse,
30 X_RANGE_DIRTY_HEADER,
31};
32use crate::common::http_client::{self, Api};
33use crate::error::{ClientError, Result};
34
35pub const CAS_ENDPOINT: &str = "http://localhost:8080";
36pub const PREFIX_DEFAULT: &str = "default";
37
38use lazy_static::lazy_static;
39
40lazy_static! {
41 static ref FN_CALL_ID: AtomicU64 = AtomicU64::new(1);
42}
43
44pub struct RemoteClient {
45 pub(crate) ctx: XetContext,
46 endpoint: String,
47 dry_run: bool,
48 http_client: Arc<ClientWithMiddleware>,
49 authenticated_http_client: Arc<ClientWithMiddleware>,
50 #[cfg(not(target_family = "wasm"))]
53 shard_upload_http_client: Arc<ClientWithMiddleware>,
54 upload_concurrency_controller: Arc<AdaptiveConcurrencyController>,
55 download_concurrency_controller: Arc<AdaptiveConcurrencyController>,
56 detected_reconstruction_api_version: AtomicU32,
58}
59
60impl RemoteClient {
61 pub fn new_with_socket(
71 ctx: XetContext,
72 endpoint: &str,
73 auth: &Option<AuthConfig>,
74 session_id: &str,
75 dry_run: bool,
76 unix_socket_path: Option<&str>,
77 custom_headers: Option<Arc<HeaderMap>>,
78 ) -> Arc<Self> {
79 Arc::new(Self {
80 ctx: ctx.clone(),
81 endpoint: endpoint.to_string(),
82 dry_run,
83 authenticated_http_client: Arc::new(
84 http_client::build_auth_http_client(&ctx, auth, session_id, unix_socket_path, custom_headers.clone())
85 .unwrap(),
86 ),
87 http_client: Arc::new(
88 http_client::build_http_client(&ctx, session_id, unix_socket_path, custom_headers.clone()).unwrap(),
89 ),
90 #[cfg(not(target_family = "wasm"))]
91 shard_upload_http_client: Arc::new(
92 http_client::build_auth_http_client_no_read_timeout(
93 &ctx,
94 auth,
95 session_id,
96 unix_socket_path,
97 custom_headers,
98 )
99 .unwrap(),
100 ),
101 upload_concurrency_controller: upload_controller(&ctx, endpoint),
102 download_concurrency_controller: download_controller(&ctx, endpoint),
103 detected_reconstruction_api_version: AtomicU32::new(0),
104 })
105 }
106
107 pub fn new(
119 ctx: XetContext,
120 endpoint: &str,
121 auth: &Option<AuthConfig>,
122 session_id: &str,
123 dry_run: bool,
124 custom_headers: Option<Arc<HeaderMap>>,
125 ) -> Arc<Self> {
126 Self::new_with_socket(ctx, endpoint, auth, session_id, dry_run, None, custom_headers)
127 }
128
129 pub fn endpoint(&self) -> &str {
131 &self.endpoint
132 }
133
134 #[cfg(feature = "simulation")]
135 pub(crate) fn http_client(&self) -> Arc<ClientWithMiddleware> {
136 self.http_client.clone()
137 }
138
139 async fn query_dedup_api(&self, prefix: &str, chunk_hash: &MerkleHash) -> Result<Option<Response>> {
140 let key = Key {
142 prefix: prefix.into(),
143 hash: *chunk_hash,
144 };
145
146 let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
147 let url = Url::parse(&format!("{}/v1/chunks/{key}", self.endpoint))?;
148 event!(
149 INFORMATION_LOG_LEVEL,
150 call_id,
151 prefix,
152 %chunk_hash,
153 "Starting query_dedup API call",
154 );
155
156 let client = self.authenticated_http_client.clone();
157 let api_tag = "cas::query_dedup";
158
159 let result = RetryWrapper::new(self.ctx.clone(), api_tag)
160 .with_429_no_retry()
161 .with_expected_404()
162 .log_errors_as_info()
163 .run(move || client.get(url.clone()).with_extension(Api(api_tag)).send())
164 .await;
165
166 if result.as_ref().is_err_and(|e| e.status().is_some()) {
167 event!(
168 INFORMATION_LOG_LEVEL,
169 call_id,
170 prefix,
171 %chunk_hash,
172 result="not_found",
173 "Completed query_dedup API call",
174 );
175 return Ok(None);
176 }
177
178 event!(
179 INFORMATION_LOG_LEVEL,
180 call_id,
181 prefix,
182 %chunk_hash,
183 result="found",
184 "Completed query_dedup API call",
185 );
186 Ok(Some(result?))
187 }
188}
189
190impl RemoteClient {
191 async fn get_reconstruction_impl<T>(
192 &self,
193 file_id: &MerkleHash,
194 bytes_range: Option<FileRange>,
195 api_version: &str,
196 ) -> Result<Option<T>>
197 where
198 T: serde::de::DeserializeOwned + 'static,
199 {
200 let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
201 let url = Url::parse(&format!("{}/{api_version}/reconstructions/{}", self.endpoint, file_id.hex()))?;
202 let api_tag = match api_version {
203 "v1" => "cas::get_reconstruction_v1",
204 "v2" => "cas::get_reconstruction_v2",
205 _ => {
206 return Err(ClientError::InternalError(anyhow!(
207 "unsupported reconstruction API version: {api_version}"
208 )));
209 },
210 };
211
212 event!(
213 INFORMATION_LOG_LEVEL,
214 call_id,
215 %file_id,
216 ?bytes_range,
217 api_version,
218 "Starting get_reconstruction API call",
219 );
220
221 let client = self.authenticated_http_client.clone();
222
223 let result: Result<T> = RetryWrapper::new(self.ctx.clone(), api_tag)
224 .with_expected_416()
225 .run_and_extract_json(move || {
226 let mut request = client.get(url.clone()).with_extension(Api(api_tag));
227 if let Some(range) = bytes_range {
228 request = request.header(RANGE, HttpRange::from(range).range_header())
229 }
230 request.send()
231 })
232 .await;
233
234 match result {
235 Ok(response) => {
236 event!(
237 INFORMATION_LOG_LEVEL,
238 call_id,
239 %file_id,
240 ?bytes_range,
241 api_version,
242 "Completed get_reconstruction API call"
243 );
244 Ok(Some(response))
245 },
246 Err(ClientError::ReqwestError(ref e, _)) if e.status() == Some(StatusCode::RANGE_NOT_SATISFIABLE) => {
247 Ok(None)
248 },
249 Err(e) => Err(e),
250 }
251 }
252
253 pub async fn get_reconstruction_v1(
255 &self,
256 file_id: &MerkleHash,
257 bytes_range: Option<FileRange>,
258 ) -> Result<Option<QueryReconstructionResponse>> {
259 self.get_reconstruction_impl(file_id, bytes_range, "v1").await
260 }
261
262 pub async fn get_reconstruction_v2(
264 &self,
265 file_id: &MerkleHash,
266 bytes_range: Option<FileRange>,
267 ) -> Result<Option<QueryReconstructionResponseV2>> {
268 self.get_reconstruction_impl(file_id, bytes_range, "v2").await
269 }
270
271 pub(crate) async fn get_reconstruction_with_version_override(
272 &self,
273 file_id: &MerkleHash,
274 bytes_range: Option<FileRange>,
275 forced_version: Option<u32>,
276 ) -> Result<Option<QueryReconstructionResponseV2>> {
277 let version = match forced_version {
280 Some(v) => v,
281 None => {
282 let detected = self.detected_reconstruction_api_version.load(Ordering::Relaxed);
283 if detected != 0 { detected } else { 2 }
284 },
285 };
286
287 match version {
288 2 => match self.get_reconstruction_v2(file_id, bytes_range).await {
289 Ok(result) => {
290 if forced_version.is_none() {
291 self.detected_reconstruction_api_version.store(2, Ordering::Relaxed);
292 }
293 Ok(result)
294 },
295 Err(e)
296 if forced_version.is_none()
297 && matches!(e.status(), Some(StatusCode::NOT_FOUND) | Some(StatusCode::NOT_IMPLEMENTED)) =>
298 {
299 info!(status = ?e.status(), "V2 reconstruction not available, falling back to V1");
300 let result = self.get_reconstruction_v1(file_id, bytes_range).await?.map(Into::into);
301 self.detected_reconstruction_api_version.store(1, Ordering::Relaxed);
303 Ok(result)
304 },
305 Err(e) => Err(e),
306 },
307 1 => Ok(self.get_reconstruction_v1(file_id, bytes_range).await?.map(Into::into)),
308 other => Err(ClientError::InternalError(anyhow!("unsupported reconstruction API version: {other}"))),
309 }
310 }
311}
312
313#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
314#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
315impl Client for RemoteClient {
316 async fn get_reconstruction(
317 &self,
318 file_id: &MerkleHash,
319 bytes_range: Option<FileRange>,
320 ) -> Result<Option<QueryReconstructionResponseV2>> {
321 let forced_version = self.ctx.config.client.reconstruction_api_version;
322 self.get_reconstruction_with_version_override(file_id, bytes_range, forced_version)
323 .await
324 }
325
326 async fn batch_get_reconstruction(&self, file_ids: &[MerkleHash]) -> Result<BatchQueryReconstructionResponse> {
327 let mut url_str = format!("{}/reconstructions?", self.endpoint);
328 let mut is_first = true;
329 let mut file_id_list = Vec::new();
330 for hash in file_ids {
331 file_id_list.push(hash.hex());
332 if is_first {
333 is_first = false;
334 } else {
335 url_str.push('&');
336 }
337 url_str.push_str("file_id=");
338 url_str.push_str(hash.hex().as_str());
339 }
340 let url: Url = url_str.parse()?;
341
342 let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
343 info!(call_id, file_ids=?file_id_list, "Starting batch_get_reconstruction API call");
344
345 let api_tag = "cas::batch_get_reconstruction";
346 let client = self.authenticated_http_client.clone();
347
348 let response: BatchQueryReconstructionResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
349 .run_and_extract_json(move || client.get(url.clone()).with_extension(Api(api_tag)).send())
350 .await?;
351
352 info!(call_id,
353 file_ids=?file_id_list,
354 response_count=response.files.len(),
355 "Completed batch_get_reconstruction API call",
356 );
357
358 Ok(response)
359 }
360
361 async fn acquire_download_permit(&self) -> Result<ConnectionPermit> {
362 self.download_concurrency_controller.acquire_connection_permit().await
363 }
364
365 async fn get_file_term_data(
366 &self,
367 url_info: Box<dyn URLProvider>,
368 download_permit: ConnectionPermit,
369 progress_callback: Option<ProgressCallback>,
370 uncompressed_size_if_known: Option<usize>,
371 ) -> Result<(Bytes, Vec<u32>)> {
372 let api_tag = "s3::get_range";
373 let http_client = self.http_client.clone();
374 let url_info = Arc::new(url_info);
375
376 let (_, url_ranges) = url_info.retrieve_url().await?;
377 let total_download_bytes: u64 = url_ranges.iter().map(|r| r.length()).sum();
378
379 let mut transfer_reporter = StreamProgressReporter::new(total_download_bytes)
380 .with_adaptive_concurrency_reporter(download_permit.get_partial_completion_reporting_function());
381 if let Some(cb) = progress_callback {
382 transfer_reporter = transfer_reporter.with_progress_callback(cb);
383 }
384
385 let result = RetryWrapper::new(self.ctx.clone(), api_tag)
386 .with_retry_on_403()
387 .with_connection_permit(download_permit, None)
388 .run_and_extract_custom(
389 move || {
390 let http_client = http_client.clone();
391 let url_info = url_info.clone();
392
393 async move {
394 let (url_string, url_ranges) = url_info
395 .retrieve_url()
396 .await
397 .map_err(|e| reqwest_middleware::Error::Middleware(e.into()))?;
398 let url =
399 Url::parse(&url_string).map_err(|e| reqwest_middleware::Error::Middleware(e.into()))?;
400
401 let range_header_value = if url_ranges.len() == 1 {
403 url_ranges[0].range_header()
404 } else {
405 let joined = url_ranges
406 .iter()
407 .map(|r| format!("{}-{}", r.start, r.end))
408 .collect::<Vec<_>>()
409 .join(",");
410 format!("bytes={joined}")
411 };
412
413 let response = http_client
414 .get(url)
415 .header(RANGE, range_header_value)
416 .with_extension(Api(api_tag))
417 .send()
418 .await?;
419
420 if response.status() == reqwest::StatusCode::FORBIDDEN {
421 url_info
422 .refresh_url()
423 .await
424 .map_err(|e| reqwest_middleware::Error::Middleware(e.into()))?;
425 }
426
427 Ok(response)
428 }
429 },
430 move |resp: Response| {
431 let transfer_reporter = transfer_reporter.clone();
432 async move {
433 let content_type = resp
434 .headers()
435 .get("content-type")
436 .and_then(|v| v.to_str().ok())
437 .unwrap_or("")
438 .to_string();
439
440 let is_multipart = content_type.contains("multipart/byteranges");
441
442 if is_multipart {
443 let body = resp
444 .bytes()
445 .await
446 .map_err(|e| RetryableReqwestError::RetryableError(ClientError::from(e)))?;
447
448 let multipart_parts = crate::cas_client::multipart::parse_multipart_byteranges(&content_type, body)
449 .map_err(RetryableReqwestError::FatalError)?;
450
451 let mut all_decompressed = Vec::with_capacity(uncompressed_size_if_known.unwrap_or(0));
452 let mut all_chunk_indices = Vec::<u32>::new();
453 let mut total_compressed_bytes = 0u64;
454
455 for part in multipart_parts {
456 total_compressed_bytes += part.data.len() as u64;
457
458 let (data, chunk_indices) =
459 xet_core_structures::xorb_object::deserialize_chunks(&mut std::io::Cursor::new(part.data.as_ref()))
460 .map_err(|e| {
461 RetryableReqwestError::RetryableError(ClientError::FormatError(e))
462 })?;
463
464 xet_core_structures::xorb_object::append_chunk_segment(
465 &mut all_decompressed,
466 &mut all_chunk_indices,
467 &data,
468 &chunk_indices,
469 );
470
471 transfer_reporter.report_progress(total_compressed_bytes as usize);
472 }
473
474 if let Some(expected) = uncompressed_size_if_known
475 && expected != all_decompressed.len()
476 {
477 return Err(RetryableReqwestError::RetryableError(ClientError::Other(format!(
478 "get_file_term_data: expected {expected} uncompressed bytes, got {}",
479 all_decompressed.len()
480 ))));
481 }
482 Ok((Bytes::from(all_decompressed), all_chunk_indices))
483 } else {
484 let incoming_stream = DownloadProgressStream::wrap_stream(
485 resp.bytes_stream().map_err(std::io::Error::other),
486 transfer_reporter,
487 );
488
489 let capacity = uncompressed_size_if_known.unwrap_or(0);
490 let mut buffer = Vec::with_capacity(capacity);
491 let mut writer = std::io::Cursor::new(&mut buffer);
492
493 let result = xet_core_structures::xorb_object::deserialize_async::deserialize_chunks_to_writer_from_stream(
494 incoming_stream,
495 &mut writer,
496 )
497 .await;
498
499 match result {
500 Ok((_compressed_len, chunk_byte_indices)) => {
501 if let Some(expected) = uncompressed_size_if_known
502 && expected != buffer.len()
503 {
504 return Err(RetryableReqwestError::RetryableError(ClientError::Other(format!(
505 "get_file_term_data: expected {expected} uncompressed bytes, got {}",
506 buffer.len()
507 ))));
508 }
509 Ok((Bytes::from(buffer), chunk_byte_indices))
510 },
511 Err(e) => Err(RetryableReqwestError::RetryableError(ClientError::FormatError(e))),
512 }
513 }
514 }
515 },
516 )
517 .await?;
518
519 Ok(result)
520 }
521
522 #[instrument(skip_all, name = "RemoteClient::get_file_reconstruction", fields(file.hash = file_hash.hex()
523 ))]
524 async fn get_file_reconstruction_info(
525 &self,
526 file_hash: &MerkleHash,
527 ) -> Result<Option<(MDBFileInfo, Option<MerkleHash>)>> {
528 let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
529 let url = Url::parse(&format!("{}/v1/reconstructions/{}", self.endpoint, file_hash.hex()))?;
530 event!(INFORMATION_LOG_LEVEL, call_id, %file_hash, "Starting get_file_reconstruction_info API call");
531
532 let api_tag = "cas::get_reconstruction_info";
533 let client = self.authenticated_http_client.clone();
534
535 let response: QueryReconstructionResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
536 .run_and_extract_json(move || client.get(url.clone()).with_extension(Api(api_tag)).send())
537 .await?;
538
539 let terms_count = response.terms.len();
540 let result = Some((
541 MDBFileInfo {
542 metadata: FileDataSequenceHeader::new(*file_hash, terms_count, false, false),
543 segments: response
544 .terms
545 .into_iter()
546 .map(|ce| {
547 FileDataSequenceEntry::new(ce.hash.into(), ce.unpacked_length, ce.range.start, ce.range.end)
548 })
549 .collect(),
550 verification: vec![],
551 metadata_ext: None,
552 },
553 None,
554 ));
555
556 event!(INFORMATION_LOG_LEVEL, call_id, %file_hash, terms_count, "Completed get_file_reconstruction_info API call");
557
558 Ok(result)
559 }
560
561 async fn query_for_global_dedup_shard(&self, prefix: &str, chunk_hash: &MerkleHash) -> Result<Option<Bytes>> {
562 let Some(response) = self.query_dedup_api(prefix, chunk_hash).await? else {
563 return Ok(None);
564 };
565
566 Ok(Some(response.bytes().await?))
567 }
568
569 async fn acquire_upload_permit(&self) -> Result<ConnectionPermit> {
570 self.upload_concurrency_controller.acquire_connection_permit().await
571 }
572
573 #[instrument(skip_all, name = "RemoteClient::upload_shard", fields(shard.len = shard_data.len()))]
574 async fn upload_shard(&self, shard_data: Bytes, upload_permit: ConnectionPermit) -> Result<bool> {
575 if self.dry_run {
576 return Ok(true);
577 }
578
579 let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
580 let n_upload_bytes = shard_data.len();
581 event!(INFORMATION_LOG_LEVEL, call_id, size = n_upload_bytes, "Starting upload_shard API call",);
582
583 let api_tag = "cas::upload_shard";
584 let url = Url::parse(&format!("{}/v1/shards", self.endpoint))?;
585
586 #[cfg(not(target_family = "wasm"))]
591 let client = self.shard_upload_http_client.clone();
592
593 #[cfg(target_family = "wasm")]
594 let client = self.authenticated_http_client.clone();
595
596 let response: UploadShardResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
597 .with_connection_permit(upload_permit, Some(shard_data.len() as u64))
598 .run_and_extract_json(move || {
599 client
600 .post(url.clone())
601 .with_extension(Api(api_tag))
602 .body(shard_data.clone())
603 .send()
604 })
605 .await?;
606
607 match response.result {
608 UploadShardResponseType::Exists => {
609 event!(
610 INFORMATION_LOG_LEVEL,
611 call_id,
612 size = n_upload_bytes,
613 result = "exists",
614 "Completed upload_shard API call",
615 );
616 Ok(false)
617 },
618 UploadShardResponseType::SyncPerformed => {
619 event!(
620 INFORMATION_LOG_LEVEL,
621 call_id,
622 size = n_upload_bytes,
623 result = "sync performed",
624 "Completed upload_shard API call",
625 );
626 Ok(true)
627 },
628 }
629 }
630
631 #[instrument(skip_all, name = "RemoteClient::upload_xorb", fields(key = Key{prefix : prefix.to_string(), hash : serialized_xorb_object.hash}.to_string(),
632 xorb.len = serialized_xorb_object.serialized_data.len(), xorb.num_chunks = serialized_xorb_object.num_chunks
633 ))]
634 async fn upload_xorb(
635 &self,
636 prefix: &str,
637 serialized_xorb_object: SerializedXorbObject,
638 progress_callback: Option<ProgressCallback>,
639 upload_permit: ConnectionPermit,
640 ) -> Result<u64> {
641 let key = Key {
642 prefix: prefix.to_string(),
643 hash: serialized_xorb_object.hash,
644 };
645
646 let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
647 let url = Url::parse(&format!("{}/v1/xorbs/{key}", self.endpoint))?;
648
649 let n_upload_bytes = serialized_xorb_object.serialized_data.len() as u64;
650 event!(
651 INFORMATION_LOG_LEVEL,
652 call_id,
653 prefix,
654 hash=%serialized_xorb_object.hash,
655 size=n_upload_bytes,
656 num_chunks=serialized_xorb_object.num_chunks,
657 "Starting upload_xorb API call",
658 );
659
660 let n_transfer_bytes = serialized_xorb_object.serialized_data.len() as u64;
661
662 let serialized_data = serialized_xorb_object.serialized_data.clone();
663
664 #[cfg(not(target_family = "wasm"))]
665 let block_size = self.ctx.config.client.upload_reporting_block_size;
666
667 let mut upload_reporter = StreamProgressReporter::new(n_transfer_bytes)
668 .with_adaptive_concurrency_reporter(upload_permit.get_partial_completion_reporting_function());
669 if let Some(cb) = progress_callback {
670 upload_reporter = upload_reporter.with_progress_callback(cb);
671 }
672
673 let xorb_uploaded = {
674 if !self.dry_run {
675 let client = self.authenticated_http_client.clone();
676
677 let api_tag = "cas::upload_xorb";
678
679 let response: UploadXorbResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
680 .with_connection_permit(upload_permit, Some(n_transfer_bytes))
681 .run_and_extract_json(move || {
682 let url = url.clone();
683 let serialized_data = serialized_data.clone();
684
685 let request = {
686 #[cfg(not(target_family = "wasm"))]
687 {
688 let upload_stream = UploadProgressStream::wrap_bytes_as_stream(
689 serialized_data,
690 block_size,
691 upload_reporter.clone(),
692 );
693 client
694 .post(url)
695 .with_extension(Api(api_tag))
696 .header(CONTENT_LENGTH, HeaderValue::from(n_upload_bytes)) .body(Body::wrap_stream(upload_stream))
698 }
699
700 #[cfg(target_family = "wasm")]
703 {
704 client.post(url).with_extension(Api(api_tag)).body(serialized_data)
705 }
706 };
707
708 request.send()
709 })
710 .await?;
711
712 #[cfg(target_family = "wasm")]
716 upload_reporter.report_progress(n_transfer_bytes as usize);
717
718 response.was_inserted
719 } else {
720 true
721 }
722 };
723
724 if !xorb_uploaded {
725 event!(
726 INFORMATION_LOG_LEVEL,
727 call_id,
728 prefix,
729 hash=%serialized_xorb_object.hash,
730 result="not_inserted",
731 "Completed upload_xorb API call",
732 );
733 } else {
734 event!(
735 INFORMATION_LOG_LEVEL,
736 call_id,
737 prefix,
738 hash=%serialized_xorb_object.hash,
739 size=n_upload_bytes,
740 result="inserted",
741 "Completed upload_xorb API call",
742 );
743 }
744
745 Ok(n_upload_bytes)
746 }
747
748 #[instrument(skip_all, name = "RemoteClient::get_file_chunk_hashes", fields(file.hash = file_id.hex(), n_ranges = dirty_ranges.len()))]
749 async fn get_file_chunk_hashes(
750 &self,
751 file_id: &MerkleHash,
752 dirty_ranges: Vec<FileRange>,
753 ) -> Result<FileChunkHashesResponse> {
754 if dirty_ranges.is_empty() {
755 return Err(ClientError::Other("get_file_chunk_hashes requires at least one dirty range".into()));
756 }
757
758 let url = Url::parse(&format!("{}/v2/file-chunk-hashes/{}", self.endpoint, file_id.hex()))?;
759
760 let header_value = HeaderValue::from_str(&format!(
763 "bytes={}",
764 dirty_ranges
765 .iter()
766 .copied()
767 .map(HttpRange::from)
768 .map(|r| r.to_string())
769 .collect::<Vec<_>>()
770 .join(",")
771 ))
772 .map_err(|err| ClientError::Other(format!("invalid X-Range-Dirty header value: {err}")))?;
773
774 let api_tag = "cas::get_file_chunk_hashes";
775 let client = self.authenticated_http_client.clone();
776
777 let response: FileChunkHashesResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
778 .run_and_extract_json(move || {
779 client
780 .get(url.clone())
781 .header(X_RANGE_DIRTY_HEADER, header_value.clone())
782 .with_extension(Api(api_tag))
783 .send()
784 })
785 .await?;
786
787 Ok(response)
788 }
789}
790
791#[cfg(test)]
792#[cfg(not(target_family = "wasm"))]
793mod tests {
794 use tracing_test::traced_test;
795 use xet_core_structures::xorb_object::CompressionScheme;
796 use xet_core_structures::xorb_object::xorb_format_test_utils::{
797 ChunkSize, build_and_verify_xorb_object, build_raw_xorb,
798 };
799
800 use super::*;
801
802 #[test]
803 fn test_clients_share_controllers_per_ctx_and_endpoint() {
804 let ctx = XetContext::default().unwrap();
805 let c1 = RemoteClient::new(ctx.clone(), "https://cas-a.example.com", &None, "", false, None);
806 let c2 = RemoteClient::new(ctx.clone(), "https://cas-a.example.com", &None, "", false, None);
807
808 assert!(Arc::ptr_eq(&c1.upload_concurrency_controller, &c2.upload_concurrency_controller));
810 assert!(Arc::ptr_eq(&c1.download_concurrency_controller, &c2.download_concurrency_controller));
811
812 let c3 = RemoteClient::new(ctx.clone(), "https://cas-b.example.com", &None, "", false, None);
814 assert!(!Arc::ptr_eq(&c1.upload_concurrency_controller, &c3.upload_concurrency_controller));
815
816 let c5 = RemoteClient::new(ctx.clone(), "https://cas-a.example.com", &None, "", false, None);
818 assert!(Arc::ptr_eq(&c1.upload_concurrency_controller, &c5.upload_concurrency_controller));
819
820 let ctx2 = XetContext::default().unwrap();
822 let c4 = RemoteClient::new(ctx2, "https://cas-a.example.com", &None, "", false, None);
823 assert!(!Arc::ptr_eq(&c1.upload_concurrency_controller, &c4.upload_concurrency_controller));
824 }
825
826 #[ignore = "requires a running CAS server"]
827 #[traced_test]
828 #[test]
829 fn test_basic_put() {
830 let prefix = PREFIX_DEFAULT;
832 let raw_xorb = build_raw_xorb(3, ChunkSize::Random(512, 10248));
833
834 let ctx = XetContext::default().unwrap();
835 let client = RemoteClient::new(ctx.clone(), CAS_ENDPOINT, &None, "", false, None);
836
837 let xorb_obj = build_and_verify_xorb_object(raw_xorb, CompressionScheme::LZ4);
838
839 let result = ctx
841 .runtime
842 .bridge_sync(async move {
843 let permit = client.acquire_upload_permit().await.unwrap();
844 client.upload_xorb(prefix, xorb_obj, None, permit).await
845 })
846 .unwrap();
847
848 assert!(result.is_ok());
850 }
851}