Add unsafe options for configuring Durable Objects (#573)

Wrangler supports multi-worker development using a dev registry. This
system allows you to reference Durable Objects defined in one process
in another. If we had a Durable Object in worker `a` we wanted to
reference in worker `b`, we could with Miniflare 3:

1. Define a proxy Durable Object in the same process as `b` that
   forwarded all requests to `a`'s process, and bind this to `b`
2. In user code, the user could obtain a Durable Object ID and
   stub for the proxy object in `b`
3. The user would send a request to this stub, which would proxy to
   a worker defined in `a`'s process.
4. This worker would have all `a`'s Durable Objects bound, and we
   could obtain a Durable Object ID and stub for the actual object.
5. Finally, we could forward the request to the real stub.

Ideally, the IDs we obtain in steps 2 and 4 would be the same. This
ensures the stub in `b` and instance in `a` share the same ID.
Unfortunately, Durable Object IDs for one object (e.g. the proxy
object in `b`) cannot be used with another (e.g. the actual object in
`a`), _unless_ the objects share the same *unique key*. In this case,
`workerd` effectively treats them as the same object. Normally, this
would be really bad, but these objects are defined in separate
processes, so I think it's ok. The `unsafeUniqueKey` option on
Durable Object bindings allows the unique key to be customised.

An issue with this is that the proxy objects will use this unique key
when persisting their data. The proxy objects don't write any data,
but they still create the databases. This can be confusing, as it'll
look like data for `a` is stored with `a` and `b` (even though `b`'s
copies will be empty). To alleviate this issue, a
`unsafeEphemeralDurableObjects` option has been added which forces
in-memory `workerd` storage for Durable Objects. This gets wiped
between reloads.

Author: mrbbot

packages/tre/src/index.ts

@@ -146,15 +146,34 @@ function getDurableObjectClassNames(
     for (const designator of Object.values(
       workerOpts.do.durableObjects ?? {}
     )) {
-      // Fallback to current worker service if name not defined
-      const [className, serviceName = workerServiceName] =
-        normaliseDurableObject(designator);
+      const {
+        className,
+        // Fallback to current worker service if name not defined
+        serviceName = workerServiceName,
+        unsafeUniqueKey,
+      } = normaliseDurableObject(designator);
+      // Get or create `Map` mapping class name to optional unsafe unique key
       let classNames = serviceClassNames.get(serviceName);
       if (classNames === undefined) {
-        classNames = new Set();
+        classNames = new Map();
         serviceClassNames.set(serviceName, classNames);
       }
-      classNames.add(className);
+      if (classNames.has(className)) {
+        // If we've already seen this class in this service, make sure the
+        // unsafe unique keys match
+        const existingUnsafeUniqueKey = classNames.get(className);
+        if (existingUnsafeUniqueKey !== unsafeUniqueKey) {
+          throw new MiniflareCoreError(
+            "ERR_DIFFERENT_UNIQUE_KEYS",
+            `Multiple unsafe unique keys defined for Durable Object "${className}" in "${serviceName}": ${JSON.stringify(
+              unsafeUniqueKey
+            )} and ${JSON.stringify(existingUnsafeUniqueKey)}`
+          );
+        }
+      } else {
+        // Otherwise, just add it
+        classNames.set(className, unsafeUniqueKey);
+      }
     }
   }
   return serviceClassNames;

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

@@ -9,6 +9,7 @@ import {
   Service,
   Worker_Binding,
   Worker_Module,
+  kVoid,
   supportedCompatibilityDate,
 } from "../../runtime";
 import {
@@ -56,6 +57,8 @@ export const CoreOptionsSchema = z.intersection(
     textBlobBindings: z.record(z.string()).optional(),
     dataBlobBindings: z.record(z.string()).optional(),
     serviceBindings: z.record(ServiceDesignatorSchema).optional(),
+
+    unsafeEphemeralDurableObjects: z.boolean().optional(),
   })
 );
 
