shardline_server/
backend.rs

1use std::sync::{
2    Arc, LazyLock, Mutex,
3    atomic::{AtomicUsize, Ordering},
4};
5use std::{num::NonZeroUsize, path::PathBuf};
6
7use axum::body::Bytes;
8use sha2::{Digest, Sha256};
9use shardline_protocol::{ByteRange, RepositoryScope, ShardlineHash};
10use shardline_storage::{
11    BeginMultipartUploadResult, DeleteOutcome, ObjectBody, ObjectIntegrity, ObjectKey,
12    ObjectMetadata, ObjectPrefix, ObjectStore, PutOutcome,
13};
14use tokio::sync::{Mutex as AsyncMutex, OwnedMutexGuard};
15
16use crate::{
17    LocalBackend, ObjectStorageAdapter, PostgresBackend, ServerConfig, ServerError,
18    ShardMetadataLimits,
19    download_stream::ServerByteStream,
20    model::{ServerStatsResponse, UploadFileResponse},
21    object_store::{ServerObjectStore, object_store_from_config},
22    overflow::checked_add,
23    protocol_support::shared_sha256_object_key,
24    reconstruction_cache::{
25        ReconstructionCacheBenchReport, benchmark_memory_reconstruction_cache_with_loader,
26    },
27    upload_ingest::{RequestBodyReader, read_body_to_bytes},
28    xet_adapter::{FileReconstructionResponse, ShardUploadResponse, XorbUploadResponse},
29};
30
31#[derive(Debug, Clone)]
32pub enum ServerBackend {
33    Local(LocalBackend),
34    Postgres(PostgresBackend),
35}
36
37/// Public benchmark-facing backend wrapper that resolves the active metadata and object
38/// adapters without exposing the private server runtime enum.
39#[derive(Debug, Clone)]
40pub struct BenchmarkBackend {
41    backend: ServerBackend,
42}
43
44static REPOSITORY_REFERENCE_PROBE_COUNT: AtomicUsize = AtomicUsize::new(0);
45static REPOSITORY_REFERENCE_PROBE_FILTER: LazyLock<Mutex<Option<String>>> =
46    LazyLock::new(|| Mutex::new(None));
47static REPOSITORY_REFERENCE_PROBE_TEST_LOCK: LazyLock<Arc<AsyncMutex<()>>> =
48    LazyLock::new(|| Arc::new(AsyncMutex::new(())));
49
50impl ServerBackend {
51    pub(crate) async fn from_config(config: &ServerConfig) -> Result<Self, ServerError> {
52        let object_store = object_store_from_config(config)?;
53        if let Some(index_postgres_url) = config.index_postgres_url() {
54            let backend =
55                PostgresBackend::new_with_object_store_and_upload_parallelism_with_frontends(
56                    config.root_dir().to_path_buf(),
57                    config.public_base_url().to_owned(),
58                    config.chunk_size(),
59                    config.upload_max_in_flight_chunks(),
60                    index_postgres_url,
61                    object_store,
62                    config.server_frontends(),
63                )
64                .await?;
65            return Ok(Self::Postgres(backend));
66        }
67
68        let backend = LocalBackend::new_with_object_store_and_upload_parallelism_with_frontends(
69            config.root_dir().to_path_buf(),
70            config.public_base_url().to_owned(),
71            config.chunk_size(),
72            config.upload_max_in_flight_chunks(),
73            object_store,
74            config.server_frontends(),
75        )
76        .await?;
77        Ok(Self::Local(backend))
78    }
79
80    pub(crate) async fn upload_xorb_stream(
81        &self,
82        expected_hash: &str,
83        body: RequestBodyReader,
84    ) -> Result<XorbUploadResponse, ServerError> {
85        match self {
86            Self::Local(backend) => backend.upload_xorb_stream(expected_hash, body).await,
87            Self::Postgres(backend) => backend.upload_xorb_stream(expected_hash, body).await,
88        }
89    }
90
91    pub(crate) async fn upload_shard_stream(
92        &self,
93        body: RequestBodyReader,
94        repository_scope: Option<&RepositoryScope>,
95        shard_metadata_limits: ShardMetadataLimits,
96    ) -> Result<ShardUploadResponse, ServerError> {
97        match self {
98            Self::Local(backend) => {
99                backend
100                    .upload_shard_stream(body, repository_scope, shard_metadata_limits)
101                    .await
102            }
103            Self::Postgres(backend) => {
104                backend
105                    .upload_shard_stream(body, repository_scope, shard_metadata_limits)
106                    .await
107            }
108        }
109    }
110
111    pub(crate) async fn put_sha256_addressed_object_stream_if_absent(
112        &self,
113        object_key: &ObjectKey,
114        digest_hex: &str,
115        body: RequestBodyReader,
116    ) -> Result<PutOutcome, ServerError> {
117        match self {
118            Self::Local(backend) => {
119                put_sha256_addressed_object_stream_if_absent_with_object_store(
120                    &backend.object_store(),
121                    object_key,
122                    digest_hex,
123                    body,
124                )
125                .await
126            }
127            Self::Postgres(backend) => {
128                put_sha256_addressed_object_stream_if_absent_with_object_store(
129                    &backend.object_store(),
130                    object_key,
131                    digest_hex,
132                    body,
133                )
134                .await
135            }
136        }
137    }
138
139    pub(crate) async fn reconstruction(
140        &self,
141        file_id: &str,
142        content_hash: Option<&str>,
143        requested_range: Option<ByteRange>,
144        repository_scope: Option<&RepositoryScope>,
145    ) -> Result<FileReconstructionResponse, ServerError> {
146        match self {
147            Self::Local(backend) => {
148                backend
149                    .reconstruction(file_id, content_hash, requested_range, repository_scope)
150                    .await
151            }
152            Self::Postgres(backend) => {
153                backend
154                    .reconstruction(file_id, content_hash, requested_range, repository_scope)
155                    .await
156            }
157        }
158    }
159
160    pub(crate) async fn file_total_bytes(
161        &self,
162        file_id: &str,
163        content_hash: Option<&str>,
164        repository_scope: Option<&RepositoryScope>,
165    ) -> Result<u64, ServerError> {
166        match self {
167            Self::Local(backend) => {
168                backend
169                    .file_total_bytes(file_id, content_hash, repository_scope)
170                    .await
171            }
172            Self::Postgres(backend) => {
173                backend
174                    .file_total_bytes(file_id, content_hash, repository_scope)
175                    .await
176            }
177        }
178    }
179
180    pub(crate) async fn chunk_length(&self, hash_hex: &str) -> Result<u64, ServerError> {
181        match self {
182            Self::Local(backend) => backend.chunk_length(hash_hex).await,
183            Self::Postgres(backend) => backend.chunk_length(hash_hex).await,
184        }
185    }
186
187    pub(crate) async fn dedupe_shard_length(&self, hash_hex: &str) -> Result<u64, ServerError> {
188        match self {
189            Self::Local(backend) => backend.dedupe_shard_length(hash_hex).await,
190            Self::Postgres(backend) => backend.dedupe_shard_length(hash_hex).await,
191        }
192    }
193
194    pub(crate) async fn xorb_length(&self, hash_hex: &str) -> Result<u64, ServerError> {
195        match self {
196            Self::Local(backend) => backend.xorb_length(hash_hex).await,
197            Self::Postgres(backend) => backend.xorb_length(hash_hex).await,
198        }
199    }
200
201    pub(crate) async fn read_xorb_range_stream(
202        &self,
203        hash_hex: &str,
204        total_length: u64,
205        range: ByteRange,
206    ) -> Result<ServerByteStream, ServerError> {
207        match self {
208            Self::Local(backend) => {
209                backend
210                    .read_xorb_range_stream(hash_hex, total_length, range)
211                    .await
212            }
213            Self::Postgres(backend) => {
214                backend
215                    .read_xorb_range_stream(hash_hex, total_length, range)
216                    .await
217            }
218        }
219    }
220
221    pub(crate) async fn read_dedupe_shard_stream(
222        &self,
223        hash_hex: &str,
224    ) -> Result<(ServerByteStream, u64), ServerError> {
225        match self {
226            Self::Local(backend) => backend.read_dedupe_shard_stream(hash_hex).await,
227            Self::Postgres(backend) => backend.read_dedupe_shard_stream(hash_hex).await,
228        }
229    }
230
231    pub(crate) async fn repository_references_xorb(
232        &self,
233        hash_hex: &str,
234        repository_scope: &RepositoryScope,
235    ) -> Result<bool, ServerError> {
236        count_repository_reference_probe_for_tests(hash_hex);
237        match self {
238            Self::Local(backend) => {
239                backend
240                    .repository_references_xorb(hash_hex, repository_scope)
241                    .await
242            }
243            Self::Postgres(backend) => {
244                backend
245                    .repository_references_xorb(hash_hex, repository_scope)
246                    .await
247            }
248        }
249    }
250
251    pub(crate) async fn stats(&self) -> Result<ServerStatsResponse, ServerError> {
252        match self {
253            Self::Local(backend) => backend.stats().await,
254            Self::Postgres(backend) => backend.stats().await,
255        }
256    }
257
258    pub(crate) async fn ready(&self) -> Result<(), ServerError> {
259        match self {
260            Self::Local(backend) => backend.ready().await,
261            Self::Postgres(backend) => backend.ready().await,
262        }
263    }
264
265    pub(crate) const fn backend_name(&self) -> &'static str {
266        match self {
267            Self::Local(_backend) => "local",
268            Self::Postgres(_backend) => "postgres",
269        }
270    }
271
272    pub(crate) const fn object_backend_name(&self) -> &'static str {
273        match self {
274            Self::Local(backend) => backend.object_backend_name(),
275            Self::Postgres(backend) => backend.object_backend_name(),
276        }
277    }
278
279    pub(crate) fn uses_s3_object_store(&self) -> bool {
280        match self {
281            Self::Local(backend) => matches!(backend.object_store(), ServerObjectStore::S3(_)),
282            Self::Postgres(backend) => matches!(backend.object_store(), ServerObjectStore::S3(_)),
283        }
284    }
285
286    pub(crate) fn put_object_bytes_if_absent(
287        &self,
288        object_key: &ObjectKey,
289        bytes: Vec<u8>,
290    ) -> Result<PutOutcome, ServerError> {
291        match self {
292            Self::Local(backend) => backend.put_object_bytes_if_absent(object_key, bytes),
293            Self::Postgres(backend) => backend.put_object_bytes_if_absent(object_key, bytes),
294        }
295    }
296
297    pub(crate) fn put_sha256_addressed_object_bytes_if_absent(
298        &self,
299        object_key: &ObjectKey,
300        digest_hex: &str,
301        bytes: Vec<u8>,
302    ) -> Result<PutOutcome, ServerError> {
303        match self {
304            Self::Local(backend) => {
305                backend.put_sha256_addressed_object_bytes_if_absent(object_key, digest_hex, bytes)
306            }
307            Self::Postgres(backend) => {
308                backend.put_sha256_addressed_object_bytes_if_absent(object_key, digest_hex, bytes)
309            }
310        }
311    }
312
313    pub(crate) fn copy_object_if_absent(
314        &self,
315        source: &ObjectKey,
316        destination: &ObjectKey,
317    ) -> Result<PutOutcome, ServerError> {
318        match self {
319            Self::Local(backend) => backend.copy_object_if_absent(source, destination),
320            Self::Postgres(backend) => backend.copy_object_if_absent(source, destination),
321        }
322    }
323
324    pub(crate) fn put_object_bytes_overwrite(
325        &self,
326        object_key: &ObjectKey,
327        bytes: Vec<u8>,
328    ) -> Result<(), ServerError> {
329        match self {
330            Self::Local(backend) => backend.put_object_bytes_overwrite(object_key, bytes),
331            Self::Postgres(backend) => backend.put_object_bytes_overwrite(object_key, bytes),
332        }
333    }
334
335    pub(crate) fn put_sha256_addressed_object_file(
336        &self,
337        object_key: &ObjectKey,
338        digest_hex: &str,
339        path: &std::path::Path,
340        integrity: &shardline_storage::ObjectIntegrity,
341    ) -> Result<PutOutcome, ServerError> {
342        match self {
343            Self::Local(backend) => {
344                backend.put_sha256_addressed_object_file(object_key, digest_hex, path, integrity)
345            }
346            Self::Postgres(backend) => {
347                backend.put_sha256_addressed_object_file(object_key, digest_hex, path, integrity)
348            }
349        }
350    }
351
352    pub(crate) async fn object_length(&self, object_key: &ObjectKey) -> Result<u64, ServerError> {
353        match self {
354            Self::Local(backend) => backend.object_length(object_key).await,
355            Self::Postgres(backend) => backend.object_length(object_key).await,
356        }
357    }
358
359    pub(crate) async fn read_object(&self, object_key: &ObjectKey) -> Result<Vec<u8>, ServerError> {
360        match self {
361            Self::Local(backend) => backend.read_object(object_key).await,
362            Self::Postgres(backend) => backend.read_object(object_key).await,
363        }
364    }
365
366    pub(crate) async fn read_object_stream(
367        &self,
368        object_key: &ObjectKey,
369        total_length: u64,
370        range: Option<ByteRange>,
371    ) -> Result<ServerByteStream, ServerError> {
372        match self {
373            Self::Local(backend) => {
374                backend
375                    .read_object_stream(object_key, total_length, range)
376                    .await
377            }
378            Self::Postgres(backend) => {
379                backend
380                    .read_object_stream(object_key, total_length, range)
381                    .await
382            }
383        }
384    }
385
386    pub(crate) fn visit_object_prefix<Visitor>(
387        &self,
388        prefix: &ObjectPrefix,
389        visitor: Visitor,
390    ) -> Result<(), ServerError>
391    where
392        Visitor: FnMut(ObjectMetadata) -> Result<(), ServerError>,
393    {
394        match self {
395            Self::Local(backend) => backend.visit_object_prefix(prefix, visitor),
396            Self::Postgres(backend) => backend.visit_object_prefix(prefix, visitor),
397        }
398    }
399
400    pub(crate) fn list_object_flat_namespace_page(
401        &self,
402        prefix: &ObjectPrefix,
403        start_after: Option<&ObjectKey>,
404        limit: usize,
405    ) -> Result<Vec<ObjectMetadata>, ServerError> {
406        match self {
407            Self::Local(backend) => {
408                backend.list_object_flat_namespace_page(prefix, start_after, limit)
409            }
410            Self::Postgres(backend) => {
411                backend.list_object_flat_namespace_page(prefix, start_after, limit)
412            }
413        }
414    }
415
416    pub(crate) async fn delete_object_if_present(
417        &self,
418        object_key: &ObjectKey,
419    ) -> Result<DeleteOutcome, ServerError> {
420        match self {
421            Self::Local(backend) => backend.delete_object_if_present(object_key).await,
422            Self::Postgres(backend) => backend.delete_object_if_present(object_key).await,
423        }
424    }
425}
426
427impl shardline_oci_adapter::OciBackend for ServerBackend {
428    async fn create_resumable_object_upload(
429        &self,
430        object_key: &ObjectKey,
431    ) -> Result<Option<String>, shardline_oci_adapter::OciAdapterError> {
432        match self {
433            Self::Local(backend) => match backend.object_store() {
434                ServerObjectStore::S3(store) => Ok(Some(
435                    store
436                        .create_resumable_upload(object_key)
437                        .await
438                        .map_err(shardline_oci_adapter::OciAdapterError::from)?,
439                )),
440                ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => Ok(None),
441            },
442            Self::Postgres(backend) => match backend.object_store() {
443                ServerObjectStore::S3(store) => Ok(Some(
444                    store
445                        .create_resumable_upload(object_key)
446                        .await
447                        .map_err(shardline_oci_adapter::OciAdapterError::from)?,
448                )),
449                ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => Ok(None),
450            },
451        }
452    }
453
454    async fn upload_resumable_object_part(
455        &self,
456        object_key: &ObjectKey,
457        upload_id: &str,
458        part_idx: usize,
459        bytes: Bytes,
460    ) -> Result<String, shardline_oci_adapter::OciAdapterError> {
461        match self {
462            Self::Local(backend) => match backend.object_store() {
463                ServerObjectStore::S3(store) => store
464                    .upload_resumable_part(object_key, upload_id, part_idx, bytes)
465                    .await
466                    .map_err(shardline_oci_adapter::OciAdapterError::from),
467                ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => {
468                    Err(shardline_oci_adapter::OciAdapterError::NotFound)
469                }
470            },
471            Self::Postgres(backend) => match backend.object_store() {
472                ServerObjectStore::S3(store) => store
473                    .upload_resumable_part(object_key, upload_id, part_idx, bytes)
474                    .await
475                    .map_err(shardline_oci_adapter::OciAdapterError::from),
476                ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => {
477                    Err(shardline_oci_adapter::OciAdapterError::NotFound)
478                }
479            },
480        }
481    }
482
483    async fn complete_resumable_object_upload(
484        &self,
485        object_key: &ObjectKey,
486        upload_id: &str,
487        part_ids: Vec<String>,
488    ) -> Result<(), shardline_oci_adapter::OciAdapterError> {
489        match self {
490            Self::Local(backend) => match backend.object_store() {
491                ServerObjectStore::S3(store) => store
492                    .complete_resumable_upload(object_key, upload_id, part_ids)
493                    .await
494                    .map_err(shardline_oci_adapter::OciAdapterError::from),
495                ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => {
496                    Err(shardline_oci_adapter::OciAdapterError::NotFound)
497                }
498            },
499            Self::Postgres(backend) => match backend.object_store() {
500                ServerObjectStore::S3(store) => store
501                    .complete_resumable_upload(object_key, upload_id, part_ids)
502                    .await
503                    .map_err(shardline_oci_adapter::OciAdapterError::from),
504                ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => {
505                    Err(shardline_oci_adapter::OciAdapterError::NotFound)
506                }
507            },
508        }
509    }
510
511    async fn abort_resumable_object_upload(
512        &self,
513        object_key: &ObjectKey,
514        upload_id: &str,
515    ) -> Result<(), shardline_oci_adapter::OciAdapterError> {
516        match self {
517            Self::Local(backend) => match backend.object_store() {
518                ServerObjectStore::S3(store) => store
519                    .abort_resumable_upload(object_key, upload_id)
520                    .await
521                    .map_err(shardline_oci_adapter::OciAdapterError::from),
522                ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => Ok(()),
523            },
524            Self::Postgres(backend) => match backend.object_store() {
525                ServerObjectStore::S3(store) => store
526                    .abort_resumable_upload(object_key, upload_id)
527                    .await
528                    .map_err(shardline_oci_adapter::OciAdapterError::from),
529                ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => Ok(()),
530            },
531        }
532    }
533
534    fn put_sha256_addressed_object_bytes_if_absent(
535        &self,
536        object_key: &ObjectKey,
537        digest_hex: &str,
538        bytes: Vec<u8>,
539    ) -> Result<PutOutcome, shardline_oci_adapter::OciAdapterError> {
540        match self {
541            Self::Local(backend) => backend
542                .put_sha256_addressed_object_bytes_if_absent(object_key, digest_hex, bytes)
543                .map_err(server_error_to_oci),
544            Self::Postgres(backend) => backend
545                .put_sha256_addressed_object_bytes_if_absent(object_key, digest_hex, bytes)
546                .map_err(server_error_to_oci),
547        }
548    }
549
550    fn copy_object_if_absent(
551        &self,
552        source: &ObjectKey,
553        destination: &ObjectKey,
554    ) -> Result<PutOutcome, shardline_oci_adapter::OciAdapterError> {
555        match self {
556            Self::Local(backend) => backend
557                .copy_object_if_absent(source, destination)
558                .map_err(server_error_to_oci),
559            Self::Postgres(backend) => backend
560                .copy_object_if_absent(source, destination)
561                .map_err(server_error_to_oci),
562        }
563    }
564
565    async fn delete_object_if_present(
566        &self,
567        object_key: &ObjectKey,
568    ) -> Result<DeleteOutcome, shardline_oci_adapter::OciAdapterError> {
569        match self {
570            Self::Local(backend) => backend
571                .delete_object_if_present(object_key)
572                .await
573                .map_err(server_error_to_oci),
574            Self::Postgres(backend) => backend
575                .delete_object_if_present(object_key)
576                .await
577                .map_err(server_error_to_oci),
578        }
579    }
580}
581
582fn server_error_to_oci(error: ServerError) -> shardline_oci_adapter::OciAdapterError {
583    use crate::error::ObjectStoreError;
584    use shardline_oci_adapter::OciAdapterError;
585    match error {
586        ServerError::Io(e) => OciAdapterError::Io(e),
587        ServerError::Json(e) => OciAdapterError::Json(e),
588        ServerError::NumericConversion(e) => OciAdapterError::NumericConversion(e),
589        ServerError::ObjectStore(ObjectStoreError::Local(e)) => {
590            OciAdapterError::LocalObjectStore(e)
591        }
592        ServerError::ObjectStore(ObjectStoreError::S3(e)) => OciAdapterError::S3ObjectStore(e),
593        ServerError::ObjectStore(ObjectStoreError::Prefix(e)) => OciAdapterError::ObjectPrefix(e),
594        ServerError::NotFound => OciAdapterError::NotFound,
595        ServerError::Overflow => OciAdapterError::Overflow,
596        ServerError::InvalidContentHash => OciAdapterError::InvalidContentHash,
597        ServerError::InvalidDigest => OciAdapterError::InvalidDigest,
598        ServerError::InvalidRepositoryName => OciAdapterError::InvalidRepositoryName,
599        ServerError::InvalidManifestReference => OciAdapterError::InvalidManifestReference,
600        ServerError::InvalidUploadSession => OciAdapterError::InvalidUploadSession,
601        ServerError::TooManyUploadSessions => OciAdapterError::TooManyUploadSessions,
602        ServerError::ExpectedBodyHashMismatch => OciAdapterError::ExpectedBodyHashMismatch,
603        ServerError::BlockingTask(e) => OciAdapterError::BlockingTask(e),
604        ref other @ (ServerError::RequestBodyRead(_)
605        | ServerError::RequestBodyTooLarge
606        | ServerError::RequestQueryTooLarge
607        | ServerError::RequestBodyFrameOutOfBounds
608        | ServerError::HashParse(_)
609        | ServerError::ObjectStore(
610            ObjectStoreError::MissingS3Config
611            | ObjectStoreError::StoredLengthMismatch
612            | ObjectStoreError::MigrationSourceHashMismatch { .. },
613        )
614        | ServerError::Index(_)
615        | ServerError::StoredFileMetadataTooLarge { .. }
616        | ServerError::StoredFileMetadataLengthMismatch
617        | ServerError::InvalidFileId
618        | ServerError::InvalidXorbPrefix
619        | ServerError::XorbHashMismatch
620        | ServerError::InvalidSerializedXorb
621        | ServerError::InvalidSerializedShard(_)
622        | ServerError::MissingReferencedXorb
623        | ServerError::TooManyShardTerms
624        | ServerError::TooManyBatchReconstructionFileIds
625        | ServerError::InvalidRangeHeader
626        | ServerError::RangeNotSatisfiable
627        | ServerError::MissingAuthorization
628        | ServerError::InvalidAuthorizationHeader
629        | ServerError::InvalidToken(_)
630        | ServerError::InsufficientScope
631        | ServerError::ProviderTokensDisabled
632        | ServerError::MissingProviderApiKey
633        | ServerError::InvalidProviderApiKey
634        | ServerError::MissingProviderSubject
635        | ServerError::InvalidProviderTokenRequest
636        | ServerError::MissingProviderWebhookAuthentication
637        | ServerError::InvalidProviderWebhookAuthentication
638        | ServerError::InvalidProviderWebhookPayload
639        | ServerError::UnknownProvider
640        | ServerError::ProviderDenied
641        | ServerError::Provider(_)
642        | ServerError::ReconstructionCache(_)
643        | ServerError::Config(_)
644        | ServerError::NotAcceptable
645        | ServerError::UnauthorizedChallenge(_)
646        | ServerError::TooManyRegistryTokenRequests
647        | ServerError::MissingReconstructionCacheRedisUrl
648        | ServerError::TransferLimiterClosed) => {
649            OciAdapterError::Io(std::io::Error::other(other.to_string()))
650        }
651    }
652}
653
654async fn put_sha256_addressed_object_stream_if_absent_with_object_store(
655    object_store: &ServerObjectStore,
656    object_key: &ObjectKey,
657    digest_hex: &str,
658    mut body: RequestBodyReader,
659) -> Result<PutOutcome, ServerError> {
660    let canonical_key = shared_sha256_object_key(digest_hex)?;
661    match object_store {
662        ServerObjectStore::S3(store) => {
663            let canonical_outcome =
664                match store.begin_content_addressed_upload(&canonical_key).await? {
665                    BeginMultipartUploadResult::AlreadyExists => PutOutcome::AlreadyExists,
666                    BeginMultipartUploadResult::Upload(mut upload) => {
667                        let mut sha256 = Sha256::new();
668                        let mut total_length = 0_u64;
669                        while let Some(bytes) = body.next_bytes().await? {
670                            sha256.update(&bytes);
671                            total_length = checked_add(total_length, u64::try_from(bytes.len())?)?;
672                            upload.write(&bytes);
673                            if let Err(error) = upload.wait_for_capacity(4).await {
674                                let _ignored = upload.abort().await;
675                                return Err(ServerError::from(error));
676                            }
677                        }
678                        let observed = hex::encode(sha256.finalize());
679                        if observed != digest_hex {
680                            let _ignored = upload.abort().await;
681                            return Err(ServerError::ExpectedBodyHashMismatch);
682                        }
683                        let _total_length = total_length;
684                        upload.finish().await?;
685                        PutOutcome::Inserted
686                    }
687                };
688            if canonical_key == *object_key {
689                return Ok(canonical_outcome);
690            }
691            Ok(object_store.copy_if_absent(&canonical_key, object_key)?)
692        }
693        ServerObjectStore::Local(_) | ServerObjectStore::Blackhole => {
694            let bytes = read_body_to_bytes(&mut body).await?;
695            let observed = hex::encode(Sha256::digest(&bytes));
696            if observed != digest_hex {
697                return Err(ServerError::ExpectedBodyHashMismatch);
698            }
699            let integrity = ObjectIntegrity::new(
700                ShardlineHash::from_bytes(*blake3::hash(&bytes).as_bytes()),
701                u64::try_from(bytes.len())?,
702            );
703            let canonical_outcome = object_store.put_if_absent(
704                &canonical_key,
705                ObjectBody::from_vec(bytes),
706                &integrity,
707            )?;
708            if canonical_key == *object_key {
709                return Ok(canonical_outcome);
710            }
711            Ok(object_store.copy_if_absent(&canonical_key, object_key)?)
712        }
713    }
714}
715
716impl BenchmarkBackend {
717    /// Creates an isolated local benchmark backend rooted at the supplied directory.
718    ///
719    /// # Errors
720    ///
721    /// Returns [`ServerError`] when local metadata or object storage cannot initialize.
722    pub async fn isolated_local(
723        root: PathBuf,
724        public_base_url: String,
725        chunk_size: NonZeroUsize,
726        upload_max_in_flight_chunks: NonZeroUsize,
727    ) -> Result<Self, ServerError> {
728        let backend = LocalBackend::new_with_upload_parallelism(
729            root,
730            public_base_url,
731            chunk_size,
732            upload_max_in_flight_chunks,
733        )
734        .await?;
735        Ok(Self {
736            backend: ServerBackend::Local(backend),
737        })
738    }
739
740    /// Creates a benchmark backend from the effective runtime configuration.
741    ///
742    /// The benchmark namespace is appended to any configured S3 key prefix so benchmark
743    /// objects cannot collide with non-benchmark data.
744    ///
745    /// # Errors
746    ///
747    /// Returns [`ServerError`] when the configured metadata or object adapters cannot
748    /// initialize.
749    pub async fn from_config(
750        config: &ServerConfig,
751        root: PathBuf,
752        benchmark_namespace: &str,
753    ) -> Result<Self, ServerError> {
754        let mut configured = config.clone().with_root_dir(root);
755        if configured.object_storage_adapter() == ObjectStorageAdapter::S3
756            && let Some(s3_config) = configured.s3_object_store_config().cloned()
757        {
758            let key_prefix =
759                compose_benchmark_object_key_prefix(s3_config.key_prefix(), benchmark_namespace);
760            configured = configured.with_object_storage(
761                ObjectStorageAdapter::S3,
762                Some(s3_config.with_key_prefix(Some(&key_prefix))),
763            );
764        }
765
766        let backend = ServerBackend::from_config(&configured).await?;
767        Ok(Self { backend })
768    }
769
770    /// Stores one logical file version.
771    ///
772    /// # Errors
773    ///
774    /// Returns [`ServerError`] when validation, object persistence, or metadata updates
775    /// fail.
776    pub async fn upload_file(
777        &self,
778        file_id: &str,
779        body: Bytes,
780        repository_scope: Option<&RepositoryScope>,
781    ) -> Result<UploadFileResponse, ServerError> {
782        match &self.backend {
783            ServerBackend::Local(backend) => {
784                backend.upload_file(file_id, body, repository_scope).await
785            }
786            ServerBackend::Postgres(backend) => {
787                backend.upload_file(file_id, body, repository_scope).await
788            }
789        }
790    }
791
792    /// Reconstructs one logical file or range.
793    ///
794    /// # Errors
795    ///
796    /// Returns [`ServerError`] when the requested version is missing or any referenced
797    /// object cannot be loaded.
798    pub async fn reconstruction(
799        &self,
800        file_id: &str,
801        content_hash: Option<&str>,
802        requested_range: Option<ByteRange>,
803        repository_scope: Option<&RepositoryScope>,
804    ) -> Result<FileReconstructionResponse, ServerError> {
805        self.backend
806            .reconstruction(file_id, content_hash, requested_range, repository_scope)
807            .await
808    }
809
810    /// Downloads one logical file version into full bytes.
811    ///
812    /// # Errors
813    ///
814    /// Returns [`ServerError`] when the requested version is missing or any referenced
815    /// object cannot be loaded.
816    pub async fn download_file(
817        &self,
818        file_id: &str,
819        content_hash: Option<&str>,
820        repository_scope: Option<&RepositoryScope>,
821    ) -> Result<Vec<u8>, ServerError> {
822        match &self.backend {
823            ServerBackend::Local(backend) => {
824                backend
825                    .download_file(file_id, content_hash, repository_scope)
826                    .await
827            }
828            ServerBackend::Postgres(backend) => {
829                backend
830                    .download_file(file_id, content_hash, repository_scope)
831                    .await
832            }
833        }
834    }
835
836    /// Returns storage statistics from the active backend.
837    ///
838    /// # Errors
839    ///
840    /// Returns [`ServerError`] when metadata or object storage cannot be traversed.
841    pub async fn stats(&self) -> Result<ServerStatsResponse, ServerError> {
842        self.backend.stats().await
843    }
844
845    /// Measures one cold reconstruction lookup followed by one hot memory-cache hit.
846    ///
847    /// # Errors
848    ///
849    /// Returns [`ServerError`] when reconstruction loading or response serialization
850    /// fails.
851    pub async fn benchmark_memory_reconstruction_cache(
852        &self,
853        file_id: &str,
854        content_hash: &str,
855        repository_scope: Option<&RepositoryScope>,
856    ) -> Result<ReconstructionCacheBenchReport, ServerError> {
857        benchmark_memory_reconstruction_cache_with_loader(
858            file_id,
859            content_hash,
860            repository_scope,
861            || async move {
862                self.reconstruction(file_id, Some(content_hash), None, repository_scope)
863                    .await
864            },
865        )
866        .await
867    }
868
869    /// Returns the metadata backend name used by this benchmark backend.
870    #[must_use]
871    pub const fn metadata_backend_name(&self) -> &'static str {
872        self.backend.backend_name()
873    }
874
875    /// Returns the immutable object-storage backend name used by this benchmark backend.
876    #[must_use]
877    pub const fn object_backend_name(&self) -> &'static str {
878        self.backend.object_backend_name()
879    }
880}
881
882fn compose_benchmark_object_key_prefix(
883    existing_key_prefix: Option<&str>,
884    benchmark_namespace: &str,
885) -> String {
886    match existing_key_prefix {
887        Some(existing_key_prefix) if !existing_key_prefix.is_empty() => {
888            format!("{existing_key_prefix}/bench/{benchmark_namespace}")
889        }
890        Some(_existing_key_prefix) => format!("bench/{benchmark_namespace}"),
891        None => format!("bench/{benchmark_namespace}"),
892    }
893}
894
895pub fn reset_repository_reference_probe_count_for_hash(hash_hex: &str) {
896    REPOSITORY_REFERENCE_PROBE_COUNT.store(0, Ordering::Relaxed);
897    let filter = REPOSITORY_REFERENCE_PROBE_FILTER.lock();
898    match filter {
899        Ok(mut filter) => *filter = Some(hash_hex.to_owned()),
900        Err(poisoned) => *poisoned.into_inner() = Some(hash_hex.to_owned()),
901    }
902}
903
904pub fn repository_reference_probe_count() -> usize {
905    REPOSITORY_REFERENCE_PROBE_COUNT.load(Ordering::Relaxed)
906}
907
908pub fn clear_repository_reference_probe_filter() {
909    let filter = REPOSITORY_REFERENCE_PROBE_FILTER.lock();
910    match filter {
911        Ok(mut filter) => *filter = None,
912        Err(poisoned) => *poisoned.into_inner() = None,
913    }
914}
915
916pub async fn lock_repository_reference_probe_test() -> OwnedMutexGuard<()> {
917    REPOSITORY_REFERENCE_PROBE_TEST_LOCK
918        .clone()
919        .lock_owned()
920        .await
921}
922
923fn count_repository_reference_probe_for_tests(hash_hex: &str) {
924    let filter = REPOSITORY_REFERENCE_PROBE_FILTER.lock();
925    let matches_filter = match filter {
926        Ok(filter) => filter
927            .as_deref()
928            .is_none_or(|expected| expected == hash_hex),
929        Err(poisoned) => poisoned
930            .into_inner()
931            .as_deref()
932            .is_none_or(|expected| expected == hash_hex),
933    };
934
935    if matches_filter {
936        REPOSITORY_REFERENCE_PROBE_COUNT.fetch_add(1, Ordering::Relaxed);
937    }
938}
939
940#[cfg(test)]
941mod tests {
942    use std::{num::NonZeroUsize, path::PathBuf};
943
944    use super::{BenchmarkBackend, compose_benchmark_object_key_prefix};
945    use crate::ServerConfig;
946
947    #[test]
948    fn benchmark_object_key_prefix_appends_namespace() {
949        assert_eq!(
950            compose_benchmark_object_key_prefix(Some("tenant-a"), "run-0001"),
951            "tenant-a/bench/run-0001"
952        );
953        assert_eq!(
954            compose_benchmark_object_key_prefix(None, "run-0001"),
955            "bench/run-0001"
956        );
957    }
958
959    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
960    async fn configured_benchmark_backend_uses_local_runtime_configuration() {
961        let chunk_size = NonZeroUsize::new(4).unwrap_or(NonZeroUsize::MIN);
962        let upload_budget = NonZeroUsize::new(4).unwrap_or(NonZeroUsize::MIN);
963        let bind_addr = "127.0.0.1:8080".parse();
964        assert!(bind_addr.is_ok());
965        let Ok(bind_addr) = bind_addr else {
966            return;
967        };
968        let config = ServerConfig::new(
969            bind_addr,
970            "http://127.0.0.1:8080".to_owned(),
971            PathBuf::from("/tmp/ignored"),
972            chunk_size,
973        )
974        .with_chunk_size(chunk_size)
975        .with_upload_max_in_flight_chunks(upload_budget);
976        let storage = tempfile::tempdir();
977        assert!(storage.is_ok());
978        let Ok(storage) = storage else {
979            return;
980        };
981
982        let backend =
983            BenchmarkBackend::from_config(&config, storage.path().to_path_buf(), "run-0001").await;
984        assert!(backend.is_ok());
985        let Ok(backend) = backend else {
986            return;
987        };
988
989        assert_eq!(backend.metadata_backend_name(), "local");
990        assert_eq!(backend.object_backend_name(), "local");
991    }
992}
shardline_server/backend.rs

shardline_server/
backend.rs
