Sign in

flagon-io/g1t

public

g1t/services/packages/src/upload.rs

444 lines17,645 bytesCodeBlame
1//! Writing a blob that arrives in chunks of any size, over any number of
2//! requests.
3//!
4//! R2 takes a multipart upload only when every part but the last is the
5//! same size, and `docker push` cuts its chunks however it likes, so bytes
6//! are gathered into parts of [`PART_BYTES`] here. Whatever is left at the
7//! end of a request (less than a part) is kept as a small object of its
8//! own, the tail, and read back by the next request. A blob that never
9//! fills a part is stored whole, at its digest's key, when it finishes.
10//! The digest is worked out as the bytes pass, so finishing never reads
11//! the blob back.
12
13use serde::{Deserialize, Serialize};
14use worker::Result;
15
16use crate::digest::{Digest, Sha256};
17use crate::store::{BlobStore, Part};
18
19/// The size of every part but the last: over R2's least (5 MiB), and
20/// small enough to hold in memory while it fills.
21pub const PART_BYTES: usize = 10 * 1024 * 1024;
22
23/// How many whole parts `buffered` bytes make, and how many are left over.
24pub fn whole_parts(buffered: usize, part: usize) -> (usize, usize) {
25 (buffered / part, buffered % part)
26}
27
28/// Where an upload's parts and tail are kept, by its id.
29pub fn parts_key(id: &str) -> String {
30 format!("blobs/parts/{id}")
31}
32
33pub fn tail_key(id: &str) -> String {
34 format!("uploads/{id}/tail")
35}
36
37/// An upload's progress, as kept on its row between requests.
38#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
39pub struct Progress {
40 pub id: String,
41 /// Set once the first whole part went up.
42 pub multipart_id: Option<String>,
43 pub parts: Vec<Part>,
44 /// Bytes received so far.
45 pub offset: u64,
46 /// Bytes in the tail object, waiting for the next part.
47 pub tail: u64,
48 #[serde(skip)]
49 pub hasher: Sha256,
50}
51
52impl Progress {
53 pub fn new(id: &str) -> Progress {
54 Progress {
55 id: id.to_owned(),
56 multipart_id: None,
57 parts: Vec::new(),
58 offset: 0,
59 tail: 0,
60 hasher: Sha256::new(),
61 }
62 }
63}
64
65/// How a finished upload ended.
66#[derive(Debug, PartialEq, Eq)]
67pub enum Finished {
68 /// Stored at `key`.
69 Stored { key: String, size: u64 },
70 /// The blob was already stored; these bytes were let go.
71 Duplicate { size: u64 },
72 /// The bytes have another digest than the client said; let go.
73 Mismatch { actual: Digest },
74}
75
76/// An upload being written by one request.
77pub struct Writer<'a, S: BlobStore> {
78 store: &'a S,
79 progress: Progress,
80 buffer: Vec<u8>,
81}
82
83impl<'a, S: BlobStore> Writer<'a, S> {
84 /// Picks up where the last request left off, the tail read back.
85 pub async fn resume(store: &'a S, progress: Progress) -> Result<Writer<'a, S>> {
86 let buffer = if progress.tail > 0 {
87 let tail = store.read(&tail_key(&progress.id)).await?.unwrap_or_default();
88 if tail.len() as u64 != progress.tail {
89 return Err(worker::Error::RustError(format!(
90 "upload {} lost its last {} bytes",
91 progress.id, progress.tail
92 )));
93 }
94 tail
95 } else {
96 Vec::with_capacity(64 * 1024)
97 };
98 Ok(Writer { store, progress, buffer })
99 }
100
101
102 /// Bytes received so far, over every request.
103 pub fn received(&self) -> u64 {
104 self.progress.offset
105 }
106
107 pub async fn write(&mut self, chunk: &[u8]) -> Result<()> {
108 self.progress.hasher.update(chunk);
109 self.progress.offset += chunk.len() as u64;
110 self.buffer.extend_from_slice(chunk);
111 let (whole, _) = whole_parts(self.buffer.len(), PART_BYTES);
112 for _ in 0..whole {
113 let rest = self.buffer.split_off(PART_BYTES);
114 let part = std::mem::replace(&mut self.buffer, rest);
115 self.send_part(part).await?;
116 }
117 Ok(())
118 }
119
120 async fn send_part(&mut self, bytes: Vec<u8>) -> Result<()> {
121 let key = parts_key(&self.progress.id);
122 let upload_id = match &self.progress.multipart_id {
123 Some(id) => id.clone(),
124 None => {
125 let id = self.store.create_multipart(&key).await?;
126 self.progress.multipart_id = Some(id.clone());
127 id
128 }
129 };
130 let number = self.progress.parts.len() as u16 + 1;
131 let part = self.store.upload_part(&key, &upload_id, number, bytes).await?;
132 self.progress.parts.push(part);
133 Ok(())
134 }
135
136 /// Ends this request: what is left is kept as the tail, and the
137 /// progress comes back to be saved. If the tail cannot be kept, the
138 /// upload is let go, so nothing it stored is left behind.
139 pub async fn pause(self) -> Result<Progress> {
140 let Writer { store, mut progress, buffer } = self;
141 let key = tail_key(&progress.id);
142 let kept = if buffer.is_empty() {
143 if progress.tail > 0 { store.delete(&key).await } else { Ok(()) }
144 } else {
145 store.put(&key, buffer.clone()).await
146 };
147 if let Err(error) = kept {
148 let _ = abort(store, &progress).await;
149 return Err(error);
150 }
151 progress.tail = buffer.len() as u64;
152 Ok(progress)
153 }
154
155 /// Gives the upload up part way: a refused or failed request. Everything
156 /// it stored goes, the parts this request sent included.
157 pub async fn abandon(self) -> Result<()> {
158 abort(self.store, &self.progress).await
159 }
160
161 /// Ends the upload. `stored` says the blob with the expected digest is
162 /// already kept, so these bytes are only checked and let go. If storing
163 /// fails, what the upload stored is let go before the error is given.
164 pub async fn finish(self, expected: &Digest, stored: bool) -> Result<Finished> {
165 let store = self.store;
166 let progress = self.progress.clone();
167 match self.finish_inner(expected, stored).await {
168 Ok(finished) => Ok(finished),
169 Err(error) => {
170 let _ = abort(store, &progress).await;
171 Err(error)
172 }
173 }
174 }
175
176 async fn finish_inner(self, expected: &Digest, stored: bool) -> Result<Finished> {
177 let Writer { store, progress, buffer } = self;
178 let size = progress.offset;
179 let actual = progress.hasher.clone().finish();
180 if &actual != expected || stored {
181 abort(store, &progress).await?;
182 return Ok(if &actual != expected { Finished::Mismatch { actual } } else { Finished::Duplicate { size } });
183 }
184 let key = match &progress.multipart_id {
185 None => {
186 let key = expected.object_key();
187 store.put(&key, buffer).await?;
188 key
189 }
190 Some(upload_id) => {
191 let key = parts_key(&progress.id);
192 let mut parts = progress.parts.clone();
193 if !buffer.is_empty() {
194 let number = parts.len() as u16 + 1;
195 parts.push(store.upload_part(&key, upload_id, number, buffer).await?);
196 }
197 store.complete_multipart(&key, upload_id, &parts).await?;
198 key
199 }
200 };
201 if progress.tail > 0 {
202 store.delete(&tail_key(&progress.id)).await?;
203 }
204 Ok(Finished::Stored { key, size })
205 }
206}
207
208/// What letting go of an unfinished upload takes.
209#[derive(Debug, PartialEq, Eq)]
210pub struct Cleanup {
211 /// The multipart upload to abort, by its key and id: started by this
212 /// request or an earlier one.
213 pub multipart: Option<(String, String)>,
214 /// The tail kept between requests, to delete.
215 pub tail: Option<String>,
216}
217
218/// What an unfinished upload has stored, from its progress as it stands.
219/// A tail is only ever written when a request pauses, so `progress.tail`
220/// says whether one is there whatever the current request did.
221pub fn cleanup(progress: &Progress) -> Cleanup {
222 Cleanup {
223 multipart: progress
224 .multipart_id
225 .as_ref()
226 .map(|id| (parts_key(&progress.id), id.clone())),
227 tail: (progress.tail > 0).then(|| tail_key(&progress.id)),
228 }
229}
230
231/// Lets go of everything an unfinished upload stored. Both steps are tried
232/// even when one fails.
233pub async fn abort<S: BlobStore>(store: &S, progress: &Progress) -> Result<()> {
234 let Cleanup { multipart, tail } = cleanup(progress);
235 let aborted = match &multipart {
236 Some((key, id)) => store.abort_multipart(key, id).await,
237 None => Ok(()),
238 };
239 let deleted = match &tail {
240 Some(key) => store.delete(key).await,
241 None => Ok(()),
242 };
243 aborted.and(deleted)
244}
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249
250 #[test]
251 fn bytes_make_whole_parts_and_a_tail() {
252 assert_eq!(whole_parts(0, PART_BYTES), (0, 0));
253 assert_eq!(whole_parts(PART_BYTES - 1, PART_BYTES), (0, PART_BYTES - 1));
254 assert_eq!(whole_parts(PART_BYTES, PART_BYTES), (1, 0));
255 assert_eq!(whole_parts(2 * PART_BYTES + 5, PART_BYTES), (2, 5));
256 const { assert!(PART_BYTES >= 5 * 1024 * 1024, "R2's least part") };
257 }
258
259 use std::cell::RefCell;
260 use std::collections::HashMap;
261 use std::future::Future;
262 use std::pin::pin;
263 use std::task::{Context, Poll, Waker};
264
265 use worker::ResponseBody;
266
267 use crate::range::Wanted;
268 use crate::store::Got;
269
270 /// Multipart uploads in progress, by id: their parts so far.
271 type Uploads = HashMap<String, Vec<(u16, Vec<u8>)>>;
272
273 /// A store in memory that keeps multipart uploads the way R2 does,
274 /// refusing parts of different sizes.
275 #[derive(Default)]
276 struct Memory {
277 objects: RefCell<HashMap<String, Vec<u8>>>,
278 parts: RefCell<Uploads>,
279 puts: RefCell<u32>,
280 }
281
282 impl BlobStore for Memory {
283 async fn put(&self, key: &str, bytes: Vec<u8>) -> Result<()> {
284 *self.puts.borrow_mut() += 1;
285 self.objects.borrow_mut().insert(key.to_owned(), bytes);
286 Ok(())
287 }
288 async fn get(&self, key: &str, _range: Option<Wanted>) -> Result<Option<Got>> {
289 Ok(self.objects.borrow().get(key).map(|b| Got { size: b.len() as u64, body: ResponseBody::Body(b.clone()) }))
290 }
291 async fn head(&self, key: &str) -> Result<Option<u64>> {
292 Ok(self.objects.borrow().get(key).map(|b| b.len() as u64))
293 }
294 async fn delete(&self, key: &str) -> Result<()> {
295 self.objects.borrow_mut().remove(key);
296 Ok(())
297 }
298 async fn create_multipart(&self, key: &str) -> Result<String> {
299 self.parts.borrow_mut().insert(format!("mp-{key}"), Vec::new());
300 Ok(format!("mp-{key}"))
301 }
302 async fn upload_part(&self, _key: &str, upload_id: &str, number: u16, bytes: Vec<u8>) -> Result<Part> {
303 self.parts.borrow_mut().get_mut(upload_id).expect("started").push((number, bytes));
304 Ok(Part { number, etag: format!("e{number}") })
305 }
306 async fn complete_multipart(&self, key: &str, upload_id: &str, parts: &[Part]) -> Result<()> {
307 let uploaded = self.parts.borrow_mut().remove(upload_id).expect("started");
308 assert_eq!(uploaded.len(), parts.len());
309 let sizes: Vec<usize> = uploaded.iter().map(|(_, b)| b.len()).collect();
310 if let Some((_, rest)) = sizes.split_last() {
311 assert!(rest.iter().all(|size| *size == PART_BYTES), "every part but the last is the same size: {sizes:?}");
312 }
313 let joined = uploaded.into_iter().flat_map(|(_, b)| b).collect();
314 self.objects.borrow_mut().insert(key.to_owned(), joined);
315 Ok(())
316 }
317 async fn abort_multipart(&self, _key: &str, upload_id: &str) -> Result<()> {
318 self.parts.borrow_mut().remove(upload_id);
319 Ok(())
320 }
321 fn presign_get(&self, _key: &str, _expires: u32, _now_ms: u64) -> Option<String> {
322 None
323 }
324 }
325
326 /// Runs a future that never waits on anything outside memory.
327 fn run<F: Future>(future: F) -> F::Output {
328 let mut future = pin!(future);
329 match future.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
330 Poll::Ready(output) => output,
331 Poll::Pending => panic!("the in-memory store never waits"),
332 }
333 }
334
335 fn bytes(length: usize) -> Vec<u8> {
336 (0..length).map(|i| (i * 31 % 251) as u8).collect()
337 }
338
339 /// Sends `data` in `requests`, each cut into the chunks given, pausing
340 /// and saving the progress (hasher state too) between requests.
341 fn upload(store: &Memory, data: &[u8], requests: &[&[usize]]) -> Finished {
342 let mut progress = Progress::new("upl_1");
343 let mut at = 0;
344 for chunks in requests {
345 let mut writer = run(Writer::resume(store, progress)).unwrap();
346 for size in *chunks {
347 let end = (at + size).min(data.len());
348 run(writer.write(&data[at..end])).unwrap();
349 at = end;
350 }
351 progress = run(writer.pause()).unwrap();
352 let text = serde_json::to_string(&progress).unwrap();
353 let saved = progress.hasher.save();
354 progress = serde_json::from_str(&text).unwrap();
355 progress.hasher = Sha256::restore(&saved).unwrap();
356 }
357 assert_eq!(at, data.len());
358 let writer = run(Writer::resume(store, progress)).unwrap();
359 run(writer.finish(&Digest::of(data), false)).unwrap()
360 }
361
362 #[test]
363 fn a_small_blob_is_stored_whole_at_its_digest() {
364 let store = Memory::default();
365 let data = bytes(100_000);
366 let finished = upload(&store, &data, &[&[1000, 4000], &[95_000]]);
367 let key = Digest::of(&data).object_key();
368 assert_eq!(finished, Finished::Stored { key: key.clone(), size: 100_000 });
369 assert_eq!(store.objects.borrow()[&key], data);
370 assert!(!store.objects.borrow().contains_key(&tail_key("upl_1")), "the tail is let go");
371 }
372
373 #[test]
374 fn chunks_of_any_size_become_parts_of_one_size() {
375 let store = Memory::default();
376 let data = bytes(2 * PART_BYTES + 12_345);
377 // Uneven chunks over three requests, one ending inside a part.
378 let finished = upload(&store, &data, &[&[3 * 1024 * 1024, 9 * 1024 * 1024], &[700_000], &[PART_BYTES * 2]]);
379 let key = parts_key("upl_1");
380 assert_eq!(finished, Finished::Stored { key: key.clone(), size: data.len() as u64 });
381 assert_eq!(store.objects.borrow()[&key], data);
382 assert_eq!(store.objects.borrow().len(), 1, "nothing else is left behind");
383 }
384
385 #[test]
386 fn a_wrong_digest_or_a_blob_already_kept_lets_the_bytes_go() {
387 let store = Memory::default();
388 let data = bytes(PART_BYTES + 10);
389 let mut writer = run(Writer::resume(&store, Progress::new("upl_2"))).unwrap();
390 run(writer.write(&data)).unwrap();
391 let finished = run(writer.finish(&Digest::of(b"something else"), false)).unwrap();
392 assert_eq!(finished, Finished::Mismatch { actual: Digest::of(&data) });
393 assert!(store.parts.borrow().is_empty(), "the multipart upload is aborted");
394 let mut writer = run(Writer::resume(&store, Progress::new("upl_3"))).unwrap();
395 run(writer.write(&data[..10])).unwrap();
396 assert_eq!(run(writer.finish(&Digest::of(&data[..10]), true)).unwrap(), Finished::Duplicate { size: 10 });
397 assert!(store.objects.borrow().is_empty());
398 }
399
400 #[test]
401 fn letting_go_takes_the_multipart_and_the_tail_there_are() {
402 let mut progress = Progress::new("upl_9");
403 assert_eq!(cleanup(&progress), Cleanup { multipart: None, tail: None });
404 progress.tail = 5;
405 assert_eq!(cleanup(&progress), Cleanup { multipart: None, tail: Some(tail_key("upl_9")) });
406 progress.multipart_id = Some("mp".into());
407 progress.tail = 0;
408 assert_eq!(
409 cleanup(&progress),
410 Cleanup { multipart: Some((parts_key("upl_9"), "mp".into())), tail: None }
411 );
412 }
413
414 #[test]
415 fn a_request_given_up_part_way_leaves_nothing_behind() {
416 let store = Memory::default();
417 let data = bytes(PART_BYTES + 100);
418 // An earlier request left a tail.
419 let mut writer = run(Writer::resume(&store, Progress::new("upl_4"))).unwrap();
420 run(writer.write(&data[..100])).unwrap();
421 let progress = run(writer.pause()).unwrap();
422 assert!(store.objects.borrow().contains_key(&tail_key("upl_4")));
423 // This one starts the multipart upload, then is refused.
424 let mut writer = run(Writer::resume(&store, progress)).unwrap();
425 run(writer.write(&data[100..])).unwrap();
426 assert_eq!(store.parts.borrow().len(), 1, "a part went up");
427 run(writer.abandon()).unwrap();
428 assert!(store.parts.borrow().is_empty(), "the multipart upload is aborted");
429 assert!(store.objects.borrow().is_empty(), "the tail is deleted");
430 }
431
432 #[test]
433 fn progress_round_trips_without_its_hasher() {
434 let mut progress = Progress::new("upl_1");
435 progress.parts.push(Part { number: 1, etag: "\"e1\"".into() });
436 progress.offset = 42;
437 let text = serde_json::to_string(&progress).unwrap();
438 let back: Progress = serde_json::from_str(&text).unwrap();
439 assert_eq!(back.parts, progress.parts);
440 assert_eq!(back.offset, 42);
441 assert_eq!(parts_key("upl_1"), "blobs/parts/upl_1");
442 assert_eq!(tail_key("upl_1"), "uploads/upl_1/tail");
443 }
444}