Skip to main content

max / makenotwork

14.2 KB · 423 lines History Blame Raw
1 //! Getting bytes to S3.
2 //!
3 //! The one family here with real logic rather than transport boilerplate: a
4 //! file over [`MULTIPART_THRESHOLD_BYTES`] goes up in parts, which means a part
5 //! geometry from the server, a bounded set of presigned part URLs refreshed on
6 //! a [`PART_URL_WINDOW`] cadence, and a completion call that names every ETag.
7
8 use super::{MnwApiClient, json_response};
9 use serde::Deserialize;
10
11 /// Response from the presign-upload internal endpoint.
12 #[derive(Debug, Deserialize)]
13 #[allow(dead_code)]
14 pub(crate) struct PresignResponse {
15 pub upload_url: String,
16 pub s3_key: String,
17 pub expires_in: u64,
18 pub cache_control: Option<String>,
19 }
20
21 /// Above this size an upload goes through a multipart session instead of one
22 /// presigned PUT. The single-PUT path reads the whole file into memory, which is
23 /// fine for a small file and unacceptable for a multi-GB one; the multipart path
24 /// holds one part at a time. It is also the only path that can carry a file past
25 /// S3's 5 GiB single-PUT ceiling, which is what the tier limits allow for.
26 pub(crate) const MULTIPART_THRESHOLD_BYTES: u64 = 64 * 1024 * 1024;
27
28 /// How many presigned part URLs to request at a time. Must not exceed the
29 /// server's own window cap.
30 const PART_URL_WINDOW: u32 = 100;
31
32 /// An opened multipart upload session.
33 #[derive(Debug, Clone, Deserialize)]
34 pub(crate) struct MultipartStart {
35 pub upload_id: String,
36 pub s3_key: String,
37 pub part_size: u64,
38 pub part_count: u32,
39 pub expires_in: u64,
40 }
41
42 /// One presigned part target, with the exact length the signature binds.
43 #[derive(Debug, Clone, Deserialize)]
44 pub(crate) struct MultipartPartUrl {
45 pub part_number: i32,
46 pub content_length: u64,
47 pub url: String,
48 }
49
50 #[derive(Debug, Deserialize)]
51 struct MultipartPartsResponse {
52 parts: Vec<MultipartPartUrl>,
53 }
54
55 impl MnwApiClient {
56 /// Get a presigned S3 upload URL.
57 pub(crate) async fn presign_upload(
58 &self,
59 user_id: &str,
60 item_id: &str,
61 file_type: &str,
62 file_name: &str,
63 content_type: &str,
64 ) -> anyhow::Result<PresignResponse> {
65 let url = format!("{}/api/internal/upload/presign", self.base_url);
66 let resp = self
67 .http
68 .post(&url)
69 .bearer_auth(&self.service_token)
70 .header("X-MNW-Actor", self.actor_header())
71 .json(&serde_json::json!({
72 "user_id": user_id,
73 "item_id": item_id,
74 "file_type": file_type,
75 "file_name": file_name,
76 "content_type": content_type,
77 }))
78 .send()
79 .await?;
80
81 json_response(resp, "presign_upload").await
82 }
83
84 /// Confirm a completed S3 upload.
85 pub(crate) async fn confirm_upload(
86 &self,
87 user_id: &str,
88 item_id: &str,
89 file_type: &str,
90 s3_key: &str,
91 ) -> anyhow::Result<bool> {
92 let url = format!("{}/api/internal/upload/confirm", self.base_url);
93 let resp = self
94 .http
95 .post(&url)
96 .bearer_auth(&self.service_token)
97 .header("X-MNW-Actor", self.actor_header())
98 .json(&serde_json::json!({
99 "user_id": user_id,
100 "item_id": item_id,
101 "file_type": file_type,
102 "s3_key": s3_key,
103 }))
104 .send()
105 .await?;
106
107 #[derive(Deserialize)]
108 struct Resp {
109 success: bool,
110 }
111 let r: Resp = json_response(resp, "confirm_upload").await?;
112 Ok(r.success)
113 }
114
115 /// Open a multipart upload session and get the part geometry.
116 pub(crate) async fn multipart_start(
117 &self,
118 item_id: &str,
119 file_type: &str,
120 file_name: &str,
121 content_type: &str,
122 file_size_bytes: u64,
123 ) -> anyhow::Result<MultipartStart> {
124 let url = format!("{}/api/internal/upload/multipart/start", self.base_url);
125 let resp = self
126 .http
127 .post(&url)
128 .bearer_auth(&self.service_token)
129 .header("X-MNW-Actor", self.actor_header())
130 .json(&serde_json::json!({
131 "item_id": item_id,
132 "file_type": file_type,
133 "file_name": file_name,
134 "content_type": content_type,
135 "file_size_bytes": file_size_bytes,
136 }))
137 .send()
138 .await?;
139
140 json_response(resp, "multipart_start").await
141 }
142
143 /// Fetch a bounded window of presigned part URLs.
144 async fn multipart_parts(
145 &self,
146 s3_key: &str,
147 upload_id: &str,
148 file_size_bytes: u64,
149 first_part: u32,
150 count: u32,
151 ) -> anyhow::Result<Vec<MultipartPartUrl>> {
152 let url = format!("{}/api/internal/upload/multipart/parts", self.base_url);
153 let resp = self
154 .http
155 .post(&url)
156 .bearer_auth(&self.service_token)
157 .header("X-MNW-Actor", self.actor_header())
158 .json(&serde_json::json!({
159 "s3_key": s3_key,
160 "upload_id": upload_id,
161 "file_size_bytes": file_size_bytes,
162 "first_part": first_part,
163 "count": count,
164 }))
165 .send()
166 .await?;
167
168 let parts: MultipartPartsResponse = json_response(resp, "multipart_parts").await?;
169 Ok(parts.parts)
170 }
171
172 /// Assemble the uploaded parts into the staging object.
173 async fn multipart_complete(
174 &self,
175 s3_key: &str,
176 upload_id: &str,
177 parts: &[(i32, String)],
178 ) -> anyhow::Result<()> {
179 let url = format!("{}/api/internal/upload/multipart/complete", self.base_url);
180 let parts: Vec<serde_json::Value> = parts
181 .iter()
182 .map(|(n, etag)| serde_json::json!({ "part_number": n, "etag": etag }))
183 .collect();
184 let resp = self
185 .http
186 .post(&url)
187 .bearer_auth(&self.service_token)
188 .header("X-MNW-Actor", self.actor_header())
189 .json(&serde_json::json!({
190 "s3_key": s3_key,
191 "upload_id": upload_id,
192 "parts": parts,
193 }))
194 .send()
195 .await?;
196
197 if !resp.status().is_success() {
198 anyhow::bail!(
199 "multipart_complete failed: HTTP {} {}",
200 resp.status(),
201 resp.text().await.unwrap_or_default()
202 );
203 }
204 Ok(())
205 }
206
207 /// Release the parts of an abandoned session. Incomplete multipart uploads
208 /// bill for their parts until aborted.
209 pub(crate) async fn multipart_abort(
210 &self,
211 s3_key: &str,
212 upload_id: &str,
213 ) -> anyhow::Result<()> {
214 let url = format!("{}/api/internal/upload/multipart/abort", self.base_url);
215 let resp = self
216 .http
217 .post(&url)
218 .bearer_auth(&self.service_token)
219 .header("X-MNW-Actor", self.actor_header())
220 .json(&serde_json::json!({ "s3_key": s3_key, "upload_id": upload_id }))
221 .send()
222 .await?;
223
224 if !resp.status().is_success() {
225 anyhow::bail!("multipart_abort failed: HTTP {}", resp.status());
226 }
227 Ok(())
228 }
229
230 /// Upload a file through a multipart session, holding one part in memory at
231 /// a time. Returns the staging key to confirm against.
232 ///
233 /// `on_progress` is called with `(bytes_uploaded, total)` after each part.
234 /// Any failure past the session opening aborts it, so a half-finished upload
235 /// does not leave parts billing indefinitely — the single abort site means a
236 /// future failure path added inside cannot forget to.
237 #[allow(clippy::too_many_arguments)]
238 pub(crate) async fn upload_file_multipart(
239 &self,
240 item_id: &str,
241 file_type: &str,
242 file_name: &str,
243 content_type: &str,
244 file_path: &std::path::Path,
245 file_size: u64,
246 mut on_progress: impl FnMut(u64, u64),
247 ) -> anyhow::Result<String> {
248 let start = self
249 .multipart_start(item_id, file_type, file_name, content_type, file_size)
250 .await?;
251
252 tracing::info!(
253 s3_key = %start.s3_key,
254 part_count = start.part_count,
255 part_size = start.part_size,
256 expires_in = start.expires_in,
257 file_size,
258 "multipart upload session opened"
259 );
260
261 match self
262 .run_multipart_upload(&start, file_path, file_size, &mut on_progress)
263 .await
264 {
265 Ok(()) => Ok(start.s3_key),
266 Err(e) => {
267 if let Err(abort_err) = self.multipart_abort(&start.s3_key, &start.upload_id).await
268 {
269 tracing::warn!(
270 error = %abort_err, s3_key = %start.s3_key,
271 "failed to abort multipart upload after a failed transfer"
272 );
273 }
274 Err(e)
275 }
276 }
277 }
278
279 /// Read and upload every part, then complete. Returns `Err` without
280 /// aborting; the caller owns the single abort.
281 async fn run_multipart_upload(
282 &self,
283 start: &MultipartStart,
284 file_path: &std::path::Path,
285 file_size: u64,
286 on_progress: &mut impl FnMut(u64, u64),
287 ) -> anyhow::Result<()> {
288 use tokio::io::AsyncReadExt;
289
290 let mut file = tokio::fs::File::open(file_path)
291 .await
292 .map_err(|e| anyhow::anyhow!("opening {} for upload: {e}", file_path.display()))?;
293
294 let mut completed: Vec<(i32, String)> = Vec::with_capacity(start.part_count as usize);
295 let mut uploaded: u64 = 0;
296 let mut next: u32 = 1;
297
298 while next <= start.part_count {
299 let count = PART_URL_WINDOW.min(start.part_count - next + 1);
300 let urls = self
301 .multipart_parts(&start.s3_key, &start.upload_id, file_size, next, count)
302 .await?;
303 if urls.is_empty() {
304 anyhow::bail!("server returned no part URLs for part {next}");
305 }
306
307 for part in urls {
308 // One part resident at a time — this is the whole point of the
309 // multipart path over the single-PUT one.
310 let mut buf = vec![0u8; part.content_length as usize];
311 file.read_exact(&mut buf).await.map_err(|e| {
312 anyhow::anyhow!(
313 "reading part {} ({} bytes) from {}: {e}",
314 part.part_number,
315 part.content_length,
316 file_path.display()
317 )
318 })?;
319
320 let etag = self.put_part(&part, buf).await?;
321 completed.push((part.part_number, etag));
322 uploaded += part.content_length;
323 on_progress(uploaded, file_size);
324 next += 1;
325 }
326 }
327
328 self.multipart_complete(&start.s3_key, &start.upload_id, &completed)
329 .await
330 }
331
332 /// PUT one part to its presigned URL, returning the ETag the completion call
333 /// needs. Retries transient failures — losing a part to a network blip
334 /// should not discard the whole transfer.
335 async fn put_part(&self, part: &MultipartPartUrl, body: Vec<u8>) -> anyhow::Result<String> {
336 // `Bytes` so a retry clones a refcount rather than re-copying the part.
337 let body = bytes::Bytes::from(body);
338 let mut attempt: u32 = 0;
339 loop {
340 attempt += 1;
341 let sent = self
342 .http
343 .put(&part.url)
344 .header(reqwest::header::CONTENT_LENGTH, part.content_length)
345 .body(body.clone())
346 .send()
347 .await;
348
349 let retriable = match sent {
350 Ok(resp) if resp.status().is_success() => {
351 let etag = resp
352 .headers()
353 .get(reqwest::header::ETAG)
354 .and_then(|v| v.to_str().ok())
355 .unwrap_or_default()
356 .to_string();
357 if etag.is_empty() {
358 anyhow::bail!(
359 "S3 returned no ETag for part {}; cannot complete the upload",
360 part.part_number
361 );
362 }
363 return Ok(etag);
364 }
365 Ok(resp) => {
366 let status = resp.status();
367 if attempt >= 3 {
368 anyhow::bail!(
369 "part {} failed after {attempt} attempts: HTTP {status}",
370 part.part_number
371 );
372 }
373 format!("HTTP {status}")
374 }
375 Err(e) => {
376 if attempt >= 3 {
377 return Err(anyhow::Error::new(e).context(format!(
378 "part {} failed after {attempt} attempts",
379 part.part_number
380 )));
381 }
382 e.to_string()
383 }
384 };
385
386 let delay_ms = 200u64 * (1u64 << ((attempt - 1) * 2));
387 tracing::warn!(
388 part_number = part.part_number, attempt, delay_ms, error = %retriable,
389 "part upload transient failure, retrying"
390 );
391 tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
392 }
393 }
394
395 /// Upload a file to S3 using a presigned URL.
396 pub(crate) async fn upload_to_s3(
397 &self,
398 presigned_url: &str,
399 file_path: &std::path::Path,
400 content_type: &str,
401 cache_control: Option<&str>,
402 ) -> anyhow::Result<()> {
403 let data = tokio::fs::read(file_path).await?;
404 let mut req = self
405 .http
406 .put(presigned_url)
407 .header("content-type", content_type)
408 .body(data);
409
410 if let Some(cc) = cache_control {
411 req = req.header("cache-control", cc);
412 }
413
414 let resp = req.send().await?;
415
416 if !resp.status().is_success() {
417 anyhow::bail!("S3 upload failed: HTTP {}", resp.status());
418 }
419
420 Ok(())
421 }
422 }
423