From 87a5e19f09f2cb4a466a4755a77ecd9ef562cc2a Mon Sep 17 00:00:00 2001
From: James M Snell <jasnell@gmail.com>
Date: Mon, 12 Dec 2022 07:31:05 -0800
Subject: [PATCH] fixup! Implement async_hooks.AsyncResource and other cleanups

---
 src/workerd/jsg/async-context.c++ | 11 +++++++++--
 src/workerd/jsg/async-context.h   |  2 ++
 src/workerd/jsg/setup.c++         |  2 +-
 src/workerd/jsg/setup.h           |  7 ++++++-
 4 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/workerd/jsg/async-context.c++ b/src/workerd/jsg/async-context.c++
index 1cec84feba3..e34cac3ce9a 100644
--- a/src/workerd/jsg/async-context.c++
+++ b/src/workerd/jsg/async-context.c++
@@ -31,6 +31,10 @@ struct AsyncResourceWrappable final: public Wrappable,
     }
     return nullptr;
   }
+
+  kj::Maybe<kj::Own<Wrappable>> maybeGetStrongRef() override {
+    return kj::addRef(*this);
+  }
 };
 
 }  // namespace
@@ -62,7 +66,7 @@ AsyncResource::~AsyncResource() noexcept(false) {
 AsyncResource& AsyncResource::current(Lock& js) {
   auto& isolateBase = IsolateBase::from(js.v8Isolate);
   KJ_ASSERT(!isolateBase.asyncResourceStack.empty());
-  return *isolateBase.asyncResourceStack.front();
+  return *isolateBase.asyncResourceStack.front().resource;
 }
 
 kj::Own<AsyncResource> AsyncResource::create(Lock& js, kj::Maybe<AsyncResource&> maybeParent) {
@@ -179,7 +183,10 @@ AsyncResource::StorageScope::~StorageScope() noexcept(false) {
 }
 
 void IsolateBase::pushAsyncResource(AsyncResource& next) {
-  asyncResourceStack.push_front(&next);
+  asyncResourceStack.push_front(AsyncResourceEntry{
+    &next,
+    next.maybeGetStrongRef()
+  });
 }
 
 void IsolateBase::popAsyncResource() {
diff --git a/src/workerd/jsg/async-context.h b/src/workerd/jsg/async-context.h
index f25ff45e0a2..87ba1bdf07e 100644
--- a/src/workerd/jsg/async-context.h
+++ b/src/workerd/jsg/async-context.h
@@ -58,6 +58,8 @@ class AsyncResource {
 
   ~AsyncResource() noexcept(false);
 
+  virtual kj::Maybe<kj::Own<Wrappable>> maybeGetStrongRef() { return nullptr; }
+
   static AsyncResource& current(Lock& js);
 
   static kj::Own<AsyncResource> create(Lock& js, kj::Maybe<AsyncResource&> maybeParent = nullptr);
diff --git a/src/workerd/jsg/setup.c++ b/src/workerd/jsg/setup.c++
index b50fa443f6e..b34352113c1 100644
--- a/src/workerd/jsg/setup.c++
+++ b/src/workerd/jsg/setup.c++
@@ -325,7 +325,7 @@ IsolateBase::IsolateBase(const V8System& system, v8::Isolate::CreateParams&& cre
       ptr(newIsolate(kj::mv(createParams))),
       heapTracer(ptr),
       rootAsyncResource(*this) {
-  asyncResourceStack.push_front(&rootAsyncResource);
+  asyncResourceStack.push_front({&rootAsyncResource});
 
   ptr->SetFatalErrorHandler(&fatalError);
   ptr->SetOOMErrorHandler(&oomError);
diff --git a/src/workerd/jsg/setup.h b/src/workerd/jsg/setup.h
index e4ff5e6ff8e..641a549e7b3 100644
--- a/src/workerd/jsg/setup.h
+++ b/src/workerd/jsg/setup.h
@@ -238,8 +238,13 @@ class IsolateBase {
   void pushAsyncResource(AsyncResource& next);
   void popAsyncResource();
 
+  struct AsyncResourceEntry {
+    AsyncResource* resource;
+    kj::Maybe<kj::Own<Wrappable>> maybeStrongRef;
+  };
+
   uint64_t nextAsyncResourceId = 0;
-  std::deque<AsyncResource*> asyncResourceStack;
+  std::deque<AsyncResourceEntry> asyncResourceStack;
   kj::HashMap<uint64_t, AsyncResource*> asyncResourcesMap;
   AsyncResource rootAsyncResource;