From 2b496fcd6fa5f804ddb64de3d2561bddb7ef7b1d Mon Sep 17 00:00:00 2001 From: Josh Howard Date: Sun, 22 Mar 2026 17:15:28 +0000 Subject: [PATCH] Implement per-actor SQLite memory metering. --- src/workerd/io/BUILD.bazel | 1 + src/workerd/io/actor-sqlite.h | 1 + src/workerd/io/limit-enforcer.h | 6 + src/workerd/util/BUILD.bazel | 21 +++ src/workerd/util/sqlite-metering-test.c++ | 160 ++++++++++++++++++++ src/workerd/util/sqlite-metering.c++ | 175 ++++++++++++++++++++++ src/workerd/util/sqlite-metering.h | 60 ++++++++ src/workerd/util/sqlite.h | 15 ++ 8 files changed, 439 insertions(+) create mode 100644 src/workerd/util/sqlite-metering-test.c++ create mode 100644 src/workerd/util/sqlite-metering.c++ create mode 100644 src/workerd/util/sqlite-metering.h diff --git a/src/workerd/io/BUILD.bazel b/src/workerd/io/BUILD.bazel index 669d3e827f7..568f943ac27 100644 --- a/src/workerd/io/BUILD.bazel +++ b/src/workerd/io/BUILD.bazel @@ -206,6 +206,7 @@ wd_cc_library( "//src/workerd/util:autogate", "//src/workerd/util:duration-exceeded-logger", "//src/workerd/util:sqlite", + "//src/workerd/util:sqlite-metering", "@capnp-cpp//src/capnp:capnp-rpc", "@capnp-cpp//src/kj:kj-async", ], diff --git a/src/workerd/io/actor-sqlite.h b/src/workerd/io/actor-sqlite.h index 69adbec6a29..b5d6793da90 100644 --- a/src/workerd/io/actor-sqlite.h +++ b/src/workerd/io/actor-sqlite.h @@ -9,6 +9,7 @@ #include #include #include +#include namespace workerd { diff --git a/src/workerd/io/limit-enforcer.h b/src/workerd/io/limit-enforcer.h index 36f2060952f..d2d58dcc085 100644 --- a/src/workerd/io/limit-enforcer.h +++ b/src/workerd/io/limit-enforcer.h @@ -188,6 +188,12 @@ class LimitEnforcer { // Only used downstream for internal metrics. virtual kj::Duration consumeTimeElapsedForPeriodicLogging() = 0; + + // Returns the last snapshotted SQLite memory usage for the current actor, in bytes. Returns 0 + // for non-actor isolates. + virtual size_t getSqliteMemoryUsage() const { + return 0; + } }; } // namespace workerd diff --git a/src/workerd/util/BUILD.bazel b/src/workerd/util/BUILD.bazel index 470c1d1f6ca..13e732fba8e 100644 --- a/src/workerd/util/BUILD.bazel +++ b/src/workerd/util/BUILD.bazel @@ -191,6 +191,17 @@ wd_cc_library( ], ) +wd_cc_library( + name = "sqlite-metering", + srcs = ["sqlite-metering.c++"], + hdrs = ["sqlite-metering.h"], + visibility = ["//visibility:public"], + deps = [ + "@capnp-cpp//src/kj", + "@sqlite3", + ], +) + wd_cc_library( name = "sqlite", srcs = [ @@ -211,6 +222,7 @@ wd_cc_library( deps = [ ":account-limits", ":sentry", + ":sqlite-metering", "@capnp-cpp//src/kj:kj-async", ], ) @@ -380,6 +392,15 @@ kj_test( ], ) +kj_test( + size = "medium", + src = "sqlite-metering-test.c++", + deps = [ + ":sqlite-metering", + "@capnp-cpp//src/kj", + ], +) + kj_test( src = "sqlite-metadata-test.c++", deps = [ diff --git a/src/workerd/util/sqlite-metering-test.c++ b/src/workerd/util/sqlite-metering-test.c++ new file mode 100644 index 00000000000..c79fcbdc468 --- /dev/null +++ b/src/workerd/util/sqlite-metering-test.c++ @@ -0,0 +1,160 @@ +// Copyright (c) 2025 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +#include "sqlite-metering.h" + +#include +#include + +#include + +namespace workerd { +namespace { + +void test_sqlite3_mem_methods(bool expectEnforcedLimits) { + // sqliteMemMalloc with a negative size returns a nullptr. + KJ_EXPECT(sqliteMemMalloc(-1) == nullptr); + + // sqliteMemMalloc with a zero size returns a nullptr. + KJ_EXPECT(sqliteMemMalloc(0) == nullptr); + + // sqliteMemMalloc with a positive size returns a non-nullptr. + void* ptr = sqliteMemMalloc(1024); + KJ_EXPECT(ptr != nullptr); + + // sqliteMemFree on a non-nullptr does not throw. + sqliteMemFree(ptr); + + // sqliteMemMalloc with a larger size returns a non-nullptr. Note that the size allocated here + // needs to be less than 1 MiB below, or the pointer address will be reused when requesting 1 MiB. + ptr = sqliteMemMalloc(1024 * 512); + if (expectEnforcedLimits) { + KJ_EXPECT(ptr == nullptr); + } else { + KJ_EXPECT(ptr != nullptr); + } + sqliteMemFree(ptr); + + // sqliteMemFree on a nullptr does not throw. + sqliteMemFree(nullptr); + + // sqliteMemRealloc(nullptr, ) should behave as sqliteMemMalloc(). + ptr = sqliteMemRealloc(nullptr, -1); + KJ_EXPECT(ptr == nullptr); + ptr = sqliteMemRealloc(nullptr, 0); + KJ_EXPECT(ptr == nullptr); + ptr = sqliteMemRealloc(nullptr, 1024); + KJ_EXPECT(ptr != nullptr); + sqliteMemFree(ptr); + + // sqliteMemRealloc should return the same pointer if the same size is already allocated by + // sqliteMemMalloc. + ptr = sqliteMemMalloc(1024); + KJ_EXPECT(ptr != nullptr); + void* new_ptr = sqliteMemRealloc(ptr, 1024); + KJ_EXPECT(ptr == new_ptr); + KJ_EXPECT(ptr != nullptr); + + // sqliteMemRealloc should return the same pointer if the same size is already allocated by + // sqliteMemRealloc. + new_ptr = sqliteMemRealloc(ptr, 1024); + KJ_EXPECT(ptr == new_ptr); + KJ_EXPECT(ptr != nullptr); + + // sqliteMemRealloc should return a different pointer if a larger size is requested. + new_ptr = sqliteMemRealloc(ptr, 1024 * 1024); + KJ_EXPECT(ptr != new_ptr); + KJ_EXPECT(ptr != nullptr); + if (expectEnforcedLimits) { + KJ_EXPECT(new_ptr == nullptr); + } else { + KJ_EXPECT(new_ptr != nullptr); + } + sqliteMemFree(new_ptr); + + // sqliteMemRealloc with a valid pointer and negative size returns a nullptr. Note that ptr is + // implicitly freed by the preceding call to sqliteMemRealloc. + ptr = sqliteMemMalloc(1024); + KJ_EXPECT(ptr != nullptr); + new_ptr = sqliteMemRealloc(ptr, -1); + KJ_EXPECT(new_ptr == nullptr); + // sqliteMemRealloc with a valid pointer and zero size returns a nullptr. Note again that ptr is + // implicitly freed by the preceding call to sqliteMemRealloc. + ptr = sqliteMemMalloc(1024); + KJ_EXPECT(ptr != nullptr); + new_ptr = sqliteMemRealloc(ptr, 0); + KJ_EXPECT(new_ptr == nullptr); +} + +KJ_TEST("sqlite3_mem_methods work without SqliteMemoryScope") { + test_sqlite3_mem_methods(/*expectEnforcedLimits*/ false); +} + +KJ_TEST("sqlite3_mem_methods work with SqliteMemoryScope and high limits") { + size_t counter = 0; + SqliteMemoryScope scope(counter, 2 * 1024 * 1024); + test_sqlite3_mem_methods(/*expectEnforcedLimits*/ false); +} + +KJ_TEST("sqlite3_mem_methods work with SqliteMemoryScope and low limits") { + size_t counter = 0; + SqliteMemoryScope scope(counter, 512 * 1024); + test_sqlite3_mem_methods(/*expectEnforcedLimits*/ true); +} + +KJ_TEST("sqlite3_mem_methods correctly modify the SqliteMemoryScope counter") { + auto runThread = [&]() { + size_t counter = 0; + SqliteMemoryScope scope(counter, 1024 * 1024); + + // sqliteMemMalloc followed by sqliteMemFree should result in a zero counter. + void* p1 = sqliteMemMalloc(1024); + KJ_EXPECT(p1 != nullptr); + KJ_EXPECT(counter >= 1024); + sqliteMemFree(p1); + KJ_EXPECT(counter == 0); + + // Multiple sqliteMemMalloc's should increase counter respectively. + void* p2 = sqliteMemMalloc(4096); + KJ_EXPECT(p2 != nullptr); + KJ_EXPECT(counter >= 4096); + void* p3 = sqliteMemMalloc(4096); + KJ_EXPECT(p3 != nullptr); + KJ_EXPECT(counter >= 8192); + size_t counter_snapshot = counter; + + // sqliteMemMalloc leaves counter unchanged when hardLimitBytes is exceeded. + KJ_EXPECT(sqliteMemMalloc(1024 * 1024) == nullptr); + KJ_EXPECT(counter == counter_snapshot); + + // sqliteMemRealloc with an increase in size should increase counter respectively, but less than + // a sqliteMemMalloc of equivalent size. + void* p4 = sqliteMemRealloc(p2, 8192); + KJ_EXPECT(p4 != nullptr); + KJ_EXPECT(counter >= 12288); + KJ_EXPECT(counter <= 16384); + + // sqliteMemRealloc with an decrease in size should decrease counter respectively. + void* p5 = sqliteMemRealloc(p3, 1024); + KJ_EXPECT(p5 != nullptr); + KJ_EXPECT(counter >= 9216); + KJ_EXPECT(counter <= 12288); + counter_snapshot = counter; + + // sqliteMemRealloc leaves counter unchanged when hardLimitBytes is exceeded. + KJ_EXPECT(sqliteMemRealloc(p5, 1024 * 1024) == nullptr); + KJ_EXPECT(counter == counter_snapshot); + + // sqliteMemFree on all active pointers should result in a zero counter. + sqliteMemFree(p4); + sqliteMemFree(p5); + KJ_EXPECT(counter == 0); + }; + + kj::Thread thread1(runThread); + kj::Thread thread2(runThread); +} + +} // namespace +} // namespace workerd diff --git a/src/workerd/util/sqlite-metering.c++ b/src/workerd/util/sqlite-metering.c++ new file mode 100644 index 00000000000..ed56b4df4a3 --- /dev/null +++ b/src/workerd/util/sqlite-metering.c++ @@ -0,0 +1,175 @@ +// Copyright (c) 2025 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +#include "sqlite-metering.h" + +#include + +#include + +#include + +#if defined(__linux__) || defined(__GLIBC__) +#include +#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) malloc_usable_size(p) +#elif defined(__APPLE__) +#include +#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) malloc_size(p) +#elif defined(_WIN32) +#include +#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) _msize(p) +#else +// Fallback: no accounting. This path is unreachable in production (Linux). +#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) 0 +#endif + +namespace workerd { + +thread_local SqliteMemoryScope* SqliteMemoryScope::threadLocalScope = nullptr; + +namespace { + +inline size_t usableSize(void* ptr) { + if (ptr == nullptr) return 0; +#if defined(SQLITE_METERING_MALLOC_USABLE_SIZE) + auto sz = SQLITE_METERING_MALLOC_USABLE_SIZE(ptr); + if (sz > 0) return sz; +#endif + // Fallback: no accounting. This path is unreachable in production (Linux). + return 0; +} + +} // namespace + +void* sqliteMemMalloc(int size) { + if (size <= 0) return nullptr; + void* ptr = std::malloc(static_cast(size)); + if (ptr == nullptr) return nullptr; + SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope; + if (scope != nullptr) { + size_t actual = usableSize(ptr); + if (actual > 0) { + if (scope->counter + actual > scope->hardLimitBytes) { + std::free(ptr); + return nullptr; + } + scope->counter += actual; + } + } + return ptr; +} + +void sqliteMemFree(void* ptr) { + if (ptr == nullptr) return; + SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope; + if (scope != nullptr) { + size_t actual = usableSize(ptr); + if (actual > 0) { + scope->counter -= actual; + } + } + std::free(ptr); +} + +void* sqliteMemRealloc(void* ptr, int newSize) { + // SQLite assumes that sqliteMemRealloc(ptr, newSize) where newSize <= 0 is equivalent to + // sqliteMemFree(ptr). + if (newSize <= 0) { + sqliteMemFree(ptr); + return nullptr; + } + + // SQLite assumes that sqliteMemRealloc(nullptr, newSize) is equivalent to + // sqliteMemMalloc(newSize). + if (ptr == nullptr) { + return sqliteMemMalloc(newSize); + } + + SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope; + size_t oldActual = (scope != nullptr) ? usableSize(ptr) : 0; + + // sqliteMemRealloc must leave the original buffer intact per SQLite's contract when returning + // a nullptr. Because std::realloc consumes the original pointer, we must check whether the + // requested size would exceed the limit before calling std::realloc. + if (scope != nullptr && oldActual > 0 && static_cast(newSize) > oldActual) { + size_t growth = static_cast(newSize) - oldActual; + if (scope->counter + growth > scope->hardLimitBytes) { + return nullptr; + } + } + + void* newPtr = std::realloc(ptr, static_cast(newSize)); + if (newPtr == nullptr) return nullptr; + + if (scope != nullptr && oldActual > 0) { + size_t newActual = usableSize(newPtr); + // Note that `scope->counter + (newActual - oldActual) > scope->hardLimitBytes` could still be + // true. If we attempt to restore a memory block of the original size then the original data + // would be lost. + scope->counter = scope->counter - oldActual + newActual; + } + return newPtr; +} + +namespace { + +int sqliteMemSize(void* ptr) { + return static_cast(usableSize(ptr)); +} + +int sqliteMemRoundup(int n) { + // Round up to the next multiple of 8 as a conservative estimate. + return (n + 7) & ~7; +} + +int sqliteMemInit(void* /*pAppData*/) { + return SQLITE_OK; +} + +void sqliteMemShutdown(void* /*pAppData*/) {} + +static const sqlite3_mem_methods kSqliteMemMethods = { + sqliteMemMalloc, + sqliteMemFree, + sqliteMemRealloc, + sqliteMemSize, + sqliteMemRoundup, + sqliteMemInit, + sqliteMemShutdown, + /*pAppData=*/nullptr, +}; + +} // namespace + +SqliteMemoryScope::SqliteMemoryScope(size_t& counter, size_t hardLimitBytes) + : counter(counter), + hardLimitBytes(hardLimitBytes) { + KJ_ASSERT( + threadLocalScope == nullptr, "Another SqliteMemoryScope is already active on this thread."); + threadLocalScope = this; +} + +SqliteMemoryScope::~SqliteMemoryScope() noexcept(false) { + KJ_ASSERT(threadLocalScope == this); + threadLocalScope = nullptr; +} + +void installSqliteCustomAllocator() { + // sqlite3_config() must be called before sqlite3_initialize(), which is itself invoked + // implicitly by the first sqlite3_vfs_register() or sqlite3_open_v2() call. We rely on callers + // (Vfs constructors, ensureMonkeyPatchedSqlite) to invoke this before any of those. + static bool installed KJ_UNUSED = []() { + int rc = sqlite3_config(SQLITE_CONFIG_MALLOC, &kSqliteMemMethods); + KJ_ASSERT(rc == SQLITE_OK, "sqlite3_config(SQLITE_CONFIG_MALLOC) failed", rc); + + // Disable the global sqlite3_mem_used counter. It carries per-call mutex overhead and is + // redundant now that per-actor accounting is done via the per-actor counter. + rc = sqlite3_config(SQLITE_CONFIG_MEMSTATUS, 0); + KJ_ASSERT(rc == SQLITE_OK, "sqlite3_config(SQLITE_CONFIG_MEMSTATUS) failed", rc); + + return true; + }(); +} + +} // namespace workerd diff --git a/src/workerd/util/sqlite-metering.h b/src/workerd/util/sqlite-metering.h new file mode 100644 index 00000000000..61f9ece9bd8 --- /dev/null +++ b/src/workerd/util/sqlite-metering.h @@ -0,0 +1,60 @@ +// Copyright (c) 2025 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 + +#include + +namespace workerd { + +// This module implements per-actor SQLite memory metering. +// +// SQLite uses a single process-wide memory allocator, but we want to account allocations against +// the specific actor (Durable Object) on whose behalf they are made, and enforce a per-actor hard +// limit for memory allocations. We do this by: +// +// 1. Installing a custom sqlite3_mem_methods that wrap the system allocator. +// +// 2. Installing a thread-local SqliteMemoryScope object, via LimitEnforcerImpl::enterJs, for each +// JS turn of an actor. The scope points at a size_t byte counter owned by ActorSqlite to meter +// memory allocations. +// +// 3. When a SqliteMemoryScope is active on the current thread, we count each memory allocation +// against the byte counter and enforce the per-actor hard limit by returning nullptr signalling +// SQLite to throw a SQLITE_NOMEM exception. +class SqliteMemoryScope { + public: + // counter: the per-actor byte counter owned by SqliteDatabase for its lifetime. + // hardLimitBytes: the per-actor cap from WorkerLimits::sqliteMaxMemoryMb. + explicit SqliteMemoryScope(size_t& counter, size_t hardLimitBytes); + ~SqliteMemoryScope() noexcept(false); + KJ_DISALLOW_COPY_AND_MOVE(SqliteMemoryScope); + + size_t& counter; + const size_t hardLimitBytes; + + private: + // Thread-local pointer to the active scope. Set to this on construction, cleared on destruction. + static thread_local SqliteMemoryScope* threadLocalScope; + + friend void* sqliteMemMalloc(int); + friend void sqliteMemFree(void*); + friend void* sqliteMemRealloc(void*, int); +}; + +// The custom sqlite3_mem_methods functions. Declared here so tests can call them directly to +// verify accounting behaviour without going through the SQLite query engine. +void* sqliteMemMalloc(int size); +void sqliteMemFree(void* ptr); +void* sqliteMemRealloc(void* ptr, int newSize); + +// Install the custom sqlite3_mem_methods. +// +// This must be called before the first sqlite3_initialize(), sqlite3_open_v2(), or +// sqlite3_vfs_register() call in the process. Idempotent (uses a static once-flag). +void installSqliteCustomAllocator(); + +} // namespace workerd diff --git a/src/workerd/util/sqlite.h b/src/workerd/util/sqlite.h index bb2b5a4a4d4..55bee5c0aa3 100644 --- a/src/workerd/util/sqlite.h +++ b/src/workerd/util/sqlite.h @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include @@ -91,6 +92,17 @@ class SqliteDatabase { kj::Maybe maybeMode = kj::none, SqliteObserver& sqliteObserver = SqliteObserver::DEFAULT, kj::Maybe actorAccountLimits = kj::none); + + // Returns a pointer to the per-actor SQLite memory byte counter. + size_t* getSqliteMemoryCounter() { + return &sqliteMemoryBytes; + } + + // Returns the current value of the per-actor SQLite memory byte counter for metrics reporting. + size_t getSqliteMemoryBytes() const { + return sqliteMemoryBytes; + } + ~SqliteDatabase() noexcept(false); KJ_DISALLOW_COPY_AND_MOVE(SqliteDatabase); @@ -317,6 +329,9 @@ class SqliteDatabase { SqliteObserver& sqliteObserver; kj::Maybe actorAccountLimits; + // The amound of memory in bytes used by this actor in bytes for use by sqlite3_mem_methods. + size_t sqliteMemoryBytes = 0; + // This pointer can be left null if a call to reset() failed to re-open the database. kj::Maybe maybeDb;