custom limits for dynamic workers

Author: git-bruh

src/workerd/api/worker-loader.c++

@@ -14,11 +14,12 @@ jsg::Ref<Fetcher> WorkerStub::getEntrypoint(jsg::Lock& js,
     jsg::Optional<kj::Maybe<kj::String>> name,
     jsg::Optional<EntrypointOptions> options) {
   Frankenvalue props;
-
+  kj::Maybe<ResourceLimits> limits;
   KJ_IF_SOME(o, options) {
     KJ_IF_SOME(p, o.props) {
       props = Frankenvalue::fromJs(js, p.getHandle(js));
     }
+    limits = o.limits;
   }
 
   kj::Maybe<kj::String> entrypointName;
@@ -30,19 +31,20 @@ jsg::Ref<Fetcher> WorkerStub::getEntrypoint(jsg::Lock& js,
     }
   }
 
-  auto subreqChannel = channel->getEntrypoint(kj::mv(entrypointName), kj::mv(props));
+  auto subreqChannel = channel->getEntrypoint(kj::mv(entrypointName), kj::mv(props), limits);
   return js.alloc<Fetcher>(IoContext::current().addObject(kj::mv(subreqChannel)));
 }
 
 jsg::Ref<DurableObjectClass> WorkerStub::getDurableObjectClass(jsg::Lock& js,
     jsg::Optional<kj::Maybe<kj::String>> name,
     jsg::Optional<EntrypointOptions> options) {
   Frankenvalue props;
-
+  kj::Maybe<ResourceLimits> limits;
   KJ_IF_SOME(o, options) {
     KJ_IF_SOME(p, o.props) {
       props = Frankenvalue::fromJs(js, p.getHandle(js));
     }
+    limits = o.limits;
   }
 
   kj::Maybe<kj::String> entrypointName;
@@ -55,7 +57,7 @@ jsg::Ref<DurableObjectClass> WorkerStub::getDurableObjectClass(jsg::Lock& js,
   }
 
   return js.alloc<DurableObjectClass>(IoContext::current().addObject(
-      channel->getActorClass(kj::mv(entrypointName), kj::mv(props))));
+      channel->getActorClass(kj::mv(entrypointName), kj::mv(props), limits)));
 }
 
 jsg::Ref<WorkerStub> WorkerLoader::get(
@@ -157,6 +159,7 @@ DynamicWorkerSource WorkerLoader::toDynamicWorkerSource(jsg::Lock& js,
 
   return {.source = kj::mv(extractedSource),
     .compatibilityFlags = compatFlags,
+    .limits = code.limits,
     .env = kj::mv(env),
     .globalOutbound = kj::mv(globalOutbound),
     .tails = kj::mv(tailChannels),

src/workerd/api/worker-loader.h

@@ -21,8 +21,9 @@ class WorkerStub: public jsg::Object {
 
   struct EntrypointOptions {
     jsg::Optional<jsg::JsRef<jsg::JsObject>> props;
+    jsg::Optional<ResourceLimits> limits;
 
-    JSG_STRUCT(props);
+    JSG_STRUCT(props, limits);
   };
 
   jsg::Ref<Fetcher> getEntrypoint(jsg::Lock& js,
@@ -90,6 +91,8 @@ class WorkerLoader: public jsg::Object {
     jsg::Optional<kj::Array<kj::String>> compatibilityFlags;
     jsg::Optional<bool> allowExperimental = false;
 
+    jsg::Optional<ResourceLimits> limits;
+
     kj::String mainModule;
 
     // Modules are specified as an object mapping names to content. If the content is just a
@@ -116,6 +119,7 @@ class WorkerLoader: public jsg::Object {
     JSG_STRUCT(compatibilityDate,
         compatibilityFlags,
         allowExperimental,
+        limits,
         mainModule,
         modules,
         env,
@@ -156,6 +160,6 @@ class WorkerLoader: public jsg::Object {
 
 #define EW_WORKER_LOADER_ISOLATE_TYPES                                                             \
   api::WorkerStub, api::WorkerStub::EntrypointOptions, api::WorkerLoader,                          \
-      api::WorkerLoader::Module, api::WorkerLoader::WorkerCode
+      api::WorkerLoader::Module, api::WorkerLoader::WorkerCode, workerd::ResourceLimits
 
 }  // namespace workerd::api

src/workerd/io/io-channels.h

@@ -297,17 +297,30 @@ class IoChannelFactory {
       ChannelTokenUsage usage, kj::ArrayPtr<const byte> token);
 };
 
+// ResourceLimits provides a means to control the resource allocation for a worker stage via a
+// set of optionally overridden parameters.
+struct ResourceLimits {
+  jsg::Optional<uint32_t> cpuMs;
+  jsg::Optional<uint32_t> subRequests;
+
+  JSG_STRUCT(cpuMs, subRequests);
+
+  ResourceLimits clone() const {
+    return {cpuMs, subRequests};
+  }
+};
+
 // Represents a dynamically-loaded Worker to which requests can be sent.
 //
 // This object is returned before the Worker actually loads, so if any errors occur while loading,
 // any requests sent to the Worker will fail, propagating the exception.
 class WorkerStubChannel {
  public:
   virtual kj::Own<IoChannelFactory::SubrequestChannel> getEntrypoint(
-      kj::Maybe<kj::String> name, Frankenvalue props) = 0;
+      kj::Maybe<kj::String> name, Frankenvalue props, kj::Maybe<ResourceLimits> limits) = 0;
 
   virtual kj::Own<IoChannelFactory::ActorClassChannel> getActorClass(
-      kj::Maybe<kj::String> name, Frankenvalue props) = 0;
+      kj::Maybe<kj::String> name, Frankenvalue props, kj::Maybe<ResourceLimits> limits) = 0;
 
   // TODO(someday): Allow caller to enumerate entrypoints?
 };
@@ -317,6 +330,8 @@ struct DynamicWorkerSource {
   WorkerSource source;
   CompatibilityFlags::Reader compatibilityFlags;
 
+  kj::Maybe<ResourceLimits> limits;
+
   // `env` object to pass to the loaded worker. Can contain anything that can be serialized to
   // a `Frankenvalue` (which should eventually include all binding types, RPC stubs, etc.).
   Frankenvalue env;