fixup! Implement async_hooks.AsyncResource and other cleanups

cloudflare · Dec 12, 2022 · 87a5e19 · 87a5e19
1 parent f1c48ed
commit 87a5e19
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git 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
@@ -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++
@@ -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
@@ -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;