1 files changed, 102 insertions, 15 deletions
diff --git a/src/js/node/async_hooks.ts b/src/js/node/async_hooks.ts
index ef77b79f7..83b313912 100644
--- a/src/js/node/async_hooks.ts
+++ b/src/js/node/async_hooks.ts
@@ -19,17 +19,53 @@
 // use. But the nature of this approach makes the implementation *itself* very low-impact on performance.
 //
 // AsyncContextData is an immutable array managed in here, formatted [key, value, key, value] where
-// each key is an AsyncLocalStorage object and the value is the associated value.
+// each key is an AsyncLocalStorage object and the value is the associated value. There are a ton of
+// calls to $assert which will verify this invariant (only during bun-debug)
 //
 const { cleanupLater, setAsyncHooksEnabled } = $lazy("async_hooks");
 
+// Only run during debug
+function assertValidAsyncContextArray(array: unknown): array is ReadonlyArray<any> | undefined {
+  // undefined is OK
+  if (array === undefined) return true;
+  // Otherwise, it must be an array
+  $assert(
+    Array.isArray(array),
+    "AsyncContextData must be an array or undefined, got",
+    Bun.inspect(array, { depth: 1 }),
+  );
+  // the array has to be even
+  $assert(array.length % 2 === 0, "AsyncContextData should be even-length, got", Bun.inspect(array, { depth: 1 }));
+  // if it is zero-length, use undefined instead
+  $assert(array.length > 0, "AsyncContextData should be undefined if empty, got", Bun.inspect(array, { depth: 1 }));
+  for (var i = 0; i < array.length; i += 2) {
+    $assert(
+      array[i] instanceof AsyncLocalStorage,
+      `Odd indexes in AsyncContextData should be an array of AsyncLocalStorage\nIndex %s was %s`,
+      i,
+      array[i],
+    );
+  }
+  return true;
+}
+
+// Only run during debug
+function debugFormatContextValue(value: ReadonlyArray<any> | undefined) {
+  if (value === undefined) return "{}";
+  let str = "{\n";
+  for (var i = 0; i < value.length; i += 2) {
+    str += `  ${value[i].__id__}: ${Bun.inspect(value[i + 1], { depth: 1, colors: Bun.enableANSIColors })}\n`;
+  }
+}
+
 function get(): ReadonlyArray<any> | undefined {
-  $debug("get", $getInternalField($asyncContext, 0));
+  $debug("get", debugFormatContextValue($getInternalField($asyncContext, 0)));
   return $getInternalField($asyncContext, 0);
 }
 
 function set(contextValue: ReadonlyArray<any> | undefined) {
-  $debug("set", contextValue);
+  $assert(assertValidAsyncContextArray(contextValue));
+  $debug("set", debugFormatContextValue(contextValue));
   return $putInternalField($asyncContext, 0, contextValue);
 }
 
