From 25b6b5f974c06849b3e0b9666d8af64008a1babd Mon Sep 17 00:00:00 2001 From: Dominik Picheta Date: Thu, 30 Jan 2025 19:32:10 +0000 Subject: [PATCH] Fixes "promise will never complete" when exceeding memory. --- src/workerd/io/BUILD.bazel | 1 + src/workerd/io/io-context.h | 31 +++++++++++++++++++++++++----- src/workerd/io/limit-enforcer.h | 2 ++ src/workerd/server/server.c++ | 4 ++++ src/workerd/tests/test-fixture.c++ | 3 +++ src/workerd/util/BUILD.bazel | 7 +++++++ src/workerd/util/exception.h | 14 ++++++++++++++ 7 files changed, 57 insertions(+), 5 deletions(-) create mode 100644 src/workerd/util/exception.h diff --git a/src/workerd/io/BUILD.bazel b/src/workerd/io/BUILD.bazel index 2ef47e77da2..984c548d3ec 100644 --- a/src/workerd/io/BUILD.bazel +++ b/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", diff --git a/src/workerd/io/io-context.h b/src/workerd/io/io-context.h index f15d7c6da7e..34266f2335e 100644 --- a/src/workerd/io/io-context.h +++ b/src/workerd/io/io-context.h @@ -16,6 +16,7 @@ #include #include #include +#include #include #include @@ -1257,23 +1258,42 @@ kj::_::ReducePromises> IoContext::awaitJs(jsg::Lock& js, jsg::Pro auto paf = kj::newPromiseAndFulfiller>(); struct RefcountedFulfiller: public Finalizeable, public kj::Refcounted { kj::Own>> fulfiller; + kj::Own> maybeIsolate; bool isDone = false; - RefcountedFulfiller(kj::Own>> fulfiller) - : fulfiller(kj::mv(fulfiller)) {} + RefcountedFulfiller(kj::Own> maybeIsolate, + kj::Own>> 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 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(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 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> IoContext::awaitJs(jsg::Lock& js, jsg::Pro } } }; - auto fulfiller = kj::refcounted(kj::mv(paf.fulfiller)); + auto& isolate = Worker::Isolate::from(js); + auto fulfiller = kj::refcounted(isolate.getWeakRef(), kj::mv(paf.fulfiller)); auto errorHandler = [fulfiller = addObject(kj::addRef(*fulfiller))]( jsg::Lock& js, jsg::Value jsExceptionRef) mutable { diff --git a/src/workerd/io/limit-enforcer.h b/src/workerd/io/limit-enforcer.h index 4525f889ce2..9572c5b2c15 100644 --- a/src/workerd/io/limit-enforcer.h +++ b/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. diff --git a/src/workerd/server/server.c++ b/src/workerd/server/server.c++ index c33a0041c91..9d67b23d9fa 100644 --- a/src/workerd/server/server.c++ +++ b/src/workerd/server/server.c++ @@ -3199,6 +3199,10 @@ kj::Own 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(); diff --git a/src/workerd/tests/test-fixture.c++ b/src/workerd/tests/test-fixture.c++ index 1925b650385..373e30bc37b 100644 --- a/src/workerd/tests/test-fixture.c++ +++ b/src/workerd/tests/test-fixture.c++ @@ -212,6 +212,9 @@ struct MockIsolateLimitEnforcer final: public IsolateLimitEnforcer { kj::Maybe checkPbkdfIterations(jsg::Lock& lock, size_t iterations) const override { return kj::none; } + bool hasExcessivelyExceededHeapLimit() const override { + return false; + } }; struct MockErrorReporter final: public Worker::ValidationErrorReporter { diff --git a/src/workerd/util/BUILD.bazel b/src/workerd/util/BUILD.bazel index bac0f29ab31..39a4200acf1 100644 --- a/src/workerd/util/BUILD.bazel +++ b/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"]) [ diff --git a/src/workerd/util/exception.h b/src/workerd/util/exception.h new file mode 100644 index 00000000000..4a3240a89d8 --- /dev/null +++ b/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 + +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