Fixes "promise will never complete" when exceeding memory.

Author: dom96

src/workerd/io/BUILD.bazel

@@ -125,6 +125,7 @@ wd_cc_library(
         "//src/workerd/jsg",
         "//src/workerd/util:autogate",
         "//src/workerd/util:completion-membrane",
+        "//src/workerd/util:exception",
         "//src/workerd/util:sqlite",
         "//src/workerd/util:thread-scopes",
         "//src/workerd/util:uuid",

src/workerd/io/io-context.h

@@ -16,6 +16,7 @@
 #include <workerd/io/trace.h>
 #include <workerd/jsg/async-context.h>
 #include <workerd/jsg/jsg.h>
+#include <workerd/util/exception.h>
 #include <workerd/util/uncaught-exception-source.h>
 #include <workerd/util/weak-refs.h>
 
@@ -1257,23 +1258,42 @@ kj::_::ReducePromises<RemoveIoOwn<T>> IoContext::awaitJs(jsg::Lock& js, jsg::Pro
   auto paf = kj::newPromiseAndFulfiller<RemoveIoOwn<T>>();
   struct RefcountedFulfiller: public Finalizeable, public kj::Refcounted {
     kj::Own<kj::PromiseFulfiller<RemoveIoOwn<T>>> fulfiller;
+    kj::Own<const AtomicWeakRef<Worker::Isolate>> maybeIsolate;
     bool isDone = false;
 
-    RefcountedFulfiller(kj::Own<kj::PromiseFulfiller<RemoveIoOwn<T>>> fulfiller)
-        : fulfiller(kj::mv(fulfiller)) {}
+    RefcountedFulfiller(kj::Own<const AtomicWeakRef<Worker::Isolate>> maybeIsolate,
+        kj::Own<kj::PromiseFulfiller<RemoveIoOwn<T>>> fulfiller)
+        : fulfiller(kj::mv(fulfiller)),
+          maybeIsolate(kj::mv(maybeIsolate)) {}
 
     ~RefcountedFulfiller() noexcept(false) {
       if (!isDone) {
+        reject();
+      }
+    }
+
+   private:
+    void reject() {
+      // We use a weak isolate reference here in case the isolate gets dropped before this code
+      // is executed. In that case we default to `false` as we cannot access the original isolate.
+      auto hasExcessivelyExceededHeapLimit = maybeIsolate->tryAddStrongRef()
+                                                 .map([](kj::Own<const Worker::Isolate> isolate) {
+        return isolate->getLimitEnforcer().hasExcessivelyExceededHeapLimit();
+      }).orDefault(false);
+      if (hasExcessivelyExceededHeapLimit) {
+        auto e = JSG_KJ_EXCEPTION(OVERLOADED, Error, "Worker has exceeded memory limit.");
+        e.setDetail(MEMORY_LIMIT_DETAIL_ID, kj::heapArray<kj::byte>(0));
+        fulfiller->reject(kj::mv(e));
+      } else {
         // The JavaScript resolver was garbage collected, i.e. JavaScript will never resolve
         // this promise.
         fulfiller->reject(JSG_KJ_EXCEPTION(FAILED, Error, "Promise will never complete."));
       }
     }
 
-   private:
     kj::Maybe<kj::StringPtr> finalize() override {
       if (!isDone) {
-        fulfiller->reject(JSG_KJ_EXCEPTION(FAILED, Error, "Promise will never complete."));
+        reject();
         isDone = true;
         return "A hanging Promise was canceled. This happens when the worker runtime is waiting "
                "for a Promise from JavaScript to resolve, but has detected that the Promise "
@@ -1284,7 +1304,8 @@ kj::_::ReducePromises<RemoveIoOwn<T>> IoContext::awaitJs(jsg::Lock& js, jsg::Pro
       }
     }
   };
-  auto fulfiller = kj::refcounted<RefcountedFulfiller>(kj::mv(paf.fulfiller));
+  auto& isolate = Worker::Isolate::from(js);
+  auto fulfiller = kj::refcounted<RefcountedFulfiller>(isolate.getWeakRef(), kj::mv(paf.fulfiller));
 
   auto errorHandler = [fulfiller = addObject(kj::addRef(*fulfiller))](
                           jsg::Lock& js, jsg::Value jsExceptionRef) mutable {

src/workerd/io/limit-enforcer.h

@@ -83,6 +83,8 @@ class IsolateLimitEnforcer: public kj::Refcounted {
   virtual size_t getBlobSizeLimit() const {
     return 128 * 1024 * 1024;  // 128 MB
   }
+
+  virtual bool hasExcessivelyExceededHeapLimit() const = 0;
 };
 
 // Abstract interface that enforces resource limits on a IoContext.

src/workerd/server/server.c++

@@ -3199,6 +3199,10 @@ kj::Own<Server::Service> Server::makeWorker(kj::StringPtr name,
       // No limit on the number of iterations in workerd
       return kj::none;
     }
+
+    bool hasExcessivelyExceededHeapLimit() const override {
+      return false;
+    }
   };
 
   auto jsgobserver = kj::atomicRefcounted<JsgIsolateObserver>();

src/workerd/tests/test-fixture.c++

@@ -212,6 +212,9 @@ struct MockIsolateLimitEnforcer final: public IsolateLimitEnforcer {
   kj::Maybe<size_t> checkPbkdfIterations(jsg::Lock& lock, size_t iterations) const override {
     return kj::none;
   }
+  bool hasExcessivelyExceededHeapLimit() const override {
+    return false;
+  }
 };
 
 struct MockErrorReporter final: public Worker::ValidationErrorReporter {

src/workerd/util/BUILD.bazel

@@ -210,6 +210,13 @@ wd_cc_library(
     deps = ["@capnp-cpp//src/capnp"],
 )
 
+wd_cc_library(
+    name = "exception",
+    hdrs = ["exception.h"],
+    visibility = ["//visibility:public"],
+    deps = ["@capnp-cpp//src/kj"],
+)
+
 exports_files(["autogate.h"])
 
 [

src/workerd/util/exception.h

@@ -0,0 +1,14 @@
+// Copyright (c) 2017-2022 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+#pragma once
+
+#include <kj/exception.h>
+
+namespace workerd {
+
+// If an exception is thrown for exceeding memory limits, it will contain this detail.
+constexpr kj::Exception::DetailTypeId MEMORY_LIMIT_DETAIL_ID = 0xbaf76dd7ce5bd8cfull;
+
+}  // namespace workerd