| 1 | import { HttpError } from "./util.js"; |
| 2 | |
| 3 | /** R2 multipart parts must all be the same size (except the last) and at least 5 MiB. */ |
| 4 | export const PART_SIZE = 8 * 1024 * 1024; |
| 5 | |
| 6 | type Bucket = Pick<R2Bucket, "put" | "createMultipartUpload">; |
| 7 | |
| 8 | /** |
| 9 | * Writes chunks to one R2 object with bounded memory: objects under one part |
| 10 | * go up with a single put, larger ones as a multipart upload of equal parts. |
| 11 | * Stops with a 413 once `limit` bytes have been written. |
| 12 | */ |
| 13 | export class R2Sink { |
| 14 | private buf: Uint8Array[] = []; |
| 15 | private buffered = 0; |
| 16 | private total = 0; |
| 17 | private upload: R2MultipartUpload | null = null; |
| 18 | private parts: R2UploadedPart[] = []; |
| 19 | |
| 20 | constructor( |
| 21 | private bucket: Bucket, |
| 22 | private key: string, |
| 23 | private contentType: string, |
| 24 | private limit: number, |
| 25 | private tooLarge = "the object is too large", |
| 26 | ) {} |
| 27 | |
| 28 | get size(): number { |
| 29 | return this.total; |
| 30 | } |
| 31 | |
| 32 | async write(chunk: Uint8Array): Promise<void> { |
| 33 | if (!chunk.length) return; |
| 34 | this.total += chunk.length; |
| 35 | if (this.total > this.limit) throw new HttpError(413, this.tooLarge); |
| 36 | this.buf.push(chunk); |
| 37 | this.buffered += chunk.length; |
| 38 | while (this.buffered >= PART_SIZE) await this.flushPart(PART_SIZE); |
| 39 | } |
| 40 | |
| 41 | /** Take exactly `n` buffered bytes (or all of them) off the front. */ |
| 42 | private take(n: number): Uint8Array { |
| 43 | const out = new Uint8Array(Math.min(n, this.buffered)); |
| 44 | let o = 0; |
| 45 | while (o < out.length) { |
| 46 | const head = this.buf[0]!; |
| 47 | const want = out.length - o; |
| 48 | if (head.length <= want) { |
| 49 | out.set(head, o); |
| 50 | o += head.length; |
| 51 | this.buf.shift(); |
| 52 | } else { |
| 53 | out.set(head.subarray(0, want), o); |
| 54 | this.buf[0] = head.subarray(want); |
| 55 | o += want; |
| 56 | } |
| 57 | } |
| 58 | this.buffered -= out.length; |
| 59 | return out; |
| 60 | } |
| 61 | |
| 62 | private async flushPart(n: number): Promise<void> { |
| 63 | this.upload ??= await this.bucket.createMultipartUpload(this.key, { httpMetadata: { contentType: this.contentType } }); |
| 64 | this.parts.push(await this.upload.uploadPart(this.parts.length + 1, this.take(n))); |
| 65 | } |
| 66 | |
| 67 | /** Finish the object; returns its size. */ |
| 68 | async close(): Promise<number> { |
| 69 | if (!this.upload) { |
| 70 | await this.bucket.put(this.key, this.take(this.buffered), { httpMetadata: { contentType: this.contentType } }); |
| 71 | return this.total; |
| 72 | } |
| 73 | if (this.buffered) await this.flushPart(this.buffered); |
| 74 | await this.upload.complete(this.parts); |
| 75 | return this.total; |
| 76 | } |
| 77 | |
| 78 | async abort(): Promise<void> { |
| 79 | this.buf = []; |
| 80 | this.buffered = 0; |
| 81 | await this.upload?.abort().catch(() => undefined); |
| 82 | } |
| 83 | } |