Merge pull request #3114 from cloudflare/dlapid/refactor_limitenforce…

…r_workerobserver

Refactor Observer and LimitEnforcer cleaner for...

src/workerd/api/pyodide/pyodide.c++

@@ -490,7 +490,7 @@ jsg::JsValue SetupEmscripten::getModule(jsg::Lock& js) {
 }
 
 void SetupEmscripten::visitForGc(jsg::GcVisitor& visitor) {
-  visitor.visit(emscriptenRuntime.emscriptenRuntime);
+  visitor.visit(const_cast<EmscriptenRuntime&>(emscriptenRuntime).emscriptenRuntime);
 }
 
 bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader modules) {

src/workerd/api/pyodide/pyodide.h

@@ -411,8 +411,8 @@ class SimplePythonLimiter: public jsg::Object {
 
 class SetupEmscripten: public jsg::Object {
  public:
-  SetupEmscripten(EmscriptenRuntime emscriptenRuntime)
-      : emscriptenRuntime(kj::mv(emscriptenRuntime)) {};
+  SetupEmscripten(const EmscriptenRuntime& emscriptenRuntime)
+      : emscriptenRuntime(emscriptenRuntime) {};
 
   jsg::JsValue getModule(jsg::Lock& js);
 
@@ -421,7 +421,7 @@ class SetupEmscripten: public jsg::Object {
   }
 
  private:
-  EmscriptenRuntime emscriptenRuntime;
+  const EmscriptenRuntime& emscriptenRuntime;
   void visitForGc(jsg::GcVisitor& visitor);
 };
 

src/workerd/io/limit-enforcer.h

@@ -17,7 +17,7 @@ static constexpr size_t DEFAULT_MAX_PBKDF2_ITERATIONS = 100'000;
 // Interface for an object that enforces resource limits on an Isolate level.
 //
 // See also LimitEnforcer, which enforces on a per-request level.
-class IsolateLimitEnforcer {
+class IsolateLimitEnforcer: public kj::Refcounted {
  public:
   // Get CreateParams to pass when constructing a new isolate.
   virtual v8::Isolate::CreateParams getCreateParams() = 0;

src/workerd/io/observer.h

@@ -141,7 +141,9 @@ class RequestObserver: public kj::Refcounted {
   }
 };
 
-class IsolateObserver: public kj::AtomicRefcounted, public jsg::IsolateObserver {
+class JsgIsolateObserver: public kj::AtomicRefcounted, public jsg::IsolateObserver {};
+
+class IsolateObserver: public kj::AtomicRefcounted {
  public:
   virtual ~IsolateObserver() noexcept(false) {}
 

src/workerd/io/worker.c++

@@ -672,6 +672,8 @@ struct Worker::Isolate::Impl {
         actorCacheLru(limitEnforcer.getActorCacheLruOptions()) {
     jsg::runInV8Stack([&](jsg::V8StackScope& stackScope) {
       auto lock = api.lock(stackScope);
+
+      limitEnforcer.customizeIsolate(lock->v8Isolate);
       if (inspectorPolicy != InspectorPolicy::DISALLOW) {
         // We just created our isolate, so we don't need to use Isolate::Impl::Lock.
         KJ_ASSERT(!isMultiTenantProcess(), "inspector is not safe in multi-tenant processes");
@@ -969,18 +971,21 @@ const HeapSnapshotDeleter HeapSnapshotDeleter::INSTANCE;
 }  // namespace
 
 Worker::Isolate::Isolate(kj::Own<Api> apiParam,
+    kj::Own<IsolateObserver> metricsParam,
     kj::StringPtr id,
+    kj::Own<IsolateLimitEnforcer> limitEnforcerParam,
     InspectorPolicy inspectorPolicy,
     ConsoleMode consoleMode)
-    : metrics(kj::atomicAddRef(apiParam->getMetrics())),
+    : metrics(kj::mv(metricsParam)),
       id(kj::str(id)),
+      limitEnforcer(kj::mv(limitEnforcerParam)),
       api(kj::mv(apiParam)),
-      limitEnforcer(api->getLimitEnforcer()),
       consoleMode(consoleMode),
       featureFlagsForFl(makeCompatJson(decompileCompatibilityFlagsForFl(api->getFeatureFlags()))),
-      impl(kj::heap<Impl>(*api, *metrics, limitEnforcer, inspectorPolicy)),
+      impl(kj::heap<Impl>(*api, *metrics, *limitEnforcer, inspectorPolicy)),
       weakIsolateRef(WeakIsolateRef::wrap(this)),
       traceAsyncContextKey(kj::refcounted<jsg::AsyncContextFrame::StorageKey>()) {
+  api->setIsolateObserver(*metrics);
   metrics->created();
   // We just created our isolate, so we don't need to use Isolate::Impl::Lock (nor an async lock).
   jsg::runInV8Stack([&](jsg::V8StackScope& stackScope) {
@@ -1329,7 +1334,7 @@ Worker::Script::Script(kj::Own<const Isolate> isolateParam,
             KJ_SWITCH_ONEOF(source) {
               KJ_CASE_ONEOF(script, ScriptSource) {
                 impl->globals =
-                    script.compileGlobals(lock, isolate->getApi(), isolate->impl->metrics);
+                    script.compileGlobals(lock, isolate->getApi(), isolate->getApi().getObserver());
 
                 {
                   // It's unclear to me if CompileUnboundScript() can get trapped in any infinite loops or
@@ -1387,7 +1392,7 @@ Worker::Isolate::~Isolate() noexcept(false) {
 
   // Update the isolate stats one last time to make sure we're accurate for cleanup in
   // `evicted()`.
-  limitEnforcer.reportMetrics(*metrics);
+  limitEnforcer->reportMetrics(*metrics);
 
   metrics->evicted();
   weakIsolateRef->invalidate();
@@ -3808,7 +3813,7 @@ kj::Own<const Worker::Script> Worker::Isolate::newScript(kj::StringPtr scriptId,
 }
 
 void Worker::Isolate::completedRequest() const {
-  limitEnforcer.completedRequest(id);
+  limitEnforcer->completedRequest(id);
 }
 
 bool Worker::Isolate::isInspectorEnabled() const {

src/workerd/io/worker.h

@@ -274,8 +274,13 @@ class Worker::Isolate: public kj::AtomicRefcounted {
   // Usually it matches the script ID. An exception is preview isolates: there, each preview
   // session has one isolate which may load many iterations of the script (this allows the
   // inspector session to stay open across them).
+  // The Isolate object owns the Api object and outlives it in order to report teardown timing.
+  // The Api object is created before the Isolate object and does not strictly require
+  // request-specific information.
   explicit Isolate(kj::Own<Api> api,
+      kj::Own<IsolateObserver> metrics,
       kj::StringPtr id,
+      kj::Own<IsolateLimitEnforcer> limitEnforcer,
       InspectorPolicy inspectorPolicy,
       ConsoleMode consoleMode = ConsoleMode::INSPECTOR_ONLY);
 
@@ -305,11 +310,11 @@ class Worker::Isolate: public kj::AtomicRefcounted {
       kj::Maybe<ValidationErrorReporter&> errorReporter = kj::none) const;
 
   inline IsolateLimitEnforcer& getLimitEnforcer() {
-    return limitEnforcer;
+    return *limitEnforcer;
   }
 
   inline const IsolateLimitEnforcer& getLimitEnforcer() const {
-    return limitEnforcer;
+    return *limitEnforcer;
   }
 
   inline Api& getApi() {
@@ -408,8 +413,8 @@ class Worker::Isolate: public kj::AtomicRefcounted {
   TeardownFinishedGuard<IsolateObserver&> teardownGuard{*metrics};
 
   kj::String id;
+  kj::Own<IsolateLimitEnforcer> limitEnforcer;
   kj::Own<Api> api;
-  IsolateLimitEnforcer& limitEnforcer;
   ConsoleMode consoleMode;
 
   // If non-null, a serialized JSON object with a single "flags" property, which is a list of
@@ -532,10 +537,8 @@ class Worker::Api {
   virtual jsg::JsObject wrapExecutionContext(
       jsg::Lock& lock, jsg::Ref<api::ExecutionContext> ref) const = 0;
 
-  virtual IsolateLimitEnforcer& getLimitEnforcer() = 0;
-  virtual const IsolateLimitEnforcer& getLimitEnforcer() const = 0;
-  virtual IsolateObserver& getMetrics() = 0;
-  virtual const IsolateObserver& getMetrics() const = 0;
+  virtual const jsg::IsolateObserver& getObserver() const = 0;
+  virtual void setIsolateObserver(IsolateObserver&) = 0;
 
   // Set the module fallback service callback, if any.
   using ModuleFallbackCallback = kj::Maybe<kj::OneOf<kj::String, jsg::ModuleRegistry::ModuleInfo>>(

src/workerd/server/server.c++

@@ -2921,27 +2921,29 @@ kj::Own<Server::Service> Server::makeWorker(kj::StringPtr name,
     }
   };
 
+  auto jsgobserver = kj::atomicRefcounted<JsgIsolateObserver>();
   auto observer = kj::atomicRefcounted<IsolateObserver>();
-  auto limitEnforcer = kj::heap<NullIsolateLimitEnforcer>();
+  auto limitEnforcer = kj::refcounted<NullIsolateLimitEnforcer>();
 
   kj::Maybe<kj::Own<jsg::modules::ModuleRegistry>> newModuleRegistry;
   if (featureFlags.getNewModuleRegistry()) {
     KJ_REQUIRE(experimental,
         "The new ModuleRegistry implementation is an experimental feature. "
         "You must run workerd with `--experimental` to use this feature.");
     newModuleRegistry = WorkerdApi::initializeBundleModuleRegistry(
-        *observer, conf, featureFlags.asReader(), pythonConfig);
+        *jsgobserver, conf, featureFlags.asReader(), pythonConfig);
   }
 
-  auto api =
-      kj::heap<WorkerdApi>(globalContext->v8System, featureFlags.asReader(), kj::mv(limitEnforcer),
-          kj::mv(observer), *memoryCacheProvider, pythonConfig, kj::mv(newModuleRegistry));
+  auto api = kj::heap<WorkerdApi>(globalContext->v8System, featureFlags.asReader(),
+      limitEnforcer->getCreateParams(), kj::mv(jsgobserver), *memoryCacheProvider, pythonConfig,
+      kj::mv(newModuleRegistry));
   auto inspectorPolicy = Worker::Isolate::InspectorPolicy::DISALLOW;
   if (inspectorOverride != kj::none) {
     // For workerd, if the inspector is enabled, it is always fully trusted.
     inspectorPolicy = Worker::Isolate::InspectorPolicy::ALLOW_FULLY_TRUSTED;
   }
-  auto isolate = kj::atomicRefcounted<Worker::Isolate>(kj::mv(api), name, inspectorPolicy,
+  auto isolate = kj::atomicRefcounted<Worker::Isolate>(kj::mv(api), kj::mv(observer), name,
+      kj::mv(limitEnforcer), inspectorPolicy,
       conf.isServiceWorkerScript() ? Worker::ConsoleMode::INSPECTOR_ONLY : consoleMode);
 
   // If we are using the inspector, we need to register the Worker::Isolate