@@ -38,6 +74,14 @@ class AsyncLocalStorage {
 
   constructor() {
     setAsyncHooksEnabled(true);
+
+    // In debug mode assign every AsyncLocalStorage a unique ID
+    if (IS_BUN_DEVELOPMENT) {
+      (this as any).__id__ =
+        Math.random().toString(36).slice(2, 8) +
+        "@" +
+        require("node:path").basename(require("bun:jsc").callerSourceOrigin());
+    }
   }
 
   static bind(fn, ...args: any) {
@@ -67,8 +111,11 @@ class AsyncLocalStorage {
       return;
     }
     var { length } = context;
+    $assert(length > 0);
+    $assert(length % 2 === 0);
     for (var i = 0; i < length; i += 2) {
       if (context[i] === this) {
+        $assert(length > i + 1);
         const clone = context.slice();
         clone[i + 1] = store;
         set(clone);
@@ -76,33 +123,42 @@ class AsyncLocalStorage {
       }
     }
     set(context.concat(this, store));
+    $assert(this.getStore() === store);
   }
 
   exit(cb, ...args) {
     return this.run(undefined, cb, ...args);
   }
 
-  run(store, callback, ...args) {
+  // This function is literred with $asserts to ensure that everything that
+  // is assumed to be true is *actually* true.
+  run(store_value, callback, ...args) {
     var context = get() as any[]; // we make sure to .slice() before mutating
     var hasPrevious = false;
-    var previous;
+    var previous_value;
     var i = 0;
-    var contextWasInit = !context;
-    if (contextWasInit) {
-      set((context = [this, store]));
+    var contextWasAlreadyInit = !context;
+    if (contextWasAlreadyInit) {
+      set((context = [this, store_value]));
     } else {
       // it's safe to mutate context now that it was cloned
       context = context!.slice();
       i = context.indexOf(this);
       if (i > -1) {
+        $assert(i % 2 === 0);
         hasPrevious = true;
-        previous = context[i + 1];
-        context[i + 1] = store;
+        previous_value = context[i + 1];
+        context[i + 1] = store_value;
       } else {
-        context.push(this, store);
+        i = context.length;
+        context.push(this, store_value);
+        $assert(i % 2 === 0);
+        $assert(context.length % 2 === 0);
       }
       set(context);
     }
+    $assert(i > -1, "i was not set");
+    $assert(this.getStore() === store_value, "run: store_value was not set");
     try {
       return callback(...args);
     } catch (e) {
@@ -111,24 +167,36 @@ class AsyncLocalStorage {
       // Note: early `return` will prevent `throw` above from working. I think...
       // Set AsyncContextFrame to undefined if we are out of context values
       if (!this.#disableCalled) {
-        var context2 = get()! as any[];
-        if (context2 === context && contextWasInit) {
+        var context2 = get()! as any[]; // we make sure to .slice() before mutating
+        if (context2 === context && contextWasAlreadyInit) {
+          $assert(context2.length === 2, "context was mutated without copy");
           set(undefined);
         } else {
           context2 = context2.slice(); // array is cloned here
+          $assert(context2[i] === this);
           if (hasPrevious) {
-            context2[i + 1] = previous;
+            context2[i + 1] = previous_value;
             set(context2);
           } else {
+            // i wonder if this is a fair assert to make
             context2.splice(i, 2);
+            $assert(context2.length % 2 === 0);
             set(context2.length ? context2 : undefined);
           }
         }
+        $assert(
+          this.getStore() === previous_value,
+          "run: previous_value",
+          Bun.inspect(previous_value),
+          "was not restored, i see",
+          this.getStore(),
+        );
       }
     }
   }
 
   disable() {
+    $debug("disable " + (this as any).__id__);
     // In this case, we actually do want to mutate the context state
     if (!this.#disableCalled) {
       var context = get() as any[];
@@ -156,11 +224,21 @@ class AsyncLocalStorage {
   }
 }
 
+if (IS_BUN_DEVELOPMENT) {
+  AsyncLocalStorage.prototype[Bun.inspect.custom] = function (depth, options) {
+    if (depth < 0) return `AsyncLocalStorage { ${Bun.inspect((this as any).__id__, options)} }`;
+    return `AsyncLocalStorage { [${options.stylize("debug id", "special")}]: ${Bun.inspect(
+      (this as any).__id__,
+      options,
+    )} }`;
+  };
+}
+
 class AsyncResource {
   type;
   #snapshot;
 
-  constructor(type, options) {
+  constructor(type, options?) {
     if (typeof type !== "string") {
       throw new TypeError('The "type" argument must be of type string. Received type ' + typeof type);
     }
@@ -200,6 +278,15 @@ class AsyncResource {
       set(prev);
     }
   }
+
+  bind(fn, thisArg) {
+    return this.runInAsyncScope.bind(this, fn, thisArg ?? this);
+  }
+
+  static bind(fn, type, thisArg) {
+    type = type || fn.name;
+    return new AsyncResource(type || "bound-anonymous-fn").bind(fn, thisArg);
+  }
 }
 
 // The rest of async_hooks is not implemented and is stubbed with no-ops and warnings.