Implement R2 multipart uploads

Closes DEVX-593

Author: mrbbot

packages/tre/src/index.ts

@@ -33,6 +33,7 @@ import {
   HEADER_CF_BLOB,
   HEADER_PROBE,
   PLUGIN_ENTRIES,
+  Persistence,
   Plugins,
   SERVICE_ENTRY,
   SOCKET_ENTRY,
@@ -63,17 +64,18 @@ import {
   serializeConfig,
 } from "./runtime";
 import {
-  Clock,
   HttpError,
   Log,
   MiniflareCoreError,
   Mutex,
   NoOpLog,
   OptionalZodTypeOf,
-  defaultClock,
+  Timers,
+  defaultTimers,
   formatResponse,
 } from "./shared";
 import { anyAbortSignal } from "./shared/signal";
+import { NewStorage } from "./storage2";
 import { waitForRequest } from "./wait";
 
 // ===== `Miniflare` User Options =====
@@ -262,7 +264,7 @@ export class Miniflare {
   #sharedOpts: PluginSharedOptions;
   #workerOpts: PluginWorkerOptions[];
   #log: Log;
-  readonly #clock: Clock;
+  readonly #timers: Timers;
 
   #runtime?: Runtime;
   #removeRuntimeExitHook?: () => void;
@@ -314,7 +316,7 @@ export class Miniflare {
     this.#sharedOpts = sharedOpts;
     this.#workerOpts = workerOpts;
     this.#log = this.#sharedOpts.core.log ?? new NoOpLog();
-    this.#clock = this.#sharedOpts.core.clock ?? defaultClock;
+    this.#timers = this.#sharedOpts.core.timers ?? defaultTimers;
     this.#initPlugins();
 
     this.#liveReloadServer = new WebSocketServer({ noServer: true });
@@ -360,7 +362,7 @@ export class Miniflare {
       if (plugin.gateway !== undefined && plugin.router !== undefined) {
         const gatewayFactory = new GatewayFactory<any>(
           this.#log,
-          this.#clock,
+          this.#timers,
           this.#sharedOpts.core.cloudflareFetch,
           key,
           plugin.gateway,
@@ -852,6 +854,17 @@ export class Miniflare {
     return response;
   }
 
+  /** @internal */
+  _getPluginStorage(
+    plugin: keyof Plugins,
+    namespace: string,
+    persist?: Persistence
+  ): NewStorage {
+    const factory = this.#gatewayFactories[plugin];
+    assert(factory !== undefined);
+    return factory.getStorage(namespace, persist).getNewStorage();
+  }
+
   async dispose(): Promise<void> {
     this.#disposeController.abort();
     try {

packages/tre/src/plugins/cache/gateway.ts

@@ -4,7 +4,7 @@ import http from "http";
 import { ReadableStream, TransformStream } from "stream/web";
 import CachePolicy from "http-cache-semantics";
 import { Headers, HeadersInit, Request, Response, fetch } from "../../http";
-import { Clock, DeferredPromise, Log } from "../../shared";
+import { DeferredPromise, Log, Timers } from "../../shared";
 import { Storage } from "../../storage";
 import {
   InclusiveRange,
@@ -26,7 +26,7 @@ interface CacheMetadata {
   size: number;
 }
 
-function getExpiration(clock: Clock, req: Request, res: Response) {
+function getExpiration(timers: Timers, req: Request, res: Response) {
   // Cloudflare ignores request Cache-Control
   const reqHeaders = normaliseHeaders(req.headers);
   delete reqHeaders["cache-control"];
@@ -59,7 +59,7 @@ function getExpiration(clock: Clock, req: Request, res: Response) {
   // @ts-expect-error `now` isn't included in CachePolicy's type definitions
   const originalNow = CachePolicy.prototype.now;
   // @ts-expect-error `now` isn't included in CachePolicy's type definitions
-  CachePolicy.prototype.now = clock;
+  CachePolicy.prototype.now = timers.now;
   try {
     const policy = new CachePolicy(cacheReq, cacheRes, { shared: true });
 
@@ -230,9 +230,13 @@ class SizingStream extends TransformStream<Uint8Array, Uint8Array> {
 export class CacheGateway {
   private readonly storage: KeyValueStorage<CacheMetadata>;
 
-  constructor(log: Log, legacyStorage: Storage, private readonly clock: Clock) {
+  constructor(
+    private readonly log: Log,
+    legacyStorage: Storage,
+    private readonly timers: Timers
+  ) {
     const storage = legacyStorage.getNewStorage();
-    this.storage = new KeyValueStorage(storage, clock);
+    this.storage = new KeyValueStorage(storage, timers);
   }
 
   async match(request: Request, cacheKey?: string): Promise<Response> {
@@ -291,7 +295,7 @@ export class CacheGateway {
     assert(body !== null);
 
     const { storable, expiration, headers } = getExpiration(
-      this.clock,
+      this.timers,
       request,
       response
     );
@@ -321,7 +325,7 @@ export class CacheGateway {
     await this.storage.put({
       key: cacheKey,
       value: body,
-      expiration: this.clock() + expiration,
+      expiration: this.timers.now() + expiration,
       metadata,
     });
     return new Response(null, { status: 204 });

packages/tre/src/plugins/core/index.ts

@@ -16,6 +16,7 @@ import {
   Log,
   LogLevel,
   MiniflareCoreError,
+  Timers,
 } from "../../shared";
 import { getCacheServiceName } from "../cache";
 import { DURABLE_OBJECTS_STORAGE_SERVICE_NAME } from "../do";
@@ -68,7 +69,7 @@ export const CoreSharedOptionsSchema = z.object({
   verbose: z.boolean().optional(),
 
   log: z.instanceof(Log).optional(),
-  clock: z.function().returns(z.number()).optional(),
+  timers: z.custom<Timers>().optional(),
   cloudflareFetch: CloudflareFetchSchema.optional(),
 
   // TODO: add back validation of cf object

packages/tre/src/plugins/do/gateway.ts

@@ -1,11 +1,11 @@
-import { Clock, Log } from "../../shared";
+import { Log, Timers } from "../../shared";
 import { Storage } from "../../storage";
 
 export class DurableObjectsStorageGateway {
   constructor(
     private readonly log: Log,
     private readonly storage: Storage,
-    private readonly clock: Clock
+    private readonly timers: Timers
   ) {}
 
   async get(_key: string) {

packages/tre/src/plugins/kv/gateway.ts

@@ -1,8 +1,8 @@
 import { ReadableStream, TransformStream } from "stream/web";
 import {
-  Clock,
   HttpError,
   Log,
+  Timers,
   maybeApply,
   millisToSeconds,
   secondsToMillis,
@@ -159,10 +159,10 @@ export class KVGateway {
   constructor(
     private readonly log: Log,
     legacyStorage: Storage,
-    private readonly clock: Clock
+    private readonly timers: Timers
   ) {
     const storage = legacyStorage.getNewStorage();
-    this.storage = new KeyValueStorage(storage, clock);
+    this.storage = new KeyValueStorage(storage, timers);
   }
 
   async get<Metadata = unknown>(
@@ -187,7 +187,7 @@ export class KVGateway {
     validateKey(key);
 
     // Normalise and validate expiration
-    const now = millisToSeconds(this.clock());
+    const now = millisToSeconds(this.timers.now());
     let expiration = normaliseInt(options.expiration);
     const expirationTtl = normaliseInt(options.expirationTtl);
     if (expirationTtl !== undefined) {

packages/tre/src/plugins/kv/remote.ts

@@ -2,7 +2,7 @@ import assert from "assert";
 import { Blob } from "buffer";
 import { z } from "zod";
 import { BodyInit, FormData, Response } from "../../http";
-import { Clock, millisToSeconds } from "../../shared";
+import { Timers, millisToSeconds } from "../../shared";
 import {
   RemoteStorage,
   StorageListOptions,
@@ -78,16 +78,17 @@ const DEFAULT_CACHE_TTL = 60;
 // Returns seconds since UNIX epoch key should expire, using the specified
 // expiration only if it is sooner than the cache TTL
 function getCacheExpiration(
-  clock: Clock,
+  timers: Timers,
   expiration?: number,
   cacheTtl = DEFAULT_CACHE_TTL
 ): number {
   // Return minimum expiration
-  const cacheExpiration = millisToSeconds(clock()) + cacheTtl;
+  const cacheExpiration = millisToSeconds(timers.now()) + cacheTtl;
   if (expiration === undefined || isNaN(expiration)) return cacheExpiration;
   else return Math.min(cacheExpiration, expiration);
 }
 
+// TODO(soon): remove all this
 export class KVRemoteStorage extends RemoteStorage {
   async get(
     key: string,
@@ -100,7 +101,7 @@ export class KVRemoteStorage extends RemoteStorage {
       // cacheTtl may have changed between the original get call that cached
       // this value and now, so check the cache is still fresh with the new TTL
       const newExpiration = cachedValue.metadata.storedAt + cacheTtl;
-      if (newExpiration >= millisToSeconds(this.clock())) {
+      if (newExpiration >= millisToSeconds(this.timers.now())) {
         // If the cache is still fresh, update the expiration and return
         await this.cache.put<RemoteCacheMetadata>(key, {
           value: cachedValue.value,
@@ -157,9 +158,9 @@ export class KVRemoteStorage extends RemoteStorage {
     const result: StoredValueMeta = { value, expiration, metadata };
     await this.cache.put<RemoteCacheMetadata>(key, {
       value: result.value,
-      expiration: getCacheExpiration(this.clock, expiration, cacheTtl),
+      expiration: getCacheExpiration(this.timers, expiration, cacheTtl),
       metadata: {
-        storedAt: millisToSeconds(this.clock()),
+        storedAt: millisToSeconds(this.timers.now()),
         actualExpiration: result.expiration,
         actualMetadata: result.metadata,
       },
@@ -176,7 +177,7 @@ export class KVRemoteStorage extends RemoteStorage {
     if (value.expiration !== undefined) {
       // Send expiration as TTL to avoid "expiration times must be at least 60s
       // in the future" issues from clock skew when setting `expirationTtl: 60`.
-      const desiredTtl = value.expiration - millisToSeconds(this.clock());
+      const desiredTtl = value.expiration - millisToSeconds(this.timers.now());
       const ttl = Math.max(desiredTtl, 60);
       searchParams.set("expiration_ttl", ttl.toString());
     }
@@ -197,9 +198,9 @@ export class KVRemoteStorage extends RemoteStorage {
     // Store this value in the cache
     await this.cache.put<RemoteCacheMetadata>(key, {
       value: value.value,
-      expiration: getCacheExpiration(this.clock, value.expiration),
+      expiration: getCacheExpiration(this.timers, value.expiration),
       metadata: {
-        storedAt: millisToSeconds(this.clock()),
+        storedAt: millisToSeconds(this.timers.now()),
         actualExpiration: value.expiration,
         actualMetadata: value.metadata,
       },
@@ -219,8 +220,11 @@ export class KVRemoteStorage extends RemoteStorage {
     // "Store" delete in cache as tombstone
     await this.cache.put<RemoteCacheMetadata>(key, {
       value: new Uint8Array(),
-      expiration: getCacheExpiration(this.clock),
-      metadata: { storedAt: millisToSeconds(this.clock()), tombstone: true },
+      expiration: getCacheExpiration(this.timers),
+      metadata: {
+        storedAt: millisToSeconds(this.timers.now()),
+        tombstone: true,
+      },
     });
 
     // Technically, it's incorrect to always say we deleted the key by returning

packages/tre/src/plugins/r2/errors.ts

@@ -14,13 +14,17 @@ enum CfCode {
   InternalError = 10001,
   NoSuchObjectKey = 10007,
   EntityTooLarge = 100100,
+  EntityTooSmall = 10011,
   MetadataTooLarge = 10012,
   InvalidObjectName = 10020,
   InvalidMaxKeys = 10022,
+  NoSuchUpload = 10024,
+  InvalidPart = 10025,
   InvalidArgument = 10029,
   PreconditionFailed = 10031,
   BadDigest = 10037,
   InvalidRange = 10039,
+  BadUpload = 10048,
 }
 
 export class R2Error extends HttpError {
@@ -110,6 +114,16 @@ export class EntityTooLarge extends R2Error {
   }
 }
 
+export class EntityTooSmall extends R2Error {
+  constructor() {
+    super(
+      Status.BadRequest,
+      "Your proposed upload is smaller than the minimum allowed object size.",
+      CfCode.EntityTooSmall
+    );
+  }
+}
+
 export class MetadataTooLarge extends R2Error {
   constructor() {
     super(
@@ -160,6 +174,26 @@ export class InvalidMaxKeys extends R2Error {
   }
 }
 
+export class NoSuchUpload extends R2Error {
+  constructor() {
+    super(
+      Status.BadRequest,
+      "The specified multipart upload does not exist.",
+      CfCode.NoSuchUpload
+    );
+  }
+}
+
+export class InvalidPart extends R2Error {
+  constructor() {
+    super(
+      Status.BadRequest,
+      "One or more of the specified parts could not be found.",
+      CfCode.InvalidPart
+    );
+  }
+}
+
 export class PreconditionFailed extends R2Error {
   constructor() {
     super(
@@ -179,3 +213,13 @@ export class InvalidRange extends R2Error {
     );
   }
 }
+
+export class BadUpload extends R2Error {
+  constructor() {
+    super(
+      Status.RangeNotSatisfiable,
+      "There was a problem with the multipart upload.",
+      CfCode.BadUpload
+    );
+  }
+}

packages/tre/src/plugins/r2/gateway.ts

@@ -96,7 +96,7 @@ export class R2ObjectBody extends R2Object {
 
   encode(): EncodedMetadata {
     const { metadataSize, value: metadata } = super.encode();
-    const identity = new TransformStream();
+    const identity = new TransformStream<Uint8Array, Uint8Array>();
     void metadata
       .pipeTo(identity.writable, { preventClose: true })
       .then(() => this.body.pipeTo(identity.writable));

packages/tre/src/plugins/r2/r2Object.ts

@@ -21,7 +21,13 @@ async function decodeMetadata(req: Request) {
   assert(req.body !== null);
   const body = req.body as ReadableStream<Uint8Array>;
 
-  // Read just metadata from body stream
+  // Read just metadata from body stream (NOTE: we can't use a `TransformStream`
+  // and buffer the first N chunks as we need this metadata to determine what
+  // to do with the rest of the body. We have to *pull* the data as opposed to
+  // passively transforming it as it's piped somewhere else. If `body` were
+  // a byte stream, we could use BYOB reads to read just enough. Even better, if
+  // this were running in the Workers runtime, we could use `readAtLeast()` to
+  // read everything at once.)
   const chunks: Uint8Array[] = [];
   let chunksLength = 0;
   for await (const chunk of body.values({ preventCancel: true })) {
@@ -40,7 +46,7 @@ async function decodeMetadata(req: Request) {
   // If we read some value when reading metadata (quite likely), create a new
   // stream, write the bit we read, then write the rest of the body stream
   if (chunksLength > metadataSize) {
-    const identity = new TransformStream();
+    const identity = new TransformStream<Uint8Array, Uint8Array>();
     const writer = identity.writable.getWriter();
     // The promise returned by `writer.write()` will only resolve once the chunk
     // is read, which won't be until after this function returns, so we can't
@@ -77,6 +83,16 @@ function encodeResult(result: R2Object | R2ObjectBody | R2Objects) {
   });
 }
 
+function encodeJSONResult(result: unknown) {
+  const encoded = JSON.stringify(result);
+  return new Response(encoded, {
+    headers: {
+      [CfHeader.MetadataSize]: `${Buffer.byteLength(encoded)}`,
+      "Content-Type": "application/json",
+    },
+  });
+}
+
 export interface R2Params {
   bucket: string;
 }
@@ -127,6 +143,35 @@ export class R2Router extends Router<R2Gateway> {
         metadata
       );
       return encodeResult(result);
+    } else if (metadata.method === "createMultipartUpload") {
+      const result = await gateway.createMultipartUpload(
+        metadata.object,
+        metadata
+      );
+      return encodeJSONResult(result);
+    } else if (metadata.method === "uploadPart") {
+      const contentLength = Number(req.headers.get("Content-Length"));
+      // `workerd` requires a known value size for R2 put requests as above
+      assert(!isNaN(contentLength));
+      const valueSize = contentLength - metadataSize;
+      const result = await gateway.uploadPart(
+        metadata.object,
+        metadata.uploadId,
+        metadata.partNumber,
+        value,
+        valueSize
+      );
+      return encodeJSONResult(result);
+    } else if (metadata.method === "completeMultipartUpload") {
+      const result = await gateway.completeMultipartUpload(
+        metadata.object,
+        metadata.uploadId,
+        metadata.parts
+      );
+      return encodeResult(result);
+    } else if (metadata.method === "abortMultipartUpload") {
+      await gateway.abortMultipartUpload(metadata.object, metadata.uploadId);
+      return new Response();
     } else {
       throw new InternalError(); // Unknown method: should never be reached
     }

packages/tre/src/plugins/r2/router.ts

@@ -1,4 +1,5 @@
 import { z } from "zod";
+import { ValueOf } from "../../shared";
 
 export interface ObjectRow {
   key: string;
@@ -11,6 +12,30 @@ export interface ObjectRow {
   http_metadata: string; // JSON-serialised `R2HTTPMetadata` (workers-types)
   custom_metadata: string; // JSON-serialised user-defined metadata
 }
+export const MultipartUploadState = {
+  IN_PROGRESS: 0,
+  COMPLETED: 1,
+  ABORTED: 2,
+} as const;
+export interface MultipartUploadRow {
+  upload_id: string;
+  key: string;
+  http_metadata: string; // JSON-serialised `R2HTTPMetadata` (workers-types)
+  custom_metadata: string; // JSON-serialised user-defined metadata
+  state: ValueOf<typeof MultipartUploadState>;
+  // NOTE: we need to keep completed/aborted uploads around for referential
+  // integrity, and because error messages are different when attempting to
+  // upload parts to them
+}
+export interface MultipartPartRow {
+  upload_id: string;
+  part_number: number;
+  blob_id: string;
+  size: number; // NOTE: used to identify which parts to read for range requests
+  etag: string; // NOTE: multipart part ETag's are not MD5 checksums
+  checksum_md5: string; // NOTE: used in construction of final object's ETag
+  object_key: string | null; // null if in-progress upload
+}
 export const SQL_SCHEMA = `
 CREATE TABLE IF NOT EXISTS _mf_objects (
     key TEXT PRIMARY KEY,
@@ -23,7 +48,27 @@ CREATE TABLE IF NOT EXISTS _mf_objects (
     http_metadata TEXT NOT NULL,
     custom_metadata TEXT NOT NULL
 );
+CREATE TABLE IF NOT EXISTS _mf_multipart_uploads (
+    upload_id TEXT PRIMARY KEY,
+    key TEXT NOT NULL,
+    http_metadata TEXT NOT NULL,
+    custom_metadata TEXT NOT NULL,
+    state TINYINT DEFAULT 0 NOT NULL
+);
+CREATE TABLE IF NOT EXISTS _mf_multipart_parts (
+    upload_id TEXT NOT NULL REFERENCES _mf_multipart_uploads(upload_id),
+    part_number INTEGER NOT NULL,
+    blob_id TEXT NOT NULL,
+    size INTEGER NOT NULL,
+    etag TEXT NOT NULL,
+    checksum_md5 TEXT NOT NULL,
+    object_key TEXT REFERENCES _mf_objects(key) DEFERRABLE INITIALLY DEFERRED,
+    PRIMARY KEY (upload_id, part_number)
+);
 `;
+// NOTE: `object_key` foreign key constraint is deferred, meaning we can delete
+// the linked object row, *then* the multipart part rows in a transaction,
+// see https://www.sqlite.org/foreignkeys.html#fk_deferred for more details
 
 // https://github.com/cloudflare/workerd/blob/4290f9717bc94647d9c8afd29602cdac97fdff1b/src/workerd/api/r2-api.capnp
 
@@ -154,13 +199,19 @@ export const R2PutRequestSchema = z
     sha512: value.sha512,
   }));
 
-// TODO: support multipart
-export const R2CreateMultipartUploadRequestSchema = z.object({
-  method: z.literal("createMultipartUpload"),
-  object: z.string(),
-  customFields: RecordSchema.optional(),
-  httpFields: R2HttpFieldsSchema.optional(),
-});
+export const R2CreateMultipartUploadRequestSchema = z
+  .object({
+    method: z.literal("createMultipartUpload"),
+    object: z.string(),
+    customFields: RecordSchema.optional(), // (renamed in transform)
+    httpFields: R2HttpFieldsSchema.optional(), // (renamed in transform)
+  })
+  .transform((value) => ({
+    method: value.method,
+    object: value.object,
+    customMetadata: value.customFields,
+    httpMetadata: value.httpFields,
+  }));
 
 export const R2UploadPartRequestSchema = z.object({
   method: z.literal("uploadPart"),
@@ -206,6 +257,7 @@ export const R2DeleteRequestSchema = z.intersection(
 
 // Not using `z.discriminatedUnion()` here, as that doesn't work with
 // intersection/transformed types.
+// TODO(someday): switch to proposed `z.switch()`: https://github.com/colinhacks/zod/issues/2106
 export const R2BindingRequestSchema = z.union([
   R2HeadRequestSchema,
   R2GetRequestSchema,
@@ -222,6 +274,9 @@ export type OmitRequest<T> = Omit<T, "method" | "object">;
 export type R2GetOptions = OmitRequest<z.infer<typeof R2GetRequestSchema>>;
 export type R2PutOptions = OmitRequest<z.infer<typeof R2PutRequestSchema>>;
 export type R2ListOptions = OmitRequest<z.infer<typeof R2ListRequestSchema>>;
+export type R2CreateMultipartUploadOptions = OmitRequest<
+  z.infer<typeof R2CreateMultipartUploadRequestSchema>
+>;
 
 export interface R2ErrorResponse {
   version: number;

packages/tre/src/plugins/r2/schemas.ts

@@ -1,5 +1,5 @@
 import assert from "assert";
-import { R2StringChecksums } from "@cloudflare/workers-types/experimental";
+import type { R2StringChecksums } from "@cloudflare/workers-types/experimental";
 import { InclusiveRange } from "../../storage2";
 import { _parseRanges } from "../shared";
 import {
@@ -102,13 +102,11 @@ export class Validator {
   }
 
   condition(
-    meta?: R2Object | Pick<R2Object, "etag" | "uploaded">,
+    meta?: Pick<R2Object, "etag" | "uploaded">,
     onlyIf?: R2Conditional
   ): Validator {
     if (onlyIf !== undefined && !_testR2Conditional(onlyIf, meta)) {
-      let error = new PreconditionFailed();
-      if (meta instanceof R2Object) error = error.attach(meta);
-      throw error;
+      throw new PreconditionFailed();
     }
     return this;
   }
@@ -121,14 +119,11 @@ export class Validator {
       const ranges = _parseRanges(options.rangeHeader, size);
       // If the header contained a single range, use it. Otherwise, if the
       // header was invalid, or contained multiple ranges, just return the full
-      // response.
-      if (ranges?.length === 1) {
-        const [start, end] = ranges[0];
-        return { start, end };
-      }
+      // response (by returning undefined from this function).
+      if (ranges?.length === 1) return ranges[0];
     } else if (options.range !== undefined) {
       let { offset, length, suffix } = options.range;
-      // Eliminate suffix is specified
+      // Eliminate suffix if specified
       if (suffix !== undefined) {
         if (suffix <= 0) throw new InvalidRange();
         if (suffix > size) suffix = size;
@@ -142,7 +137,7 @@ export class Validator {
       // Clamp length to maximum
       if (offset + length > size) length = size - offset;
       // Convert to inclusive range
-      return { start: offset, end: offset + length - 1 }; // TODO: check -1
+      return { start: offset, end: offset + length - 1 };
     }
   }
 

packages/tre/src/plugins/r2/validator.ts

@@ -4,9 +4,9 @@ import { z } from "zod";
 import { RequestInit, Response } from "../../http";
 import {
   Awaitable,
-  Clock,
   Log,
   MiniflareCoreError,
+  Timers,
   sanitisePath,
 } from "../../shared";
 import {
@@ -34,7 +34,7 @@ export const CloudflareFetchSchema =
 export type CloudflareFetch = z.infer<typeof CloudflareFetchSchema>;
 
 export interface GatewayConstructor<Gateway> {
-  new (log: Log, storage: Storage, clock: Clock): Gateway;
+  new (log: Log, storage: Storage, timers: Timers): Gateway;
 }
 
 export interface RemoteStorageConstructor {
@@ -62,7 +62,7 @@ export class GatewayFactory<Gateway> {
 
   constructor(
     private readonly log: Log,
-    private readonly clock: Clock,
+    private readonly timers: Timers,
     private readonly cloudflareFetch: CloudflareFetch | undefined,
     private readonly pluginName: string,
     private readonly gatewayClass: GatewayConstructor<Gateway>,
@@ -74,7 +74,7 @@ export class GatewayFactory<Gateway> {
     if (storage !== undefined) return storage;
     this.#memoryStorages.set(
       namespace,
-      (storage = new MemoryStorage(undefined, this.clock))
+      (storage = new MemoryStorage(undefined, this.timers.now))
     );
     return storage;
   }
@@ -98,9 +98,13 @@ export class GatewayFactory<Gateway> {
         const root = path.join(fileURLToPath(url), sanitisedNamespace);
         const unsanitise =
           url.searchParams.get(PARAM_FILE_UNSANITISE) === "true";
-        return new FileStorage(root, !unsanitise, this.clock);
+        return new FileStorage(root, !unsanitise, this.timers.now);
       } else if (url.protocol === "sqlite:") {
-        return new SqliteStorage(url.pathname, sanitisedNamespace, this.clock);
+        return new SqliteStorage(
+          url.pathname,
+          sanitisedNamespace,
+          this.timers.now
+        );
       }
       // TODO: support Redis storage?
       if (url.protocol === "remote:") {
@@ -132,15 +136,15 @@ export class GatewayFactory<Gateway> {
       persist === true
         ? path.join(DEFAULT_PERSIST_ROOT, this.pluginName, sanitisedNamespace)
         : path.join(persist, sanitisedNamespace);
-    return new FileStorage(root, undefined, this.clock);
+    return new FileStorage(root, undefined, this.timers.now);
   }
 
   get(namespace: string, persist: Persistence): Gateway {
     const cached = this.#gateways.get(namespace);
     if (cached !== undefined && cached[0] === persist) return cached[1];
 
     const storage = this.getStorage(namespace, persist);
-    const gateway = new this.gatewayClass(this.log, storage, this.clock);
+    const gateway = new this.gatewayClass(this.log, storage, this.timers);
     this.#gateways.set(namespace, [persist, gateway]);
     return gateway;
   }

packages/tre/src/plugins/shared/gateway.ts

@@ -1,9 +1,8 @@
-export * from "./clock";
+export * from "./timers";
 export * from "./data";
-export * from "./deferred";
 export * from "./error";
 export * from "./event";
 export * from "./log";
 export * from "./matcher";
-export * from "./mutex";
+export * from "./sync";
 export * from "./types";

packages/tre/src/shared/clock.ts

@@ -0,0 +1,89 @@
+import assert from "assert";
+import { Awaitable } from "./types";
+
+export type DeferredPromiseResolve<T> = (value: T | PromiseLike<T>) => void;
+export type DeferredPromiseReject = (reason?: any) => void;
+
+export class DeferredPromise<T> extends Promise<T> {
+  readonly resolve: DeferredPromiseResolve<T>;
+  readonly reject: DeferredPromiseReject;
+
+  constructor(
+    executor: (
+      resolve: (value: T | PromiseLike<T>) => void,
+      reject: (reason?: any) => void
+    ) => void = () => {}
+  ) {
+    let promiseResolve: DeferredPromiseResolve<T>;
+    let promiseReject: DeferredPromiseReject;
+    super((resolve, reject) => {
+      promiseResolve = resolve;
+      promiseReject = reject;
+      return executor(resolve, reject);
+    });
+    // Cannot access `this` until after `super`
+    this.resolve = promiseResolve!;
+    this.reject = promiseReject!;
+  }
+}
+
+export class Mutex {
+  private locked = false;
+  private resolveQueue: (() => void)[] = [];
+
+  private lock(): Awaitable<void> {
+    if (!this.locked) {
+      this.locked = true;
+      return;
+    }
+    return new Promise((resolve) => this.resolveQueue.push(resolve));
+  }
+
+  private unlock(): void {
+    assert(this.locked);
+    if (this.resolveQueue.length > 0) {
+      this.resolveQueue.shift()?.();
+    } else {
+      this.locked = false;
+    }
+  }
+
+  get hasWaiting(): boolean {
+    return this.resolveQueue.length > 0;
+  }
+
+  async runWith<T>(closure: () => Awaitable<T>): Promise<T> {
+    const acquireAwaitable = this.lock();
+    if (acquireAwaitable instanceof Promise) await acquireAwaitable;
+    try {
+      const awaitable = closure();
+      if (awaitable instanceof Promise) return await awaitable;
+      return awaitable;
+    } finally {
+      this.unlock();
+    }
+  }
+}
+
+export class WaitGroup {
+  private counter = 0;
+  private resolveQueue: (() => void)[] = [];
+
+  add(): void {
+    this.counter++;
+  }
+
+  done(): void {
+    assert(this.counter > 0);
+    this.counter--;
+    if (this.counter === 0) {
+      let resolve: (() => void) | undefined;
+      while ((resolve = this.resolveQueue.shift()) !== undefined) resolve();
+    }
+  }
+
+  wait(): Promise<void> {
+    if (this.counter === 0) return Promise.resolve();
+    return new Promise((resolve) => this.resolveQueue.push(resolve));
+  }
+}

packages/tre/src/shared/deferred.ts

@@ -0,0 +1,24 @@
+import { Awaitable } from "./types";
+
+export interface Timers {
+  now(): number; // milliseconds
+  queueMicrotask(closure: () => Awaitable<unknown>): void;
+  // TODO(soon): `setTimeout`, for Queues batching
+}
+
+export const defaultTimers: Timers = {
+  now: () => Date.now(),
+  queueMicrotask,
+};
+
+// TODO(soon): remove once we remove the old storage system
+export type Clock = Timers["now"];
+export const defaultClock = defaultTimers.now;
+
+export function millisToSeconds(millis: number): number {
+  return Math.floor(millis / 1000);
+}
+
+export function secondsToMillis(seconds: number): number {
+  return seconds * 1000;
+}

packages/tre/src/shared/index.ts

@@ -4,7 +4,7 @@ import {
   Awaitable,
   base64Decode,
   base64Encode,
-  defaultClock,
+  defaultTimers,
   lexicographicCompare,
   nonCircularClone,
 } from "../shared";
@@ -219,7 +219,7 @@ export abstract class RemoteStorage extends Storage {
     protected readonly cache: Storage,
     protected readonly cloudflareFetch: CloudflareFetch,
     protected readonly namespace: string,
-    protected readonly clock = defaultClock
+    protected readonly timers = defaultTimers
   ) {
     super();
   }

packages/tre/src/shared/mutex.ts

@@ -1,6 +1,11 @@
 import assert from "assert";
 import { ReadableStream } from "stream/web";
-import { Awaitable, base64Decode, base64Encode, defaultClock } from "../shared";
+import {
+  Awaitable,
+  base64Decode,
+  base64Encode,
+  defaultTimers,
+} from "../shared";
 import {
   InclusiveRange,
   MultipartOptions,
@@ -116,15 +121,15 @@ export class KeyValueStorage<Metadata = unknown> {
 
   constructor(
     private readonly storage: NewStorage,
-    private readonly clock = defaultClock
+    private readonly timers = defaultTimers
   ) {
     storage.db.pragma("case_sensitive_like = TRUE");
     storage.db.exec(SQL_SCHEMA);
     this.#stmts = sqlStmts(storage.db);
   }
 
   #hasExpired(entry: Pick<Row, "expiration">) {
-    return entry.expiration !== null && entry.expiration <= this.clock();
+    return entry.expiration !== null && entry.expiration <= this.timers.now();
   }
 
   #backgroundDelete(blobId: string) {
@@ -216,7 +221,7 @@ export class KeyValueStorage<Metadata = unknown> {
 
   async list(opts: KeyEntriesQuery): Promise<KeyEntries<Metadata>> {
     // Find non-expired entries matching query after cursor
-    const now = this.clock();
+    const now = this.timers.now();
     const rows = this.#stmts.list.all({
       now,
       escaped_prefix: escapePrefix(opts.prefix ?? ""),

packages/tre/src/shared/sync.ts

@@ -226,9 +226,9 @@ const expireMacro = test.macro({
   async exec(t, opts: { headers: HeadersInit; expectedTtl: number }) {
     // Reset clock to known time, restoring afterwards.
     // Note this macro must be used with `test.serial` to avoid races.
-    const originalTimestamp = t.context.clock.timestamp;
-    t.teardown(() => (t.context.clock.timestamp = originalTimestamp));
-    t.context.clock.timestamp = 1_000_000; // 1000s
+    const originalTimestamp = t.context.timers.timestamp;
+    t.teardown(() => (t.context.timers.timestamp = originalTimestamp));
+    t.context.timers.timestamp = 1_000_000; // 1000s
 
     const key = "http://localhost/cache-expire";
     await t.context.mf.dispatchFetch(key, {
@@ -240,11 +240,11 @@ const expireMacro = test.macro({
     let res = await t.context.mf.dispatchFetch(key);
     t.is(res.status, 200);
 
-    t.context.clock.timestamp += opts.expectedTtl / 2;
+    t.context.timers.timestamp += opts.expectedTtl / 2;
     res = await t.context.mf.dispatchFetch(key);
     t.is(res.status, 200);
 
-    t.context.clock.timestamp += opts.expectedTtl / 2;
+    t.context.timers.timestamp += opts.expectedTtl / 2;
     res = await t.context.mf.dispatchFetch(key);
     t.is(res.status, 404);
   },
@@ -276,7 +276,7 @@ const isCachedMacro = test.macro({
       .digest("hex");
     const key = `http://localhost/cache-is-cached-${headersHash}`;
 
-    const expires = new Date(t.context.clock.timestamp + 2000).toUTCString();
+    const expires = new Date(t.context.timers.timestamp + 2000).toUTCString();
     await t.context.mf.dispatchFetch(key, {
       method: "PUT",
       headers: {
@@ -406,15 +406,15 @@ test.serial("operations log warning on workers.dev subdomain", async (t) => {
 test.serial("operations persist cached data", async (t) => {
   // Create new temporary file-system persistence directory
   const tmp = await useTmp(t);
-  const clock = () => t.context.clock.timestamp;
+  const clock = () => t.context.timers.timestamp;
   // TODO(soon): clean up this mess once we've migrated all gateways
   const legacyStorage = new FileStorage(
     path.join(tmp, "default"),
     undefined,
     clock
   );
   const newStorage = legacyStorage.getNewStorage();
-  const kvStorage = new KeyValueStorage(newStorage, clock);
+  const kvStorage = new KeyValueStorage(newStorage, t.context.timers);
 
   // Set option, then reset after test
   await t.context.setOptions({ cachePersist: tmp });

packages/tre/src/shared/timers.ts

@@ -16,7 +16,7 @@ import {
   TIME_EXPIRING,
   TIME_NOW,
   createJunkStream,
-  testClock,
+  testTimers,
 } from "../../test-shared";
 
 interface Context {
@@ -28,10 +28,10 @@ const test = anyTest as TestFn<Context>;
 
 test.beforeEach((t) => {
   // TODO(soon): clean up this mess once we've migrated all gateways
-  const legacyStorage = new MemoryStorage(undefined, testClock);
+  const legacyStorage = new MemoryStorage(undefined, testTimers.now);
   const newStorage = legacyStorage.getNewStorage();
-  const gateway = new KVGateway(new NoOpLog(), legacyStorage, testClock);
-  const kvStorage = new KeyValueStorage(newStorage, testClock);
+  const gateway = new KVGateway(new NoOpLog(), legacyStorage, testTimers);
+  const kvStorage = new KeyValueStorage(newStorage, testTimers);
   t.context = { storage: kvStorage, gateway };
 });
 

packages/tre/src/storage/storage.ts

@@ -8,7 +8,7 @@ import {
   Response,
   RouteHandler,
   Router,
-  defaultClock,
+  defaultTimers,
 } from "@miniflare/tre";
 import test from "ava";
 
@@ -21,7 +21,7 @@ class TestRouter extends Router<TestGateway> {
     const log = new NoOpLog();
     super(
       log,
-      new GatewayFactory(log, defaultClock, undefined, "test", TestGateway)
+      new GatewayFactory(log, defaultTimers, undefined, "test", TestGateway)
     );
   }
 

packages/tre/src/storage2/keyvalue.ts

@@ -0,0 +1,104 @@
+import { setTimeout } from "timers/promises";
+import { DeferredPromise, Mutex, WaitGroup } from "@miniflare/tre";
+import test from "ava";
+
+test("DeferredPromise: waits for resolve/reject callbacks", async (t) => {
+  // Check resolves with regular value
+  let promise = new DeferredPromise<number>();
+  promise.resolve(42);
+  t.is(await promise, 42);
+
+  // Check resolves with another Promise
+  promise = new DeferredPromise<number>();
+  promise.resolve(Promise.resolve(0));
+  t.is(await promise, 0);
+
+  // Check rejects with error
+  promise = new DeferredPromise<number>();
+  promise.reject(new Error("🤯"));
+  await t.throwsAsync(promise, { message: "🤯" });
+});
+
+test("Mutex: runs closures exclusively", async (t) => {
+  const mutex = new Mutex();
+  const events: number[] = [];
+  await Promise.all([
+    mutex.runWith(async () => {
+      events.push(1);
+      await setTimeout();
+      events.push(2);
+    }),
+    mutex.runWith(async () => {
+      events.push(3);
+    }),
+  ]);
+  t.deepEqual(events, events[0] === 1 ? [1, 2, 3] : [3, 1, 2]);
+});
+test("Mutex: lock can be acquired synchronously", (t) => {
+  const mutex = new Mutex();
+  let acquired = false;
+  mutex.runWith(() => (acquired = true));
+  t.true(acquired);
+});
+
+test("WaitGroup: waits for all tasks to complete", async (t) => {
+  const group = new WaitGroup();
+
+  // Check doesn't wait if no tasks added
+  await group.wait();
+
+  // Check waits for single task
+  let resolved = false;
+  group.add(); // count -> 1
+  group.wait().then(() => (resolved = true));
+  await Promise.resolve();
+  t.false(resolved);
+
+  group.done(); // count -> 0 (complete)
+  await Promise.resolve();
+  t.true(resolved);
+
+  // Check waits for multiple tasks, including those added whilst waiting
+  resolved = false;
+  group.add(); // count -> 1
+  group.add(); // count -> 2
+  group.wait().then(() => (resolved = true));
+  group.add(); // count -> 3
+  await Promise.resolve();
+  t.false(resolved);
+
+  group.done(); // count -> 2
+  await Promise.resolve();
+  t.false(resolved);
+
+  group.done(); // count -> 1
+  await Promise.resolve();
+  t.false(resolved);
+
+  group.add(); // count -> 2
+  await Promise.resolve();
+  t.false(resolved);
+
+  group.done(); // count -> 1
+  await Promise.resolve();
+  t.false(resolved);
+
+  group.done(); // count -> 0 (complete)
+  await Promise.resolve();
+  t.true(resolved);
+
+  // Check allows multiple waiters
+  resolved = false;
+  let resolved2 = false;
+  group.add(); // count -> 1
+  group.wait().then(() => (resolved = true));
+  group.wait().then(() => (resolved2 = true));
+  await Promise.resolve();
+  t.false(resolved);
+  t.false(resolved2);
+
+  group.done(); // count -> 0 (complete)
+  await Promise.resolve();
+  t.true(resolved);
+  t.true(resolved2);
+});

packages/tre/test/plugins/cache/index.spec.ts

@@ -4,7 +4,7 @@ import type {
   Request as WorkerRequest,
   Response as WorkerResponse,
 } from "@cloudflare/workers-types/experimental";
-import { Awaitable, Miniflare, MiniflareOptions } from "@miniflare/tre";
+import { Awaitable, Miniflare, MiniflareOptions, Timers } from "@miniflare/tre";
 import anyTest, { TestFn } from "ava";
 import { getPort } from "./http";
 import { TestLog } from "./log";
@@ -16,8 +16,25 @@ export type TestMiniflareHandler<Env> = (
   ctx: ExecutionContext
 ) => Awaitable<WorkerResponse>;
 
-export interface TestClock {
-  timestamp: number;
+export class TestTimers implements Timers {
+  timestamp = 1_000_000; // 1000s
+  #pendingMicrotasks = new Set<Promise<unknown>>();
+
+  now = () => {
+    return this.timestamp;
+  };
+
+  queueMicrotask(closure: () => Awaitable<unknown>) {
+    const result = closure();
+    if (result instanceof Promise) {
+      this.#pendingMicrotasks.add(result);
+      result.then(() => this.#pendingMicrotasks.delete(result));
+    }
+  }
+
+  async waitForMicrotasks() {
+    await Promise.all(this.#pendingMicrotasks);
+  }
 }
 
 export interface MiniflareTestContext {
@@ -27,7 +44,7 @@ export interface MiniflareTestContext {
   // Warning: if mutating or calling any of the following, `test.serial` must be
   // used to prevent races.
   log: TestLog;
-  clock: TestClock;
+  timers: TestTimers;
   setOptions(opts: Partial<MiniflareOptions>): Promise<void>;
 }
 
@@ -72,14 +89,13 @@ export function miniflareTest<
   const test = anyTest as TestFn<Context>;
   test.before(async (t) => {
     const log = new TestLog(t);
-    const clock: TestClock = { timestamp: 1_000_000 }; // 1000s
-    const clockFunction = () => clock.timestamp;
+    const timers = new TestTimers();
 
     const opts: Partial<MiniflareOptions> = {
       ...scriptOpts,
       port: await getPort(),
       log,
-      clock: clockFunction,
+      timers,
       verbose: true,
     };
 
@@ -88,7 +104,7 @@ export function miniflareTest<
     // `userOpts`, a `handler` has been provided.
     t.context.mf = new Miniflare({ ...userOpts, ...opts } as MiniflareOptions);
     t.context.log = log;
-    t.context.clock = clock;
+    t.context.timers = timers;
     t.context.setOptions = (userOpts) =>
       t.context.mf.setOptions({ ...userOpts, ...opts } as MiniflareOptions);
     t.context.url = await t.context.mf.ready;

packages/tre/test/plugins/kv/gateway.spec.ts

@@ -3,7 +3,7 @@ import fs from "fs/promises";
 import path from "path";
 import { ReadableStream } from "stream/web";
 import { TextDecoder, TextEncoder } from "util";
-import { Clock, sanitisePath, unwrapBYOBRequest } from "@miniflare/tre";
+import { Timers, sanitisePath, unwrapBYOBRequest } from "@miniflare/tre";
 import { ExecutionContext } from "ava";
 
 const encoder = new TextEncoder();
@@ -27,7 +27,11 @@ export const TIME_NOW = 750;
 // Tests will check the expiry is within 120s of this.
 export const TIME_EXPIRING = 1000;
 
-export const testClock: Clock = () => TIME_NOW * 1000;
+// TODO(soon): remove once we remove the old storage system
+export const testTimers: Timers = {
+  now: () => TIME_NOW * 1000,
+  queueMicrotask,
+};
 
 const tmpRoot = path.resolve(".tmp");
 export async function useTmp(t: ExecutionContext): Promise<string> {