Open History to see the last commit
1// @effect-diagnostics cryptoRandomUUID:off globalConsole:off - vendored upstream protocol, not an Effect service.
2import { cborResponse, decodeEnvelope, type StoreMetadata } from './cbor.ts';
3import {
4 abandon,
5 publish,
6 recordMultipart,
7 reserve,
8 type Generation,
9 type Scope,
10} from './database.ts';
11import { badRequest, tooLarge, unavailable } from './errors.ts';
12import { PART_SIZE, type Limits } from './limits.ts';
13import { boundaryFrom, multipart } from './multipart.ts';
14import { collect, contentLength, Deadline, Input } from './streams.ts';
15import type { WriteIdentity } from './auth.ts';
16import type { Observations } from './observations.ts';
17
18async function uploadBlob(
19 env: Env,
20 generation: Generation,
21 source: AsyncIterable<Uint8Array>,
22 limit: number,
23 deadline: Deadline,
24 stats: Observations,
25): Promise<void> {
26 let buffer = new Uint8Array(PART_SIZE);
27 let used = 0;
28 let upload: R2MultipartUpload | undefined;
29 const parts: R2UploadedPart[] = [];
30 for await (const chunk of source) {
31 generation.blob_size += chunk.length;
32 if (generation.blob_size > limit) tooLarge();
33 let offset = 0;
34 while (offset < chunk.length) {
35 if (used === PART_SIZE) {
36 if (!upload) {
37 stats.r2_operations++;
38 upload = await deadline.run(env.ARTIFACTS.createMultipartUpload(generation.blob_object));
39 try {
40 await deadline.run(recordMultipart(env.INDEX, generation, upload.uploadId, stats));
41 } catch (error) {
42 // Also cover creation succeeding but D1 failing to record its ID.
43 try {
44 await deadline.run(upload.abort());
45 } catch {
46 /* Lifecycle is the final backstop. */
47 }
48 throw error;
49 }
50 }
51 stats.r2_operations++;
52 parts.push(await deadline.run(upload.uploadPart(parts.length + 1, buffer)));
53 buffer = new Uint8Array(PART_SIZE);
54 used = 0;
55 }
56 const size = Math.min(PART_SIZE - used, chunk.length - offset);
57 buffer.set(chunk.subarray(offset, offset + size), used);
58 used += size;
59 offset += size;
60 }
61 }
62 if (upload) {
63 stats.r2_operations += Number(used > 0) + 1;
64 if (used)
65 parts.push(await deadline.run(upload.uploadPart(parts.length + 1, buffer.subarray(0, used))));
66 await deadline.run(upload.complete(parts));
67 } else {
68 stats.r2_operations++;
69 const result = await deadline.run(
70 env.ARTIFACTS.put(generation.blob_object, buffer.subarray(0, used)),
71 );
72 if (!result) unavailable();
73 }
74 generation.blob_id = crypto.randomUUID();
75}
76
77export async function store(
78 request: Request,
79 env: Env,
80 ctx: Pick<ExecutionContext, 'waitUntil'>,
81 scope: Scope,
82 identity: WriteIdentity,
83 limits: Limits,
84 deadline: Deadline,
85 stats: Observations,
86): Promise<Response> {
87 const boundary = boundaryFrom(request.headers.get('Content-Type'));
88 const length = contentLength(request, limits.store);
89 const generation = await deadline.run(
90 reserve(env.INDEX, scope, identity.exp, length ?? limits.store, stats),
91 );
92 let input: Input | undefined;
93 try {
94 input = new Input(request.body, Math.min(length ?? limits.store, limits.store), deadline);
95 let metadata: StoreMetadata | undefined;
96 for await (const part of multipart(input, boundary, limits.headers)) {
97 if (part.name === 'metadata') {
98 metadata = decodeEnvelope(await collect(part.body, limits.metadata), true, limits);
99 generation.value_size = metadata.value.length;
100 stats.r2_operations++;
101 const result = await deadline.run(
102 env.ARTIFACTS.put(generation.value_object, metadata.value),
103 );
104 if (!result) unavailable();
105 } else {
106 await uploadBlob(env, generation, part.body, limits.blob, deadline, stats);
107 }
108 }
109 if (!metadata || (length !== null && input.bytes !== length)) badRequest();
110 deadline.check();
111 await deadline.run(publish(env.INDEX, generation, metadata.key, metadata.secondary_key, stats));
112 return cborResponse({ blob_id: generation.blob_id });
113 } catch (error) {
114 // Publication is never deferred. Only cleanup can continue after the response.
115 ctx.waitUntil(
116 abandon(env.INDEX, generation.generation_id).catch(() => {
117 console.warn(
118 JSON.stringify({
119 operation: 'cleanup',
120 error: 'abandon_failed',
121 generation: generation.generation_id,
122 }),
123 );
124 }),
125 );
126 throw error;
127 } finally {
128 stats.request_bytes = input?.bytes ?? 0;
129 input?.close();
130 }
131}