feat(lib/storage): update S3 error and crate config

Modified lib/storage/src/error.rs +9 −0 View lib/storage/src/error.rs

@@ -7,6 +7,7 @@

77pub enum StorageError {

88 InvalidConfig(String),

99 S3(S3Error),

10 Io(std::io::Error),

1011 UnexpectedStatus {

1112 operation: &'static str,

1213 status: u16,

@@ -18,6 +19,7 @@

1819 match self {

1920 StorageError::InvalidConfig(msg) => write!(f, "invalid s3 config: {msg}"),

2021 StorageError::S3(err) => write!(f, "s3 error: {err}"),

22 StorageError::Io(err) => write!(f, "io error: {err}"),

2123 StorageError::UnexpectedStatus { operation, status } => {

2224 write!(f, "{operation} returned unexpected status {status}")

2325 }

@@ -29,6 +31,7 @@

2931 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {

3032 match self {

3133 StorageError::S3(err) => Some(err),

34 StorageError::Io(err) => Some(err),

3235 _ => None,

3336 }

3437 }

@@ -40,6 +43,12 @@

4043 }

4144}

4245

46impl From<std::io::Error> for StorageError {

47 fn from(err: std::io::Error) -> Self {

48 StorageError::Io(err)

49 }

50}

51

4352impl From<CredentialsError> for StorageError {

4453 fn from(err: CredentialsError) -> Self {

4554 StorageError::S3(S3Error::Credentials(err))

Modified lib/storage/src/s3.rs +222 −0 View lib/storage/src/s3.rs

@@ -1,11 +1,17 @@

11use std::collections::HashMap;

2use std::io;

3use std::pin::Pin;

4use std::task::{Context as TaskContext, Poll};

25

36use config::s3::S3Config;

47use http::HeaderMap;

8use s3::bucket::{CHUNK_SIZE, Part};

59use s3::creds::Credentials;

610use s3::region::Region;

711use s3::serde_types::{DateTime as S3DateTime, Object as S3Object};

812use s3::Bucket;

13use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};

14use futures::StreamExt as _;

915

1016use crate::error::StorageError;

1117

@@ -93,6 +99,97 @@

9399 }

94100 }

95101

102 /// True streaming upload with bounded memory: consumes the reader in

103 /// 8 MiB chunks and uploads them serially (single PUT below one chunk,

104 /// multipart above), hashing nothing but counting bytes through a tee.

105 ///

106 /// Returns the number of bytes actually read from the reader.

107 pub async fn put_stream<R: AsyncRead + Unpin>(

108 &self,

109 key: &str,

110 reader: &mut R,

111 content_type: &str,

112 ) -> Result<u64, StorageError> {

113 let first = read_chunk(reader).await?;

114 if first.len() < CHUNK_SIZE {

115 self.put_with_content_type(key, &first, content_type).await?;

116 return Ok(first.len() as u64);

117 }

118 let msg = self.bucket.initiate_multipart_upload(key, content_type).await?;

119 let path = msg.key;

120 let upload_id = msg.upload_id;

121

122 let mut chunk = first;

123 let mut part_number: u32 = 0;

124 let mut parts: Vec<Part> = Vec::new();

125 let mut total: u64 = 0;

126 loop {

127 total += chunk.len() as u64;

128 part_number += 1;

129 match self

130 .bucket

131 .put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)

132 .await

133 {

134 Ok(part) => parts.push(part),

135 Err(e) => {

136 abort_upload(&self.bucket, &path, &upload_id).await;

137 return Err(StorageError::from(e));

138 }

139 }

140

141 chunk = match read_chunk(reader).await {

142 Ok(c) => c,

143 Err(e) => {

144 abort_upload(&self.bucket, &path, &upload_id).await;

145 return Err(e);

146 }

147 };

148 if chunk.is_empty() {

149 abort_upload(&self.bucket, &path, &upload_id).await;

150 return Err(StorageError::UnexpectedStatus {

151 operation: "put_stream",

152 status: 400,

153 });

154 }

155 if chunk.len() < CHUNK_SIZE {

156 total += chunk.len() as u64;

157 part_number += 1;

158 match self

159 .bucket

160 .put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)

161 .await

162 {

163 Ok(part) => parts.push(part),

164 Err(e) => {

165 abort_upload(&self.bucket, &path, &upload_id).await;

166 return Err(StorageError::from(e));

167 }

168 }

169 break;

170 }

171 }

172

173 match self

174 .bucket

175 .complete_multipart_upload(&path, &upload_id, parts)

176 .await

177 {

178 Ok(r) if (200..300).contains(&r.status_code()) => Ok(total),

179 Ok(r) => {

180 abort_upload(&self.bucket, &path, &upload_id).await;

181 Err(StorageError::UnexpectedStatus {

182 operation: "complete_multipart_upload",

183 status: r.status_code(),

184 })

185 }

186 Err(e) => {

187 abort_upload(&self.bucket, &path, &upload_id).await;

188 Err(StorageError::from(e))

189 }

190 }

191 }

192

96193 pub async fn get(&self, key: &str) -> Result<Vec<u8>, StorageError> {

97194 let response = self.bucket.get_object(key).await?;

98195 let status = response.status_code();

@@ -105,6 +202,47 @@

105202 Ok(response.as_slice().to_vec())

106203 }

107204

205 /// Stream an object into `writer` without buffering it in memory; the

