kvbm_engine/object/s3/
client.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! S3-compatible object storage client for block management.
5//!
6//! This module provides an implementation of [`ObjectBlockOps`] using the AWS S3 SDK.
7//! It supports S3-compatible storage services including MinIO.
8
9use anyhow::{Result, anyhow};
10use aws_sdk_s3::Client;
11use aws_sdk_s3::error::ProvideErrorMetadata;
12use aws_sdk_s3::primitives::ByteStream;
13use bytes::Bytes;
14use futures::future::BoxFuture;
15use futures::stream::StreamExt;
16
17use crate::object::{DefaultKeyFormatter, KeyFormatter, LayoutConfigExt, ObjectBlockOps};
18use crate::{BlockId, SequenceHash};
19use kvbm_common::LogicalLayoutHandle;
20use kvbm_physical::transfer::PhysicalLayout;
21use std::sync::Arc;
22
23/// Configuration for S3 object storage client.
24#[derive(Debug, Clone)]
25pub struct S3Config {
26    /// Custom endpoint URL for S3-compatible services (e.g., MinIO).
27    /// If None, uses the default AWS S3 endpoint.
28    pub endpoint_url: Option<String>,
29
30    /// S3 bucket name for storing blocks.
31    pub bucket: String,
32
33    /// AWS region.
34    pub region: String,
35
36    /// Use path-style URLs instead of virtual-hosted-style.
37    /// Required for MinIO and some S3-compatible services.
38    pub force_path_style: bool,
39
40    /// Maximum number of concurrent S3 requests.
41    pub max_concurrent_requests: usize,
42}
43
44impl Default for S3Config {
45    /// Returns a default configuration suitable for local MinIO testing.
46    fn default() -> Self {
47        Self {
48            endpoint_url: Some("http://localhost:9000".into()),
49            bucket: "kvbm-blocks".into(),
50            region: "us-east-1".into(),
51            force_path_style: true,
52            max_concurrent_requests: 16,
53        }
54    }
55}
56
57impl S3Config {
58    /// Create a new S3Config for AWS S3 (not MinIO).
59    pub fn aws(bucket: String, region: String) -> Self {
60        Self {
61            endpoint_url: None,
62            bucket,
63            region,
64            force_path_style: false,
65            max_concurrent_requests: 16,
66        }
67    }
68
69    /// Create a new S3Config for MinIO.
70    pub fn minio(endpoint_url: String, bucket: String) -> Self {
71        Self {
72            endpoint_url: Some(endpoint_url),
73            bucket,
74            region: "us-east-1".into(),
75            force_path_style: true,
76            max_concurrent_requests: 16,
77        }
78    }
79
80    /// Create from kvbm-config's S3ObjectConfig.
81    pub fn from_object_config(config: &kvbm_config::S3ObjectConfig) -> Self {
82        Self {
83            endpoint_url: config.endpoint_url.clone(),
84            bucket: config.bucket.clone(),
85            region: config.region.clone(),
86            force_path_style: config.force_path_style,
87            max_concurrent_requests: config.max_concurrent_requests,
88        }
89    }
90
91    /// Set the maximum number of concurrent requests.
92    pub fn with_max_concurrent_requests(mut self, max: usize) -> Self {
93        self.max_concurrent_requests = max;
94        self
95    }
96}
97
98/// S3-compatible object storage client for block operations.
99///
100/// This client implements [`ObjectBlockOps`] using the AWS S3 SDK.
101/// It supports parallel block operations and uses rayon for CPU-bound memory copies.
102///
103/// # Key Formatting
104///
105/// Uses a [`KeyFormatter`] to convert `SequenceHash` to object keys. The formatter
106/// can embed rank, namespace, or other prefixes for key uniqueness across workers.
107pub struct S3ObjectBlockClient {
108    /// AWS S3 client
109    client: Client,
110
111    /// S3 configuration
112    config: S3Config,
113
114    /// Key formatter for converting SequenceHash to object keys.
115    key_formatter: Arc<dyn KeyFormatter>,
116}
117
118impl S3ObjectBlockClient {
119    /// Create a new S3ObjectBlockClient with default key formatting.
120    ///
121    /// # Arguments
122    /// * `config` - S3 configuration
123    ///
124    /// # Errors
125    /// Returns an error if the S3 client cannot be initialized.
126    pub async fn new(config: S3Config) -> Result<Self> {
127        let client = build_s3_client(&config).await?;
128        Ok(Self {
129            client,
130            config,
131            key_formatter: Arc::new(DefaultKeyFormatter),
132        })
133    }
134
135    /// Create a new S3ObjectBlockClient with a custom key formatter.
136    ///
137    /// # Arguments
138    /// * `config` - S3 configuration
139    /// * `key_formatter` - Custom key formatter for SequenceHash → String conversion
140    ///
141    /// # Errors
142    /// Returns an error if the S3 client cannot be initialized.
143    pub async fn with_key_formatter(
144        config: S3Config,
145        key_formatter: Arc<dyn KeyFormatter>,
146    ) -> Result<Self> {
147        let client = build_s3_client(&config).await?;
148        Ok(Self {
149            client,
150            config,
151            key_formatter,
152        })
153    }
154
155    /// Create from an existing AWS S3 client with default key formatting.
156    pub fn from_client(client: Client, config: S3Config) -> Self {
157        Self {
158            client,
159            config,
160            key_formatter: Arc::new(DefaultKeyFormatter),
161        }
162    }
163
164    /// Create from an existing AWS S3 client with a custom key formatter.
165    pub fn from_client_with_formatter(
166        client: Client,
167        config: S3Config,
168        key_formatter: Arc<dyn KeyFormatter>,
169    ) -> Self {
170        Self {
171            client,
172            config,
173            key_formatter,
174        }
175    }
176
177    /// Get a reference to the S3 client.
178    pub fn client(&self) -> &Client {
179        &self.client
180    }
181
182    /// Get a reference to the configuration.
183    pub fn config(&self) -> &S3Config {
184        &self.config
185    }
186
187    /// Get a reference to the key formatter.
188    pub fn key_formatter(&self) -> &Arc<dyn KeyFormatter> {
189        &self.key_formatter
190    }
191
192    /// Get a reference to the bucket name.
193    pub fn bucket(&self) -> &str {
194        &self.config.bucket
195    }
196
197    /// Ensure the bucket exists, creating it if necessary.
198    pub async fn ensure_bucket_exists(&self) -> Result<()> {
199        match self
200            .client
201            .head_bucket()
202            .bucket(&self.config.bucket)
203            .send()
204            .await
205        {
206            Ok(_) => Ok(()),
207            Err(_) => {
208                // Bucket doesn't exist, create it
209                self.client
210                    .create_bucket()
211                    .bucket(&self.config.bucket)
212                    .send()
213                    .await
214                    .map_err(|e| {
215                        anyhow!("failed to create bucket '{}': {}", self.config.bucket, e)
216                    })?;
217                Ok(())
218            }
219        }
220    }
221
222    /// Put an object with a conditional check (If-None-Match: *).
223    ///
224    /// This performs an atomic write that only succeeds if the object does not
225    /// already exist. Returns:
226    /// - `Ok(true)` if the object was created successfully
227    /// - `Ok(false)` if the object already exists (PreconditionFailed)
228    /// - `Err(...)` for other errors
229    ///
230    /// # Arguments
231    /// * `key` - Object key
232    /// * `data` - Object data to write
233    pub async fn put_if_not_exists(&self, key: &str, data: Bytes) -> Result<bool> {
234        match self
235            .client
236            .put_object()
237            .bucket(&self.config.bucket)
238            .key(key)
239            .if_none_match("*")
240            .body(ByteStream::from(data))
241            .send()
242            .await
243        {
244            Ok(_) => Ok(true),
245            Err(e) => {
246                // Check if this is a precondition failed error (HTTP 412)
247                let service_error = e.into_service_error();
248                if service_error.code() == Some("PreconditionFailed") {
249                    Ok(false)
250                } else {
251                    Err(anyhow!(
252                        "S3 conditional put failed for key '{}': {}",
253                        key,
254                        service_error
255                    ))
256                }
257            }
258        }
259    }
260
261    /// Get an object's raw bytes.
262    ///
263    /// # Arguments
264    /// * `key` - Object key
265    ///
266    /// # Returns
267    /// - `Ok(Some(bytes))` if the object exists
268    /// - `Ok(None)` if the object does not exist
269    /// - `Err(...)` for other errors
270    pub async fn get_object(&self, key: &str) -> Result<Option<Bytes>> {
271        match self
272            .client
273            .get_object()
274            .bucket(&self.config.bucket)
275            .key(key)
276            .send()
277            .await
278        {
279            Ok(resp) => {
280                let data = resp
281                    .body
282                    .collect()
283                    .await
284                    .map_err(|e| anyhow!("failed to collect S3 response body: {}", e))?
285                    .into_bytes();
286                Ok(Some(data))
287            }
288            Err(e) => {
289                let service_error = e.into_service_error();
290                if service_error.code() == Some("NoSuchKey") {
291                    Ok(None)
292                } else {
293                    Err(anyhow!(
294                        "S3 get_object failed for key '{}': {}",
295                        key,
296                        service_error
297                    ))
298                }
299            }
300        }
301    }
302
303    /// Delete an object.
304    ///
305    /// # Arguments
306    /// * `key` - Object key
307    ///
308    /// # Returns
309    /// - `Ok(true)` if the object was deleted
310    /// - `Ok(false)` if the object did not exist
311    /// - `Err(...)` for other errors
312    pub async fn delete_object(&self, key: &str) -> Result<bool> {
313        match self
314            .client
315            .delete_object()
316            .bucket(&self.config.bucket)
317            .key(key)
318            .send()
319            .await
320        {
321            Ok(_) => Ok(true),
322            Err(e) => {
323                let service_error = e.into_service_error();
324                if service_error.code() == Some("NoSuchKey") {
325                    Ok(false)
326                } else {
327                    Err(anyhow!(
328                        "S3 delete_object failed for key '{}': {}",
329                        key,
330                        service_error
331                    ))
332                }
333            }
334        }
335    }
336
337    /// Check if an object exists (HEAD request).
338    ///
339    /// # Arguments
340    /// * `key` - Object key
341    ///
342    /// # Returns
343    /// - `Ok(true)` if the object exists
344    /// - `Ok(false)` if the object does not exist
345    /// - `Err(...)` for other errors
346    pub async fn has_object(&self, key: &str) -> Result<bool> {
347        match self
348            .client
349            .head_object()
350            .bucket(&self.config.bucket)
351            .key(key)
352            .send()
353            .await
354        {
355            Ok(_) => Ok(true),
356            Err(e) => {
357                let service_error = e.into_service_error();
358                // HeadObject returns "NotFound" when object doesn't exist
359                if service_error.code() == Some("NotFound") {
360                    Ok(false)
361                } else {
362                    Err(anyhow!(
363                        "S3 head_object failed for key '{}': {}",
364                        key,
365                        service_error
366                    ))
367                }
368            }
369        }
370    }
371
372    /// Put an object unconditionally (overwrite if exists).
373    ///
374    /// # Arguments
375    /// * `key` - Object key
376    /// * `data` - Object data to write
377    pub async fn put_object(&self, key: &str, data: Bytes) -> Result<()> {
378        self.client
379            .put_object()
380            .bucket(&self.config.bucket)
381            .key(key)
382            .body(ByteStream::from(data))
383            .send()
384            .await
385            .map_err(|e| anyhow!("S3 put_object failed for key '{}': {}", key, e))?;
386        Ok(())
387    }
388
389    /// Put blocks to object storage using a physical layout.
390    ///
391    /// This is the internal implementation that workers call after resolving
392    /// the logical layout handle to a physical layout.
393    ///
394    /// # Arguments
395    /// * `keys` - Sequence hashes identifying each block
396    /// * `layout` - Physical layout containing the block data
397    /// * `block_ids` - Block IDs within the layout to upload
398    ///
399    /// Returns a vector of results for each block:
400    /// - Ok(hash) indicates the block was successfully stored
401    /// - Err(hash) indicates the block failed to store
402    pub fn put_blocks_with_layout(
403        &self,
404        keys: Vec<SequenceHash>,
405        layout: PhysicalLayout,
406        block_ids: Vec<BlockId>,
407    ) -> BoxFuture<'static, Vec<Result<SequenceHash, SequenceHash>>> {
408        let config = layout.layout().config();
409        let block_size = config.block_size_bytes();
410        let region_size = config.region_size();
411        let is_contiguous = layout.layout().is_fully_contiguous();
412        let max_concurrent = self.config.max_concurrent_requests;
413        let client = self.client.clone();
414        let bucket = self.config.bucket.clone();
415        let formatter = self.key_formatter.clone();
416
417        Box::pin(async move {
418            let work_items: Vec<_> = keys.into_iter().zip(block_ids).collect();
419
420            let tasks = work_items.into_iter().map(|(key, block_id)| {
421                let client = client.clone();
422                let bucket = bucket.clone();
423                let key_str = formatter.format_key(&key);
424                let layout = layout.clone();
425
426                async move {
427                    let result: Result<(), anyhow::Error> = async {
428                        // Copy block data to bytes on rayon thread pool
429                        let data = tokio_rayon::spawn(move || {
430                            copy_block_to_bytes(
431                                &layout,
432                                block_id,
433                                block_size,
434                                region_size,
435                                is_contiguous,
436                            )
437                        })
438                        .await?;
439
440                        // Upload to S3
441                        client
442                            .put_object()
443                            .bucket(&bucket)
444                            .key(&key_str)
445                            .body(ByteStream::from(data))
446                            .send()
447                            .await
448                            .map_err(|e| anyhow!("S3 put_object failed: {}", e))?;
449
450                        Ok(())
451                    }
452                    .await;
453
454                    match result {
455                        Ok(()) => Ok(key),
456                        Err(e) => {
457                            tracing::warn!(key = %key, error = %e, "put block transfer failed");
458                            Err(key)
459                        }
460                    }
461                }
462            });
463
464            futures::stream::iter(tasks)
465                .buffer_unordered(max_concurrent)
466                .collect()
467                .await
468        })
469    }
470
471    /// Get blocks from object storage into a physical layout.
472    ///
473    /// This is the internal implementation that workers call after resolving
474    /// the logical layout handle to a physical layout.
475    ///
476    /// # Arguments
477    /// * `keys` - Sequence hashes identifying each block
478    /// * `layout` - Physical layout to write the block data into
479    /// * `block_ids` - Block IDs within the layout to download into
480    ///
481    /// Returns a vector of results for each block:
482    /// - Ok(hash) indicates the block was successfully retrieved
483    /// - Err(hash) indicates the block failed to retrieve
484    pub fn get_blocks_with_layout(
485        &self,
486        keys: Vec<SequenceHash>,
487        layout: PhysicalLayout,
488        block_ids: Vec<BlockId>,
489    ) -> BoxFuture<'static, Vec<Result<SequenceHash, SequenceHash>>> {
490        let config = layout.layout().config();
491        let block_size = config.block_size_bytes();
492        let region_size = config.region_size();
493        let is_contiguous = layout.layout().is_fully_contiguous();
494        let max_concurrent = self.config.max_concurrent_requests;
495        let client = self.client.clone();
496        let bucket = self.config.bucket.clone();
497        let formatter = self.key_formatter.clone();
498
499        Box::pin(async move {
500            let work_items: Vec<_> = keys.into_iter().zip(block_ids).collect();
501
502            let tasks = work_items.into_iter().map(|(key, block_id)| {
503                let client = client.clone();
504                let bucket = bucket.clone();
505                let key_str = formatter.format_key(&key);
506                let layout = layout.clone();
507
508                async move {
509                    let result: Result<(), anyhow::Error> = async {
510                        // Download from S3
511                        let resp = client
512                            .get_object()
513                            .bucket(&bucket)
514                            .key(&key_str)
515                            .send()
516                            .await
517                            .map_err(|e| anyhow!("S3 get_object failed: {}", e))?;
518
519                        let data = resp
520                            .body
521                            .collect()
522                            .await
523                            .map_err(|e| anyhow!("failed to collect S3 response body: {}", e))?
524                            .into_bytes();
525
526                        // Copy bytes to block on rayon thread pool
527                        tokio_rayon::spawn(move || {
528                            copy_bytes_to_block(
529                                &data,
530                                &layout,
531                                block_id,
532                                block_size,
533                                region_size,
534                                is_contiguous,
535                            )
536                        })
537                        .await?;
538
539                        Ok(())
540                    }
541                    .await;
542
543                    match result {
544                        Ok(()) => Ok(key),
545                        Err(e) => {
546                            tracing::warn!(key = %key, error = %e, "get block transfer failed");
547                            Err(key)
548                        }
549                    }
550                }
551            });
552
553            futures::stream::iter(tasks)
554                .buffer_unordered(max_concurrent)
555                .collect()
556                .await
557        })
558    }
559
560    /// Get an object's raw bytes along with its ETag.
561    ///
562    /// Used for conditional updates (CAS-style operations) where the caller
563    /// needs the current ETag to perform a conditional PUT.
564    ///
565    /// # Returns
566    /// - `Ok(Some((bytes, etag)))` if the object exists
567    /// - `Ok(None)` if the object does not exist
568    /// - `Err(...)` for other errors
569    pub async fn get_object_with_etag(&self, key: &str) -> Result<Option<(Bytes, Option<String>)>> {
570        match self
571            .client
572            .get_object()
573            .bucket(&self.config.bucket)
574            .key(key)
575            .send()
576            .await
577        {
578            Ok(resp) => {
579                let etag = resp.e_tag().map(|s| s.to_string());
580                let data = resp
581                    .body
582                    .collect()
583                    .await
584                    .map_err(|e| anyhow!("failed to collect S3 response body: {}", e))?
585                    .into_bytes();
586                Ok(Some((data, etag)))
587            }
588            Err(e) => {
589                let service_error = e.into_service_error();
590                if service_error.code() == Some("NoSuchKey") {
591                    Ok(None)
592                } else {
593                    Err(anyhow!(
594                        "S3 get_object failed for key '{}': {}",
595                        key,
596                        service_error
597                    ))
598                }
599            }
600        }
601    }
602
603    /// Put an object with an ETag precondition (If-Match).
604    ///
605    /// This performs a conditional write that only succeeds if the object's current
606    /// ETag matches the provided value. Used for CAS-style atomic updates.
607    ///
608    /// # Returns
609    /// - `Ok(true)` if the write succeeded (ETag matched)
610    /// - `Ok(false)` if the ETag did not match (412 PreconditionFailed — lost the race)
611    /// - `Err(...)` for other errors
612    pub async fn put_object_if_match(&self, key: &str, data: Bytes, etag: &str) -> Result<bool> {
613        match self
614            .client
615            .put_object()
616            .bucket(&self.config.bucket)
617            .key(key)
618            .if_match(etag)
619            .body(ByteStream::from(data))
620            .send()
621            .await
622        {
623            Ok(_) => Ok(true),
624            Err(e) => {
625                let service_error = e.into_service_error();
626                if service_error.code() == Some("PreconditionFailed") {
627                    Ok(false)
628                } else {
629                    Err(anyhow!(
630                        "S3 conditional put (if-match) failed for key '{}': {}",
631                        key,
632                        service_error
633                    ))
634                }
635            }
636        }
637    }
638}
639
640/// Build an S3 client from configuration.
641async fn build_s3_client(config: &S3Config) -> Result<Client> {
642    let sdk_config = aws_config::defaults(aws_config::BehaviorVersion::latest())
643        .region(aws_sdk_s3::config::Region::new(config.region.clone()))
644        .load()
645        .await;
646
647    let mut s3_config_builder = aws_sdk_s3::config::Builder::from(&sdk_config);
648
649    if let Some(endpoint) = &config.endpoint_url {
650        s3_config_builder = s3_config_builder.endpoint_url(endpoint);
651    }
652
653    if config.force_path_style {
654        s3_config_builder = s3_config_builder.force_path_style(true);
655    }
656
657    let s3_config = s3_config_builder.build();
658    Ok(Client::from_conf(s3_config))
659}
660
661/// Copy block data from a layout to a Bytes buffer.
662///
663/// For fully contiguous layouts, this is a single memcpy.
664/// For layer-separate layouts, this iterates over all regions.
665fn copy_block_to_bytes(
666    layout: &PhysicalLayout,
667    block_id: BlockId,
668    block_size: usize,
669    region_size: usize,
670    is_contiguous: bool,
671) -> Result<Bytes> {
672    if is_contiguous {
673        // Fast path: single contiguous region — the layout guarantees that
674        // block_size bytes are contiguous from region(block_id, 0, 0).addr().
675        let region = layout.memory_region(block_id, 0, 0)?;
676        let slice = unsafe { std::slice::from_raw_parts(region.addr() as *const u8, block_size) };
677        Ok(Bytes::copy_from_slice(slice))
678    } else {
679        // Slow path: iterate over all regions
680        let mut buf = Vec::with_capacity(block_size);
681        let inner_layout = layout.layout();
682        for layer_id in 0..inner_layout.num_layers() {
683            for outer_id in 0..inner_layout.outer_dim() {
684                let region = layout.memory_region(block_id, layer_id, outer_id)?;
685                if region.size() < region_size {
686                    return Err(anyhow!(
687                        "memory region too small: got {} bytes, need {}",
688                        region.size(),
689                        region_size
690                    ));
691                }
692                let slice =
693                    unsafe { std::slice::from_raw_parts(region.addr() as *const u8, region_size) };
694                buf.extend_from_slice(slice);
695            }
696        }
697        Ok(Bytes::from(buf))
698    }
699}
700
701/// Copy data from a Bytes buffer to a layout.
702///
703/// For fully contiguous layouts, this is a single memcpy.
704/// For layer-separate layouts, this iterates over all regions.
705fn copy_bytes_to_block(
706    data: &[u8],
707    layout: &PhysicalLayout,
708    block_id: BlockId,
709    block_size: usize,
710    region_size: usize,
711    is_contiguous: bool,
712) -> Result<()> {
713    if is_contiguous {
714        // Fast path: single contiguous region — the layout guarantees that
715        // block_size bytes are contiguous from region(block_id, 0, 0).addr().
716        if data.len() < block_size {
717            return Err(anyhow!(
718                "S3 data too short: got {} bytes, expected {}",
719                data.len(),
720                block_size
721            ));
722        }
723        let region = layout.memory_region(block_id, 0, 0)?;
724        unsafe {
725            std::ptr::copy_nonoverlapping(data.as_ptr(), region.addr() as *mut u8, block_size);
726        }
727    } else {
728        // Slow path: iterate over all regions
729        let mut offset = 0;
730        let inner_layout = layout.layout();
731        for layer_id in 0..inner_layout.num_layers() {
732            for outer_id in 0..inner_layout.outer_dim() {
733                if offset + region_size > data.len() {
734                    return Err(anyhow!(
735                        "S3 data too short at offset {}: need {} more bytes, only {} remain",
736                        offset,
737                        region_size,
738                        data.len().saturating_sub(offset)
739                    ));
740                }
741                let region = layout.memory_region(block_id, layer_id, outer_id)?;
742                if region.size() < region_size {
743                    return Err(anyhow!(
744                        "memory region too small: got {} bytes, need {}",
745                        region.size(),
746                        region_size
747                    ));
748                }
749                unsafe {
750                    std::ptr::copy_nonoverlapping(
751                        data[offset..].as_ptr(),
752                        region.addr() as *mut u8,
753                        region_size,
754                    );
755                }
756                offset += region_size;
757            }
758        }
759    }
760    Ok(())
761}
762
763impl ObjectBlockOps for S3ObjectBlockClient {
764    fn has_blocks(
765        &self,
766        keys: Vec<SequenceHash>,
767    ) -> BoxFuture<'static, Vec<(SequenceHash, Option<usize>)>> {
768        let max_concurrent = self.config.max_concurrent_requests;
769        let client = self.client.clone();
770        let bucket = self.config.bucket.clone();
771        let formatter = self.key_formatter.clone();
772
773        Box::pin(async move {
774            let tasks = keys.into_iter().map(|key| {
775                let client = client.clone();
776                let bucket = bucket.clone();
777                let key_str = formatter.format_key(&key);
778
779                async move {
780                    match client
781                        .head_object()
782                        .bucket(&bucket)
783                        .key(&key_str)
784                        .send()
785                        .await
786                    {
787                        Ok(resp) => (key, resp.content_length().map(|l| l as usize)),
788                        Err(e) => {
789                            tracing::warn!(key = %key, error = %e, "head_object failed, treating as missing");
790                            (key, None)
791                        }
792                    }
793                }
794            });
795
796            futures::stream::iter(tasks)
797                .buffer_unordered(max_concurrent)
798                .collect()
799                .await
800        })
801    }
802
803    fn put_blocks(
804        &self,
805        keys: Vec<SequenceHash>,
806        _src_layout: LogicalLayoutHandle,
807        _block_ids: Vec<BlockId>,
808    ) -> BoxFuture<'static, Vec<Result<SequenceHash, SequenceHash>>> {
809        // S3ObjectBlockClient cannot resolve LogicalLayoutHandle to PhysicalLayout.
810        // Workers should use put_blocks_with_layout() instead after resolving the handle.
811        tracing::error!(
812            "S3ObjectBlockClient::put_blocks called with LogicalLayoutHandle - \
813             use put_blocks_with_layout() via DirectWorker instead"
814        );
815        Box::pin(async move { keys.into_iter().map(Err).collect() })
816    }
817
818    fn get_blocks(
819        &self,
820        keys: Vec<SequenceHash>,
821        _dst_layout: LogicalLayoutHandle,
822        _block_ids: Vec<BlockId>,
823    ) -> BoxFuture<'static, Vec<Result<SequenceHash, SequenceHash>>> {
824        // S3ObjectBlockClient cannot resolve LogicalLayoutHandle to PhysicalLayout.
825        // Workers should use get_blocks_with_layout() instead after resolving the handle.
826        tracing::error!(
827            "S3ObjectBlockClient::get_blocks called with LogicalLayoutHandle - \
828             use get_blocks_with_layout() via DirectWorker instead"
829        );
830        Box::pin(async move { keys.into_iter().map(Err).collect() })
831    }
832
833    fn put_blocks_with_layout(
834        &self,
835        keys: Vec<SequenceHash>,
836        layout: PhysicalLayout,
837        block_ids: Vec<BlockId>,
838    ) -> BoxFuture<'static, Vec<Result<SequenceHash, SequenceHash>>> {
839        // Delegate to the inherent method
840        S3ObjectBlockClient::put_blocks_with_layout(self, keys, layout, block_ids)
841    }
842
843    fn get_blocks_with_layout(
844        &self,
845        keys: Vec<SequenceHash>,
846        layout: PhysicalLayout,
847        block_ids: Vec<BlockId>,
848    ) -> BoxFuture<'static, Vec<Result<SequenceHash, SequenceHash>>> {
849        // Delegate to the inherent method
850        S3ObjectBlockClient::get_blocks_with_layout(self, keys, layout, block_ids)
851    }
852}
853
854#[cfg(test)]
855mod tests {
856    use super::*;
857
858    #[test]
859    fn test_s3_config_default() {
860        let config = S3Config::default();
861        assert_eq!(config.endpoint_url, Some("http://localhost:9000".into()));
862        assert_eq!(config.bucket, "kvbm-blocks");
863        assert_eq!(config.region, "us-east-1");
864        assert!(config.force_path_style);
865        assert_eq!(config.max_concurrent_requests, 16);
866    }
867
868    #[test]
869    fn test_s3_config_aws() {
870        let config = S3Config::aws("my-bucket".into(), "us-west-2".into());
871        assert_eq!(config.endpoint_url, None);
872        assert_eq!(config.bucket, "my-bucket");
873        assert_eq!(config.region, "us-west-2");
874        assert!(!config.force_path_style);
875    }
876
877    #[test]
878    fn test_s3_config_minio() {
879        let config = S3Config::minio("http://minio:9000".into(), "test-bucket".into());
880        assert_eq!(config.endpoint_url, Some("http://minio:9000".into()));
881        assert_eq!(config.bucket, "test-bucket");
882        assert!(config.force_path_style);
883    }
884}
885
886#[cfg(all(test, feature = "testing"))]
887mod bounds_check_tests {
888    use super::*;
889    use crate::object::LayoutConfigExt;
890    use kvbm_physical::testing::{create_fc_layout, create_lw_layout, create_test_agent};
891    use kvbm_physical::transfer::StorageKind;
892
893    #[test]
894    fn test_copy_bytes_to_block_rejects_short_data_contiguous() {
895        let agent = create_test_agent("test_short_data_fc");
896        let layout = create_fc_layout(agent, StorageKind::System, 2);
897        let config = layout.layout().config();
898        let block_size = config.block_size_bytes();
899        let region_size = config.region_size();
900
901        // Data is one byte short
902        let short_data = vec![0u8; block_size - 1];
903        let err = copy_bytes_to_block(&short_data, &layout, 0, block_size, region_size, true)
904            .expect_err("should reject short data");
905        assert!(
906            err.to_string().contains("S3 data too short"),
907            "unexpected error: {}",
908            err
909        );
910    }
911
912    #[test]
913    fn test_copy_bytes_to_block_rejects_short_data_non_contiguous() {
914        let agent = create_test_agent("test_short_data_lw");
915        let layout = create_lw_layout(agent, StorageKind::System, 2);
916        let config = layout.layout().config();
917        let block_size = config.block_size_bytes();
918        let region_size = config.region_size();
919
920        // Data is one region short
921        let short_data = vec![0u8; block_size - region_size];
922        let err = copy_bytes_to_block(&short_data, &layout, 0, block_size, region_size, false)
923            .expect_err("should reject short data in non-contiguous path");
924        assert!(
925            err.to_string().contains("S3 data too short"),
926            "unexpected error: {}",
927            err
928        );
929    }
930
931    #[test]
932    fn test_copy_bytes_to_block_accepts_exact_size() {
933        let agent = create_test_agent("test_exact_fc");
934        let layout = create_fc_layout(agent, StorageKind::System, 2);
935        let config = layout.layout().config();
936        let block_size = config.block_size_bytes();
937        let region_size = config.region_size();
938
939        let data = vec![42u8; block_size];
940        copy_bytes_to_block(&data, &layout, 0, block_size, region_size, true)
941            .expect("exact-size data should succeed");
942    }
943
944    #[test]
945    fn test_copy_block_to_bytes_roundtrip_contiguous() {
946        let agent = create_test_agent("test_roundtrip_fc");
947        let layout = create_fc_layout(agent, StorageKind::System, 2);
948        let config = layout.layout().config();
949        let block_size = config.block_size_bytes();
950        let region_size = config.region_size();
951
952        // Write known data
953        let data = vec![0xAB_u8; block_size];
954        copy_bytes_to_block(&data, &layout, 0, block_size, region_size, true).unwrap();
955
956        // Read it back
957        let out = copy_block_to_bytes(&layout, 0, block_size, region_size, true).unwrap();
958        assert_eq!(out.as_ref(), &data[..]);
959    }
960
961    #[test]
962    fn test_copy_block_to_bytes_roundtrip_non_contiguous() {
963        let agent = create_test_agent("test_roundtrip_lw");
964        let layout = create_lw_layout(agent, StorageKind::System, 2);
965        let config = layout.layout().config();
966        let block_size = config.block_size_bytes();
967        let region_size = config.region_size();
968
969        let data = vec![0xCD_u8; block_size];
970        copy_bytes_to_block(&data, &layout, 0, block_size, region_size, false).unwrap();
971
972        let out = copy_block_to_bytes(&layout, 0, block_size, region_size, false).unwrap();
973        assert_eq!(out.as_ref(), &data[..]);
974    }
975}
976
977#[cfg(all(test, feature = "testing-s3"))]
978pub mod s3_integration {
979    use super::*;
980
981    /// Create an S3ObjectBlockClient connected to the test MinIO instance.
982    ///
983    /// Reads `S3_TEST_ENDPOINT` from the environment (set by `test-s3.sh`).
984    /// Falls back to `http://localhost:9876`.
985    pub async fn create_test_client(bucket: &str) -> S3ObjectBlockClient {
986        let endpoint =
987            std::env::var("S3_TEST_ENDPOINT").unwrap_or_else(|_| "http://localhost:9876".into());
988        let config = S3Config::minio(endpoint, bucket.to_string());
989        let client = S3ObjectBlockClient::new(config).await.unwrap();
990        client.ensure_bucket_exists().await.unwrap();
991        client
992    }
993
994    #[tokio::test]
995    async fn test_put_get_roundtrip() {
996        let client = create_test_client("test-roundtrip").await;
997        let key = format!("roundtrip-{}", uuid::Uuid::new_v4());
998        let payload = Bytes::from("hello world");
999
1000        client.put_object(&key, payload.clone()).await.unwrap();
1001
1002        let result = client.get_object(&key).await.unwrap();
1003        assert_eq!(result, Some(payload));
1004
1005        // Cleanup
1006        client.delete_object(&key).await.unwrap();
1007    }
1008
1009    #[tokio::test]
1010    async fn test_put_object_if_match_rejects_stale_etag() {
1011        let client = create_test_client("test-if-match").await;
1012        let key = format!("if-match-{}", uuid::Uuid::new_v4());
1013
1014        // Write initial object
1015        client
1016            .put_object(&key, Bytes::from("version1"))
1017            .await
1018            .unwrap();
1019
1020        // Get with ETag
1021        let (_, etag) = client
1022            .get_object_with_etag(&key)
1023            .await
1024            .unwrap()
1025            .expect("object should exist");
1026        let etag = etag.expect("should have etag");
1027
1028        // Overwrite the object to change its ETag
1029        client
1030            .put_object(&key, Bytes::from("version2"))
1031            .await
1032            .unwrap();
1033
1034        // Conditional put with stale ETag should fail
1035        let won = client
1036            .put_object_if_match(&key, Bytes::from("version3"), &etag)
1037            .await
1038            .unwrap();
1039        assert!(!won, "conditional put with stale ETag should return false");
1040
1041        // Verify the object still has version2
1042        let data = client.get_object(&key).await.unwrap().unwrap();
1043        assert_eq!(data, Bytes::from("version2"));
1044
1045        // Cleanup
1046        client.delete_object(&key).await.unwrap();
1047    }
1048
1049    #[tokio::test]
1050    async fn test_put_object_if_match_accepts_current_etag() {
1051        let client = create_test_client("test-if-match-ok").await;
1052        let key = format!("if-match-ok-{}", uuid::Uuid::new_v4());
1053
1054        client
1055            .put_object(&key, Bytes::from("version1"))
1056            .await
1057            .unwrap();
1058
1059        let (_, etag) = client
1060            .get_object_with_etag(&key)
1061            .await
1062            .unwrap()
1063            .expect("object should exist");
1064        let etag = etag.expect("should have etag");
1065
1066        // Conditional put with current ETag should succeed
1067        let won = client
1068            .put_object_if_match(&key, Bytes::from("version2"), &etag)
1069            .await
1070            .unwrap();
1071        assert!(won, "conditional put with current ETag should succeed");
1072
1073        let data = client.get_object(&key).await.unwrap().unwrap();
1074        assert_eq!(data, Bytes::from("version2"));
1075
1076        // Cleanup
1077        client.delete_object(&key).await.unwrap();
1078    }
1079}
kvbm_engine/object/s3/client.rs

kvbm_engine/object/s3/
client.rs