@@ -239,9 +242,9 @@ export const CORE_PLUGIN: Plugin<
     }
 
     const name = getUserServiceName(options.name);
-    const classNames = Array.from(
-      durableObjectClassNames.get(name) ?? new Set<string>()
-    );
+    const classNames = durableObjectClassNames.get(name);
+    const classNamesEntries = Array.from(classNames ?? []);
+
     const compatibilityDate = validateCompatibilityDate(
       log,
       options.compatibilityDate ?? FALLBACK_COMPATIBILITY_DATE
@@ -255,16 +258,23 @@ export const CORE_PLUGIN: Plugin<
           compatibilityDate,
           compatibilityFlags: options.compatibilityFlags,
           bindings: workerBindings,
-          durableObjectNamespaces: classNames.map((className) => ({
-            className,
-            // This `uniqueKey` will (among other things) be used as part of the
-            // path when persisting to the file-system. `-` is invalid in
-            // JavaScript class names, but safe on filesystems (incl. Windows).
-            uniqueKey: `${options.name ?? ""}-${className}`,
-          })),
+          durableObjectNamespaces: classNamesEntries.map(
+            ([className, unsafeUniqueKey]) => {
+              return {
+                className,
+                // This `uniqueKey` will (among other things) be used as part of the
+                // path when persisting to the file-system. `-` is invalid in
+                // JavaScript class names, but safe on filesystems (incl. Windows).
+                uniqueKey:
+                  unsafeUniqueKey ?? `${options.name ?? ""}-${className}`,
+              };
+            }
+          ),
           durableObjectStorage:
-            classNames.length === 0
+            classNamesEntries.length === 0
               ? undefined
+              : options.unsafeEphemeralDurableObjects
+              ? { inMemory: kVoid }
               : { localDisk: DURABLE_OBJECTS_STORAGE_SERVICE_NAME },
           cacheApiOutbound: { name: getCacheServiceName(workerIndex) },
         },

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

@@ -21,6 +21,11 @@ export const DurableObjectsOptionsSchema = z.object({
         z.object({
           className: z.string(),
           scriptName: z.string().optional(),
+          // Allow `uniqueKey` to be customised. We use in Wrangler when setting
+          // up stub Durable Objects that proxy requests to Durable Objects in
+          // another `workerd` process, to ensure the IDs created by the stub
+          // object can be used by the real object too.
+          unsafeUniqueKey: z.string().optional(),
         }),
       ])
     )
@@ -34,14 +39,15 @@ export function normaliseDurableObject(
   designator: NonNullable<
     z.infer<typeof DurableObjectsOptionsSchema>["durableObjects"]
   >[string]
-): [className: string, serviceName: string | undefined] {
+): { className: string; serviceName?: string; unsafeUniqueKey?: string } {
   const isObject = typeof designator === "object";
   const className = isObject ? designator.className : designator;
   const serviceName =
     isObject && designator.scriptName !== undefined
       ? getUserServiceName(designator.scriptName)
       : undefined;
-  return [className, serviceName];
+  const unsafeUniqueKey = isObject ? designator.unsafeUniqueKey : undefined;
+  return { className, serviceName, unsafeUniqueKey };
 }
 
 export const DURABLE_OBJECTS_PLUGIN_NAME = "do";
