Simplify instrumentation, store env in jsg::Lock

Rather than persisting the env in Worker::Impl, store it in
jsg::Lock so that it is easier to access from places that do
not have access to Worker::Impl.

src/workerd/api/BUILD.bazel

@@ -19,6 +19,7 @@ filegroup(
             "pyodide/pyodide.c++",
             "pyodide/setup-emscripten.c++",
             "memory-cache.c++",
+            "modules.c++",
             "r2*.c++",
             "rtti.c++",
             "url.c++",
@@ -84,7 +85,10 @@ wd_cc_library(
 
 wd_cc_library(
     name = "rtti",
-    srcs = ["rtti.c++"],
+    srcs = [
+        "modules.c++",
+        "rtti.c++",
+    ],
     hdrs = [
         "modules.h",
         "rtti.h",

src/workerd/api/global-scope.c++

@@ -20,7 +20,6 @@
 #include <workerd/io/io-context.h>
 #include <workerd/jsg/async-context.h>
 #include <workerd/jsg/ser.h>
-#include <workerd/jsg/setup.h>
 #include <workerd/jsg/util.h>
 #include <workerd/util/sentry.h>
 #include <workerd/util/stream-utils.h>
@@ -231,7 +230,6 @@ kj::Promise<DeferredProxy<void>> ServiceWorkerGlobalScope::request(kj::HttpMetho
   bool useDefaultHandling;
   KJ_IF_SOME(h, exportedHandler) {
     KJ_IF_SOME(f, h.fetch) {
-      auto envStorageScope = lock.maybeStoreEnvAsyncContext(h.env.addRef(lock));
       auto promise = f(lock, event->getRequest(), h.env.addRef(js), h.getCtx());
       event->respondWith(lock, kj::mv(promise));
       useDefaultHandling = false;
@@ -347,12 +345,10 @@ void ServiceWorkerGlobalScope::sendTraces(kj::ArrayPtr<kj::Own<Trace>> traces,
 
   KJ_IF_SOME(h, exportedHandler) {
     KJ_IF_SOME(f, h.tail) {
-      auto envStorageScope = lock.maybeStoreEnvAsyncContext(h.env.addRef(lock));
       auto tailEvent = jsg::alloc<TailEvent>(lock, "tail"_kj, traces);
       auto promise = f(lock, tailEvent->getEvents(), h.env.addRef(isolate), h.getCtx());
       tailEvent->waitUntil(kj::mv(promise));
     } else KJ_IF_SOME(f, h.trace) {
-      auto envStorageScope = lock.maybeStoreEnvAsyncContext(h.env.addRef(lock));
       auto traceEvent = jsg::alloc<TailEvent>(lock, "trace"_kj, traces);
       auto promise = f(lock, traceEvent->getEvents(), h.env.addRef(isolate), h.getCtx());
       traceEvent->waitUntil(kj::mv(promise));
@@ -387,7 +383,6 @@ void ServiceWorkerGlobalScope::startScheduled(kj::Date scheduledTime,
 
   KJ_IF_SOME(h, exportedHandler) {
     KJ_IF_SOME(f, h.scheduled) {
-      auto envStorageScope = lock.maybeStoreEnvAsyncContext(h.env.addRef(lock));
       auto promise = f(
           lock, jsg::alloc<ScheduledController>(event.addRef()), h.env.addRef(isolate), h.getCtx());
       event->waitUntil(kj::mv(promise));
@@ -434,8 +429,6 @@ kj::Promise<WorkerInterface::AlarmResult> ServiceWorkerGlobalScope::runAlarm(kj:
 
       auto& alarm = KJ_ASSERT_NONNULL(handler.alarm);
 
-      auto envStorageScope = lock.maybeStoreEnvAsyncContext(handler.env.addRef(lock));
-
       return context
           .run([exportedHandler, &context, timeout, retryCount, &alarm,
                    maybeAsyncContext = jsg::AsyncContextFrame::currentRef(lock)](
@@ -580,7 +573,6 @@ jsg::Promise<void> ServiceWorkerGlobalScope::test(
   auto& testHandler =
       JSG_REQUIRE_NONNULL(eh.test, Error, "Entrypoint does not export a test() function.");
 
-  auto envStorageScope = lock.maybeStoreEnvAsyncContext(eh.env.addRef(lock));
   return testHandler(lock, jsg::alloc<TestController>(), eh.env.addRef(lock), eh.getCtx());
 }
 
@@ -619,7 +611,6 @@ void ServiceWorkerGlobalScope::sendHibernatableWebSocketMessage(
 
   KJ_IF_SOME(h, exportedHandler) {
     KJ_IF_SOME(handler, h.webSocketMessage) {
-      auto envStorageScope = lock.maybeStoreEnvAsyncContext(h.env.addRef(lock));
       event->waitUntil(setHibernatableEventTimeout(
           handler(lock, kj::mv(websocket), kj::mv(message)), eventTimeoutMs));
     }
@@ -644,7 +635,6 @@ void ServiceWorkerGlobalScope::sendHibernatableWebSocketClose(HibernatableSocket
       kj::mv(releasePackage.tags), api::WebSocket::HibernatableReleaseState::CLOSE);
   KJ_IF_SOME(h, exportedHandler) {
     KJ_IF_SOME(handler, h.webSocketClose) {
-      auto envStorageScope = lock.maybeStoreEnvAsyncContext(h.env.addRef(lock));
       event->waitUntil(setHibernatableEventTimeout(
           handler(lock, kj::mv(websocket), close.code, kj::mv(close.reason), close.wasClean),
           eventTimeoutMs));
@@ -672,7 +662,6 @@ void ServiceWorkerGlobalScope::sendHibernatableWebSocketError(kj::Exception e,
 
   KJ_IF_SOME(h, exportedHandler) {
     KJ_IF_SOME(handler, h.webSocketError) {
-      auto envStorageScope = lock.maybeStoreEnvAsyncContext(h.env.addRef(lock));
       event->waitUntil(setHibernatableEventTimeout(
           handler(js, kj::mv(websocket), js.exceptionToJs(kj::mv(e))), eventTimeoutMs));
     }

src/workerd/api/modules.c++

@@ -0,0 +1,36 @@
+#include "modules.h"
+
+#include <workerd/jsg/setup.h>
+
+namespace workerd::api {
+
+kj::Maybe<jsg::JsObject> EnvModule::getCurrent(jsg::Lock& js) {
+  auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
+  KJ_IF_SOME(frame, jsg::AsyncContextFrame::current(js)) {
+    KJ_IF_SOME(value, frame.get(key)) {
+      auto handle = value.getHandle(js);
+      if (handle->IsObject()) {
+        return jsg::JsObject(handle.As<v8::Object>());
+      }
+    }
+  }
+  // If the compat flag is set to disable importable env, then this
+  // will return nothing.
+  if (FeatureFlags::get(js).getDisableImportableEnv()) return kj::none;
+
+  // Otherwise, fallback to provide the stored environment.
+  return js.getWorkerEnv().map([&](const jsg::Value& val) -> jsg::JsObject {
+    auto handle = val.getHandle(js);
+    JSG_REQUIRE(handle->IsObject(), TypeError, "Expected environment to be an object.");
+    return jsg::JsObject(handle.As<v8::Object>());
+  });
+}
+
+jsg::JsValue EnvModule::withEnv(
+    jsg::Lock& js, jsg::Value newEnv, jsg::Function<jsg::JsValue()> fn) {
+  auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
+  jsg::AsyncContextFrame::StorageScope storage(js, key, kj::mv(newEnv));
+  return fn(js);
+}
+
+}  // namespace workerd::api

src/workerd/api/modules.h

@@ -13,32 +13,16 @@
 #include <workerd/io/features.h>
 #include <workerd/io/worker.h>
 #include <workerd/jsg/modules-new.h>
-#include <workerd/jsg/setup.h>
 
 #include <cloudflare/cloudflare.capnp.h>
 
 namespace workerd::api {
 
 class EnvModule final: public jsg::Object {
  public:
-  kj::Maybe<jsg::JsObject> getCurrent(jsg::Lock& js) {
-    auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
-    KJ_IF_SOME(frame, jsg::AsyncContextFrame::current(js)) {
-      KJ_IF_SOME(value, frame.get(key)) {
-        auto handle = value.getHandle(js);
-        if (handle->IsObject()) {
-          return jsg::JsObject(handle.As<v8::Object>());
-        }
-      }
-    }
-    return kj::none;
-  }
+  kj::Maybe<jsg::JsObject> getCurrent(jsg::Lock& js);
 
-  jsg::JsValue withEnv(jsg::Lock& js, jsg::Value newEnv, jsg::Function<jsg::JsValue()> fn) {
-    auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
-    jsg::AsyncContextFrame::StorageScope storage(js, key, kj::mv(newEnv));
-    return fn(js);
-  }
+  jsg::JsValue withEnv(jsg::Lock& js, jsg::Value newEnv, jsg::Function<jsg::JsValue()> fn);
 
   JSG_RESOURCE_TYPE(EnvModule) {
     JSG_METHOD(getCurrent);

src/workerd/api/node/util.c++

@@ -242,7 +242,7 @@ jsg::JsValue UtilModule::getBuiltinModule(jsg::Lock& js, kj::String specifier) {
 }
 
 jsg::JsObject UtilModule::getEnvObject(jsg::Lock& js) {
-  return js.getEnv(true);
+  return js.getProcessEnv(true);
 }
 
 namespace {

src/workerd/api/tests/importable-env-test.js

@@ -7,6 +7,9 @@ import {
 } from 'node:assert';
 import { env, withEnv } from 'cloudflare:workers';
 
+import { AsyncLocalStorage } from 'node:async_hooks';
+const check = new AsyncLocalStorage();
+
 // The env is populated at the top level scope.
 strictEqual(env.FOO, 'BAR');
 
@@ -60,5 +63,39 @@ export const importableEnv = {
 
     // Original env is unmodified
     strictEqual(env.BAZ, undefined);
+
+    // Verify that JSRPC calls appropriately see the environment.
+    const remote = await argEnv.RPC.rpcTarget(undefined);
+    await Promise.all([remote.test(), remote.test2()]);
+  },
+
+  get fetch() {
+    // It's a weird edge case, yes, but let's make sure that the env is
+    // available in getters when extracting the default handlers.
+    strictEqual(env.FOO, 'BAR');
+    return async () => {};
+  },
+
+  // Verifies that the environment is available and correctly propagated
+  // in custom events like jsrpc.
+  rpcTarget() {
+    strictEqual(env.FOO, 'BAR');
+    return withEnv({ FOO: 'BAZ' }, () => {
+      // Arguably, returned RPC targets should probably automatically capture
+      // the current async context and propagate that on subsequent calls,
+      // but they currently do not... therefore we have to manually capture
+      // the async context and be sure to enter it ourselves below.
+      const runInAsyncContext = AsyncLocalStorage.snapshot();
+      return {
+        test() {
+          // This one runs with the modified env
+          runInAsyncContext(() => strictEqual(env.FOO, 'BAZ'));
+        },
+        test2() {
+          // This one runs with the original env
+          strictEqual(env.FOO, 'BAR');
+        },
+      };
+    });
   },
 };

src/workerd/api/tests/importable-env-test.wd-test

@@ -14,6 +14,7 @@ const unitTests :Workerd.Config = (
           "nodejs_compat_v2",
         ],
         bindings = [
+          (name = "RPC", service = ( name = "importable-env-test", entrypoint="importableEnv")),
           (name = "FOO", text = "BAR"),
           (name = "CACHE", memoryCache = (
             id = "abc123",

src/workerd/io/worker.c++

@@ -481,7 +481,6 @@ struct Worker::Impl {
   kj::Maybe<jsg::JsContext<api::ServiceWorkerGlobalScope>> context;
 
   // The environment blob to pass to handlers.
-  kj::Maybe<jsg::Value> env;
   kj::Maybe<jsg::Value> ctxExports;
 
   // Note: The default export is given the string name "default", because that's what V8 tells us,
@@ -1629,6 +1628,7 @@ Worker::Worker(kj::Own<const Script> scriptParam,
             if (script->isModular()) {
               // Use `env` variable.
               bindingsScope = v8::Object::New(lock.v8Isolate);
+              lock.setWorkerEnv(lock.v8Ref(bindingsScope.As<v8::Value>()));
             } else {
               // Use global-scope bindings.
               bindingsScope = context->Global();
@@ -1649,14 +1649,6 @@ Worker::Worker(kj::Own<const Script> scriptParam,
             SpanBuilder currentUserSpan =
                 spans.userParentSpan.newChild("lw:top_level_execution"_kjc);
 
-            auto envStorageScope =
-                ([&]() -> kj::Maybe<kj::Own<jsg::AsyncContextFrame::StorageScope>> {
-              if (FeatureFlags::get(js).getDisableImportableEnv()) return kj::none;
-              return kj::heap<jsg::AsyncContextFrame::StorageScope>(js,
-                  jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey(),
-                  js.v8Ref(bindingsScope.As<v8::Value>()));
-            })();
-
             KJ_SWITCH_ONEOF(script->impl->unboundScriptOrMainModule) {
               KJ_CASE_ONEOF(unboundScript, jsg::NonModuleScript) {
                 auto limitScope =
@@ -1666,7 +1658,6 @@ Worker::Worker(kj::Own<const Script> scriptParam,
               KJ_CASE_ONEOF(mainModule, kj::Path) {
                 KJ_IF_SOME(ns,
                     tryResolveMainModule(lock, mainModule, *jsContext, *script, limitErrorOrTime)) {
-                  impl->env = lock.v8Ref(bindingsScope.As<v8::Value>());
                   impl->ctxExports = lock.v8Ref(ctxExports.As<v8::Value>());
 
                   auto& api = script->isolate->getApi();
@@ -1995,7 +1986,7 @@ kj::Maybe<kj::Own<api::ExportedHandler>> Worker::Lock::getExportedHandler(
     auto handler = kj::heap(cls(js,
         jsg::alloc<api::ExecutionContext>(js,
             jsg::JsValue(KJ_ASSERT_NONNULL(worker.impl->ctxExports).getHandle(js)), props.toJs(js)),
-        KJ_ASSERT_NONNULL(worker.impl->env).addRef(js)));
+        KJ_ASSERT_NONNULL(js.getWorkerEnv())));
 
     // HACK: We set handler.env and handler.ctx to undefined because we already passed the real
     //   env and ctx into the constructor, and we want the handler methods to act like they take
@@ -2040,13 +2031,6 @@ jsg::AsyncContextFrame::StorageKey& Worker::Lock::getEnvAsyncContextKey() {
   return isolate.getEnvAsyncContextKey();
 }
 
-kj::Maybe<kj::Own<jsg::AsyncContextFrame::StorageScope>> Worker::Lock::maybeStoreEnvAsyncContext(
-    jsg::Value env) {
-  if (FeatureFlags::get(*this).getDisableImportableEnv()) return kj::none;
-  return kj::heap<jsg::AsyncContextFrame::StorageScope>(
-      *this, getEnvAsyncContextKey(), kj::mv(env));
-}
-
 bool Worker::Lock::isInspectorEnabled() {
   return worker.script->isolate->impl->inspector != kj::none;
 }
@@ -3502,7 +3486,7 @@ kj::Promise<void> Worker::Actor::ensureConstructedImpl(IoContext& context, Actor
             jsg::JsRef<jsg::JsValue>(
                 js, KJ_ASSERT_NONNULL(lock.getWorker().impl->ctxExports).addRef(js)),
             kj::mv(storage), kj::mv(impl->container), containerRunning),
-        KJ_ASSERT_NONNULL(lock.getWorker().impl->env).addRef(js));
+        KJ_ASSERT_NONNULL(js.getWorkerEnv()));
 
     // HACK: We set handler.env to undefined because we already passed the real env into the
     //   constructor, and we want the handler methods to act like they take just one parameter.

src/workerd/io/worker.h

@@ -656,8 +656,6 @@ class Worker::Lock {
   jsg::AsyncContextFrame::StorageKey& getTraceAsyncContextKey();
 
   jsg::AsyncContextFrame::StorageKey& getEnvAsyncContextKey();
-  kj::Maybe<kj::Own<jsg::AsyncContextFrame::StorageScope>> maybeStoreEnvAsyncContext(
-      jsg::Value env);
 
  private:
   explicit Lock(const Worker& worker, LockType lockType, jsg::V8StackScope&);

src/workerd/jsg/jsg.h

@@ -2682,10 +2682,16 @@ class Lock {
 
   // Sets an env value that will be expressed on the process.env
   // if/when nodejs-compat mode is used.
-  virtual void setEnvField(const JsValue& name, const JsValue& value) = 0;
+  virtual void setProcessEnvField(const JsValue& name, const JsValue& value) = 0;
 
-  // Returns the env base object.
-  virtual JsObject getEnv(bool release = false) = 0;
+  // Returns the process.env base object.
+  virtual JsObject getProcessEnv(bool release = false) = 0;
+
+  // Store the worker environment.
+  virtual void setWorkerEnv(Value value) = 0;
+
+  // Retrieve the worker environment.
+  virtual kj::Maybe<Value> getWorkerEnv() = 0;
 
  private:
   // Mark the jsg::Lock as being disallowed from being passed as a parameter into

src/workerd/jsg/setup.c++

@@ -434,6 +434,7 @@ void IsolateBase::dropWrappers(kj::FunctionParam<void()> drop) {
     KJ_DEFER(symbolAsyncDispose.Reset());
     KJ_DEFER(opaqueTemplate.Reset());
     KJ_DEFER(envObj.Reset());
+    KJ_DEFER(workerEnvObj.Reset());
 
     // Make sure the TypeWrapper is destroyed under lock by declaring a new copy of the variable
     // that is destroyed before the lock is released.

src/workerd/jsg/setup.h

@@ -278,6 +278,9 @@ class IsolateBase {
   // Object that is used as the underlying target of process.env when nodejs-compat mode is used.
   v8::Global<v8::Object> envObj;
 
+  // Object used as the underlying storage for a workers environment.
+  v8::Global<v8::Object> workerEnvObj;
+
   // Polyfilled Symbol.asyncDispose.
   v8::Global<v8::Symbol> symbolAsyncDispose;
 
@@ -690,12 +693,12 @@ class Isolate: public IsolateBase {
 
     // Sets an env value that will be expressed on the process.env
     // if/when nodejs-compat mode is used.
-    void setEnvField(const JsValue& name, const JsValue& value) override {
-      getEnv().set(*this, name, value);
+    void setProcessEnvField(const JsValue& name, const JsValue& value) override {
+      getProcessEnv().set(*this, name, value);
     }
 
     // Returns the env base object.
-    JsObject getEnv(bool release = false) override {
+    JsObject getProcessEnv(bool release = false) override {
       KJ_DEFER({
         if (release) jsgIsolate.envObj.Reset();
       });
@@ -706,6 +709,17 @@ class Isolate: public IsolateBase {
       return JsObject(jsgIsolate.envObj.Get(v8Isolate));
     }
 
+    void setWorkerEnv(Value value) override {
+      auto handle = value.getHandle(*this);
+      KJ_ASSERT(handle->IsObject());
+      jsgIsolate.workerEnvObj.Reset(v8Isolate, handle.template As<v8::Object>());
+    }
+
+    kj::Maybe<Value> getWorkerEnv() override {
+      if (jsgIsolate.workerEnvObj.IsEmpty()) return kj::none;
+      return v8Ref<v8::Value>(jsgIsolate.workerEnvObj.Get(v8Isolate));
+    }
+
    private:
     Isolate& jsgIsolate;