agentos_vfs/s3/
block_store.rs

1use async_trait::async_trait;
2use aws_sdk_s3::primitives::{ByteStream, ByteStreamError};
3use aws_sdk_s3::types::{Delete, ObjectIdentifier};
4use aws_sdk_s3::Client;
5use vfs::engine::{BlockKey, BlockStore, VfsError, VfsResult};
6
7#[derive(Debug, Clone, Default)]
8pub struct S3BlockStoreOptions {
9    pub prefix: String,
10}
11
12#[derive(Debug, Clone)]
13pub struct S3BlockStore {
14    client: Client,
15    bucket: String,
16    options: S3BlockStoreOptions,
17}
18
19impl S3BlockStore {
20    pub fn new(client: Client, bucket: impl Into<String>) -> Self {
21        Self::with_options(client, bucket, S3BlockStoreOptions::default())
22    }
23
24    pub fn with_options(
25        client: Client,
26        bucket: impl Into<String>,
27        options: S3BlockStoreOptions,
28    ) -> Self {
29        Self {
30            client,
31            bucket: bucket.into(),
32            options,
33        }
34    }
35
36    fn key_for(&self, key: &BlockKey) -> String {
37        format!("{}{}", self.options.prefix, key.0)
38    }
39}
40
41#[async_trait]
42impl BlockStore for S3BlockStore {
43    async fn get(&self, key: &BlockKey) -> VfsResult<Vec<u8>> {
44        let object_key = self.key_for(key);
45        let response = self
46            .client
47            .get_object()
48            .bucket(&self.bucket)
49            .key(&object_key)
50            .send()
51            .await
52            .map_err(|err| s3_error(format!("get s3 block '{object_key}': {err}")))?;
53        collect_body(response.body, &object_key).await
54    }
55
56    async fn get_range(&self, key: &BlockKey, off: u64, len: u64) -> VfsResult<Vec<u8>> {
57        if len == 0 {
58            return Ok(Vec::new());
59        }
60        let object_key = self.key_for(key);
61        let end = off
62            .checked_add(len)
63            .and_then(|value| value.checked_sub(1))
64            .ok_or_else(|| VfsError::einval("invalid S3 byte range"))?;
65        let response = self
66            .client
67            .get_object()
68            .bucket(&self.bucket)
69            .key(&object_key)
70            .range(format!("bytes={off}-{end}"))
71            .send()
72            .await
73            .map_err(|err| s3_error(format!("get s3 block range '{object_key}': {err}")))?;
74        collect_body(response.body, &object_key).await
75    }
76
77    async fn put(&self, key: &BlockKey, data: &[u8]) -> VfsResult<()> {
78        let object_key = self.key_for(key);
79        self.client
80            .put_object()
81            .bucket(&self.bucket)
82            .key(&object_key)
83            .body(ByteStream::from(data.to_vec()))
84            .send()
85            .await
86            .map_err(|err| s3_error(format!("put s3 block '{object_key}': {err}")))?;
87        Ok(())
88    }
89
90    async fn exists(&self, key: &BlockKey) -> VfsResult<bool> {
91        let object_key = self.key_for(key);
92        match self
93            .client
94            .head_object()
95            .bucket(&self.bucket)
96            .key(&object_key)
97            .send()
98            .await
99        {
100            Ok(_) => Ok(true),
101            Err(err)
102                if err.as_service_error().is_some_and(|err| {
103                    err.is_not_found()
104                        || matches!(err.meta().code(), Some("NotFound" | "NoSuchKey"))
105                }) =>
106            {
107                Ok(false)
108            }
109            Err(err) => Err(s3_error(format!("head s3 block '{object_key}': {err}"))),
110        }
111    }
112
113    async fn delete_many(&self, keys: &[BlockKey]) -> VfsResult<()> {
114        let mut errors = Vec::new();
115        for chunk in keys.chunks(1000) {
116            let mut object_keys = Vec::with_capacity(chunk.len());
117            let mut objects = Vec::with_capacity(chunk.len());
118            for key in chunk {
119                let object_key = self.key_for(key);
120                let object = ObjectIdentifier::builder()
121                    .key(&object_key)
122                    .build()
123                    .map_err(|err| {
124                        s3_error(format!("build s3 delete key '{object_key}': {err}"))
125                    })?;
126                object_keys.push(object_key);
127                objects.push(object);
128            }
129            let delete = Delete::builder()
130                .set_objects(Some(objects))
131                .quiet(true)
132                .build()
133                .map_err(|err| s3_error(format!("build s3 delete batch: {err}")))?;
134            match self
135                .client
136                .delete_objects()
137                .bucket(&self.bucket)
138                .delete(delete)
139                .send()
140                .await
141            {
142                Ok(response) => {
143                    for error in response.errors() {
144                        let key = error.key().unwrap_or("<unknown>");
145                        let code = error.code().unwrap_or("unknown");
146                        let message = error.message().unwrap_or("unknown error");
147                        errors.push(format!("delete s3 block '{key}': {code}: {message}"));
148                    }
149                }
150                Err(err) => errors.push(format!(
151                    "delete s3 block batch [{}]: {err}",
152                    object_keys.join(", ")
153                )),
154            }
155        }
156        if errors.is_empty() {
157            Ok(())
158        } else {
159            Err(s3_error(format!(
160                "delete {} s3 blocks failed: {}",
161                errors.len(),
162                errors.join("; ")
163            )))
164        }
165    }
166
167    async fn copy(&self, src: &BlockKey, dst: &BlockKey) -> VfsResult<()> {
168        let src_key = self.key_for(src);
169        let dst_key = self.key_for(dst);
170        self.client
171            .copy_object()
172            .bucket(&self.bucket)
173            .copy_source(format!("{}/{}", self.bucket, src_key))
174            .key(&dst_key)
175            .send()
176            .await
177            .map_err(|err| s3_error(format!("copy s3 block '{src_key}' to '{dst_key}': {err}")))?;
178        Ok(())
179    }
180}
181
182pub(crate) async fn collect_body(body: ByteStream, key: &str) -> VfsResult<Vec<u8>> {
183    body.collect()
184        .await
185        .map_err(|err| body_error(key, err))
186        .map(|bytes| bytes.into_bytes().to_vec())
187}
188
189pub(crate) fn s3_error(message: impl Into<String>) -> VfsError {
190    VfsError::eio(message)
191}
192
193fn body_error(key: &str, err: ByteStreamError) -> VfsError {
194    VfsError::eio(format!("read s3 object '{key}': {err}"))
195}
agentos_vfs/s3/block_store.rs

agentos_vfs/s3/
block_store.rs
