diff --git a/src/workerd/jsg/async-context-test.c++ b/src/workerd/jsg/async-context-test.c++ new file mode 100644 index 00000000000..d1e51bed832 --- /dev/null +++ b/src/workerd/jsg/async-context-test.c++ @@ -0,0 +1,59 @@ +// Copyright (c) 2026 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +#include "async-context.h" +#include "jsg-test.h" + +namespace workerd::jsg::test { +namespace { + +V8System v8System({"--expose-gc"_kj}); + +struct AsyncContextTestContext: public Object, public ContextGlobal { + JSG_RESOURCE_TYPE(AsyncContextTestContext) {} +}; +JSG_DECLARE_ISOLATE_TYPE(AsyncContextTestIsolate, AsyncContextTestContext); + +void collectGarbage(Lock& js) { + auto script = check(v8::Script::Compile(js.v8Context(), js.str("gc(); gc();"_kj))); + check(script->Run(js.v8Context())); +} + +KJ_TEST("AsyncContextFrame::Scope retains its prior frame") { + Evaluator evaluator(v8System); + evaluator.run([&](auto& js) { + auto key = kj::arc(); + kj::Maybe> prior = AsyncContextFrame::create(js, + AsyncContextFrame::StorageEntry( + kj::mv(key), js.v8Ref(v8Str(js.v8Isolate, "prior"_kj).template As()))); + auto weakPrior = KJ_ASSERT_NONNULL(prior).getWeakRef(js); + + auto priorScope = [&]() { + v8::HandleScope handleScope(js.v8Isolate); + return kj::heap(js, *KJ_ASSERT_NONNULL(prior).get()); + }(); + { + auto rootScope = [&]() { + v8::HandleScope handleScope(js.v8Isolate); + return kj::heap(js, kj::none); + }(); + prior = kj::none; + collectGarbage(js); + KJ_EXPECT(weakPrior.isAlive(), "scope did not retain the prior async context frame"); + } + + KJ_IF_SOME(current, AsyncContextFrame::current(js)) { + KJ_IF_SOME(restored, weakPrior.tryGet()) { + KJ_EXPECT(¤t == &restored); + } else { + KJ_FAIL_EXPECT("restored async context frame was destroyed"); + } + } else { + KJ_FAIL_EXPECT("prior async context frame was not restored"); + } + }); +} + +} // namespace +} // namespace workerd::jsg::test diff --git a/src/workerd/jsg/async-context.c++ b/src/workerd/jsg/async-context.c++ index fa2c01643c4..0ae4761fab7 100644 --- a/src/workerd/jsg/async-context.c++ +++ b/src/workerd/jsg/async-context.c++ @@ -57,8 +57,11 @@ kj::Maybe AsyncContextFrame::current(Lock& js) { } kj::Maybe> AsyncContextFrame::currentRef(Lock& js) { - return jsg::AsyncContextFrame::current(js).map( - [](jsg::AsyncContextFrame& frame) { return frame.addRef(); }); + return currentRef(js.v8Isolate); +} + +kj::Maybe> AsyncContextFrame::currentRef(v8::Isolate* isolate) { + return current(isolate).map([](AsyncContextFrame& frame) { return frame.addRef(); }); } kj::Maybe AsyncContextFrame::current(v8::Isolate* isolate) { @@ -158,7 +161,7 @@ AsyncContextFrame::Scope::Scope(Lock& js, kj::Maybe resource AsyncContextFrame::Scope::Scope(v8::Isolate* ptr, kj::Maybe maybeFrame) : isolate(ptr), - prior(AsyncContextFrame::current(ptr)) { + prior(AsyncContextFrame::currentRef(ptr)) { maybeSetV8ContinuationContext(isolate, maybeFrame); } @@ -168,7 +171,8 @@ AsyncContextFrame::Scope::Scope(Lock& js, kj::Maybe>& res })) {} AsyncContextFrame::Scope::~Scope() noexcept(false) { - maybeSetV8ContinuationContext(isolate, prior); + maybeSetV8ContinuationContext(isolate, + prior.map([](Ref& frame) -> AsyncContextFrame& { return *frame; })); } AsyncContextFrame::StorageScope::StorageScope(Lock& js, kj::Arc key, Value store) diff --git a/src/workerd/jsg/async-context.h b/src/workerd/jsg/async-context.h index c2b76c371dc..40205c7eb37 100644 --- a/src/workerd/jsg/async-context.h +++ b/src/workerd/jsg/async-context.h @@ -135,6 +135,7 @@ class AsyncContextFrame final: public Wrappable { // Convenience variation on current() that returns the result wrapped in a Ref for when we // need to make sure the frame stays alive. static kj::Maybe> currentRef(Lock& js); + static kj::Maybe> currentRef(v8::Isolate* isolate); // Create a new AsyncContextFrame. The new frame inherits the storage context of the current // frame (if any) and the given StorageEntry is added. @@ -172,7 +173,7 @@ class AsyncContextFrame final: public Wrappable { // stack until the scope is destroyed. struct Scope { v8::Isolate* isolate; - kj::Maybe prior; + kj::Maybe> prior; // If frame is nullptr, the root frame is assumed. Scope(Lock& js, kj::Maybe frame = kj::none); // If frame is nullptr, the root frame is assumed.