@@ -97,7 +103,7 @@ export const DURABLE_OBJECTS_PLUGIN: Plugin<
   getBindings(options) {
     return Object.entries(options.durableObjects ?? {}).map<Worker_Binding>(
       ([name, klass]) => {
-        const [className, serviceName] = normaliseDurableObject(klass);
+        const { className, serviceName } = normaliseDurableObject(klass);
         return {
           name,
           durableObjectNamespace: { className, serviceName },

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

@@ -5,7 +5,10 @@ import { GatewayConstructor, RemoteStorageConstructor } from "./gateway";
 import { RouterConstructor } from "./router";
 
 // Maps **service** names to the Durable Object class names exported by them
-export type DurableObjectClassNames = Map<string, Set<string>>;
+export type DurableObjectClassNames = Map<
+  string,
+  Map</* className */ string, /* unsafeUniqueKey */ string | undefined>
+>;
 
 export const QueueConsumerOptionsSchema = z.object({
   // https://developers.cloudflare.com/queues/platform/configuration/#consumer

packages/tre/src/shared/error.ts

@@ -24,7 +24,8 @@ export type MiniflareCoreErrorCode =
   | "ERR_PERSIST_REMOTE_UNSUPPORTED" // Remote storage is not supported for this database
   | "ERR_FUTURE_COMPATIBILITY_DATE" // Compatibility date in the future
   | "ERR_NO_WORKERS" // No workers defined
-  | "ERR_DUPLICATE_NAME"; // Multiple workers defined with same name
+  | "ERR_DUPLICATE_NAME" // Multiple workers defined with same name
+  | "ERR_DIFFERENT_UNIQUE_KEYS"; // Multiple Durable Object bindings declared for same class with different unsafe unique keys
 export class MiniflareCoreError extends MiniflareError<MiniflareCoreErrorCode> {}
 
 export class HttpError extends MiniflareError<number> {

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

@@ -60,6 +60,18 @@ test("persists Durable Object data in-memory between options reloads", async (t)
   mf = new Miniflare(opts);
   res = await mf.dispatchFetch("http://localhost");
   t.is(await res.text(), "Options #5: 1");
+
+  // Check doesn't persist with `unsafeEphemeralDurableObjects` enabled
+  opts.script = COUNTER_SCRIPT("Options #6: ");
+  opts.unsafeEphemeralDurableObjects = true;
+  await mf.setOptions(opts);
+  res = await mf.dispatchFetch("http://localhost");
+  t.is(await res.text(), "Options #6: 1");
+  res = await mf.dispatchFetch("http://localhost");
+  t.is(await res.text(), "Options #6: 2");
+  await mf.setOptions(opts);
+  res = await mf.dispatchFetch("http://localhost");
+  t.is(await res.text(), "Options #6: 1");
 });
 
 test("persists Durable Object data on file-system", async (t) => {
@@ -172,3 +184,60 @@ test("multiple Workers access same Durable Object data", async (t) => {
   });
   t.is(await res.text(), "via B: b: 2");
 });
+
+test("can use Durable Object ID from one object in another", async (t) => {
+  const opts: MiniflareOptions = {
+    port: await getPort(),
+    workers: [
+      {
+        name: "a",
+        routes: ["*/id"],
+        unsafeEphemeralDurableObjects: true,
+        durableObjects: {
+          OBJECT_B: { className: "b_B", unsafeUniqueKey: "b-B" },
+        },
+        modules: true,
+        script: `
+        export class b_B {}
+        export default {
+          fetch(request, env) {
+            const id = env.OBJECT_B.newUniqueId();
+            return new Response(id);
+          }
+        }
+        `,
+      },
+      {
+        name: "b",
+        routes: ["*/*"],
+        durableObjects: { OBJECT_B: "B" },
+        modules: true,
+        script: `
+        export class B {
+          constructor(state) {
+            this.state = state;
+          }
+          fetch() {
+            return new Response("id:" + this.state.id);
+          }
+        }
+        export default {
+          fetch(request, env) {
+            const url = new URL(request.url);
+            const id = env.OBJECT_B.idFromString(url.pathname.substring(1));
+            const stub = env.OBJECT_B.get(id);
+            return stub.fetch(request);
+          }
+        }
+        `,
+      },
+    ],
+  };
+  const mf = new Miniflare(opts);
+  t.teardown(() => mf.dispose());
+
+  let res = await mf.dispatchFetch("http://localhost/id");
+  const id = await res.text();
+  res = await mf.dispatchFetch(`http://localhost/${id}`);
+  t.is(await res.text(), `id:${id}`);
+});