206 /// writer is expected to be Send + Unpin so HTTP body writers can be fed

207 /// directly. Returns the number of bytes written.

208 pub async fn get_to_writer<W: AsyncWrite + Send + Unpin>(

209 &self,

210 key: &str,

211 writer: &mut W,

212 ) -> Result<u64, StorageError> {

213 let mut counting = CountingWriter { inner: writer, count: 0 };

214 let status = self

215 .bucket

216 .get_object_to_writer(key, &mut counting)

217 .await?;

218 if !(200..300).contains(&status) {

219 return Err(StorageError::UnexpectedStatus {

220 operation: "get_object",

221 status,

222 });

223 }

224 Ok(counting.count)

225 }

226

227 /// Stream an object out as a byte stream (ready for an HTTP response

228 /// body). The returned stream yields `Result<Bytes, StorageError>`.

229 pub async fn get_stream(

230 &self,

231 key: &str,

232 ) -> Result<

233 Pin<Box<dyn futures::Stream<Item = Result<bytes::Bytes, StorageError>> + Send>>,

234 StorageError,

235 > {

236 let stream = self.bucket.get_object_stream(key).await?;

237 if !(200..300).contains(&stream.status_code) {

238 return Err(StorageError::UnexpectedStatus {

239 operation: "get_object",

240 status: stream.status_code,

241 });

242 }

243 Ok(Box::pin(stream.bytes.map(move |item| item.map_err(StorageError::from))))

244 }

245

108246 pub async fn delete(&self, key: &str) -> Result<(), StorageError> {

109247 let response = self.bucket.delete_object(key).await?;

110248 let status = response.status_code();

@@ -286,3 +424,87 @@

286424 }

287425 Ok(())

288426}

427

428/// Read up to one `CHUNK_SIZE` block from the reader (EOF-tolerant).

429async fn read_chunk<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec<u8>, StorageError> {

430 let mut chunk = Vec::with_capacity(CHUNK_SIZE);

431 let mut take = reader.take(CHUNK_SIZE as u64);

432 let read = AsyncReadExt::read_to_end(&mut take, &mut chunk).await?;

433 chunk.truncate(read);

434 Ok(chunk)

435}

436

437/// Fire-and-forget multipart abort (errors are intentionally swallowed: the

438/// original failure is what the caller cares about).

439async fn abort_upload(bucket: &Bucket, path: &str, upload_id: &str) {

440 let _ = bucket.abort_upload(path, upload_id).await;

441}

442

443struct CountingWriter<'a, W> {

444 inner: &'a mut W,

445 count: u64,

446}

447

448impl<W: AsyncWrite + Unpin> AsyncWrite for CountingWriter<'_, W> {

449 fn poll_write(

450 mut self: Pin<&mut Self>,

451 cx: &mut TaskContext<'_>,

452 buf: &[u8],

453 ) -> Poll<io::Result<usize>> {

454 let this = &mut *self;

455 match Pin::new(&mut *this.inner).poll_write(cx, buf) {

456 Poll::Ready(Ok(n)) => {

457 this.count += n as u64;

458 Poll::Ready(Ok(n))

459 }

460 Poll::Ready(Err(e)) => Poll::Ready(Err(e)),

461 Poll::Pending => Poll::Pending,

462 }

463 }

464

465 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {

466 Pin::new(&mut *self.inner).poll_flush(cx)

467 }

468

469 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {

470 Pin::new(&mut *self.inner).poll_shutdown(cx)

471 }

472

473 fn is_write_vectored(&self) -> bool {

474 self.inner.is_write_vectored()

475 }

476}

477

478#[cfg(test)]

479mod tests {

480 use super::*;

481 use tokio::io::AsyncWriteExt;

482

483 #[tokio::test]

484 async fn read_chunk_respects_chunk_size() {

485 let mut data = (0..CHUNK_SIZE as u32).map(|i| (i % 251) as u8).collect::<Vec<_>>();

486 let mut cursor = io::Cursor::new(&data);

487 let first = read_chunk(&mut cursor).await.expect("read");

488 assert_eq!(first.len(), CHUNK_SIZE);

489 let second = read_chunk(&mut cursor).await.expect("read");

490 assert!(second.is_empty());

491

492 data.truncate(CHUNK_SIZE - 1);

493 let mut cursor = io::Cursor::new(&data);

494 let first = read_chunk(&mut cursor).await.expect("read");

495 assert_eq!(first.len(), CHUNK_SIZE - 1);

496 }

497

498 #[tokio::test]

499 async fn counting_writer_tallies_bytes() {

500 let mut buf = Vec::new();

501 {

502 let mut cw = CountingWriter { inner: &mut buf, count: 0 };

503 cw.write_all(b"hello ").await.expect("write");

504 cw.write_all(b"world").await.expect("write");

505 cw.flush().await.expect("flush");

506 assert_eq!(cw.count, 11);

507 }

508 assert_eq!(buf, b"hello world");

509 }

510}

Modified lib/storage/Cargo.toml +3 −1 View lib/storage/Cargo.toml

@@ -6,5 +6,7 @@

66[dependencies]

77config = { path = "../config" }

88rust-s3-async = { workspace = true }

9tokio = { workspace = true }

9tokio = { workspace = true, features = ["macros", "rt", "io-util"] }

1010http = "0.2"

11futures = "0.3"

12bytes = "1"