diff --git a/src/workerd/io/limit-enforcer.h b/src/workerd/io/limit-enforcer.h index 36f2060952f..0598368d955 100644 --- a/src/workerd/io/limit-enforcer.h +++ b/src/workerd/io/limit-enforcer.h @@ -188,6 +188,9 @@ class LimitEnforcer { // Only used downstream for internal metrics. virtual kj::Duration consumeTimeElapsedForPeriodicLogging() = 0; + + // Only used for internal metrics. + virtual size_t getSqliteMemoryUsage() const = 0; }; } // namespace workerd diff --git a/src/workerd/server/server.c++ b/src/workerd/server/server.c++ index c1c9581b566..fb25b4a6937 100644 --- a/src/workerd/server/server.c++ +++ b/src/workerd/server/server.c++ @@ -3568,6 +3568,9 @@ class Server::WorkerService final: public Service, kj::Duration consumeTimeElapsedForPeriodicLogging() override { return 0 * kj::SECONDS; } + size_t getSqliteMemoryUsage() const override { + return 0; + } }; struct FutureSubrequestChannel { diff --git a/src/workerd/tests/test-fixture.c++ b/src/workerd/tests/test-fixture.c++ index 91c3110c3cc..9e0ac35f21d 100644 --- a/src/workerd/tests/test-fixture.c++ +++ b/src/workerd/tests/test-fixture.c++ @@ -207,6 +207,9 @@ struct MockLimitEnforcer final: public LimitEnforcer { kj::Duration consumeTimeElapsedForPeriodicLogging() override { return 0 * kj::SECONDS; } + size_t getSqliteMemoryUsage() const override { + return 0; + } }; struct MockIsolateLimitEnforcer final: public IsolateLimitEnforcer { diff --git a/src/workerd/util/BUILD.bazel b/src/workerd/util/BUILD.bazel index 453aa286512..9f8ded578bc 100644 --- a/src/workerd/util/BUILD.bazel +++ b/src/workerd/util/BUILD.bazel @@ -212,10 +212,24 @@ wd_cc_library( deps = [ ":account-limits", ":sentry", + ":sqlite-metering", "@capnp-cpp//src/kj:kj-async", ], ) +wd_cc_library( + name = "sqlite-metering", + srcs = ["sqlite-metering.c++"], + hdrs = ["sqlite-metering.h"], + implementation_deps = [ + "@sqlite3", + ], + visibility = ["//visibility:public"], + deps = [ + "@capnp-cpp//src/kj", + ], +) + wd_cc_library( name = "test-util", testonly = True, @@ -388,6 +402,14 @@ kj_test( ], ) +kj_test( + src = "sqlite-metering-test.c++", + deps = [ + ":sqlite-metering", + "@capnp-cpp//src/kj", + ], +) + kj_test( src = "test-test.c++", deps = [ diff --git a/src/workerd/util/sqlite-kv-test.c++ b/src/workerd/util/sqlite-kv-test.c++ index fee36a92820..820e8221867 100644 --- a/src/workerd/util/sqlite-kv-test.c++ +++ b/src/workerd/util/sqlite-kv-test.c++ @@ -9,6 +9,14 @@ namespace workerd { namespace { +struct GlobalInit { + GlobalInit() { + installSqliteCustomAllocator(); + } +}; + +static GlobalInit init; + KJ_TEST("SQLite-KV") { class TestSqliteObserver: public SqliteObserver { public: @@ -24,8 +32,8 @@ KJ_TEST("SQLite-KV") { auto dir = kj::newInMemoryDirectory(kj::nullClock()); SqliteDatabase::Vfs vfs(*dir); TestSqliteObserver sqliteObserver; - SqliteDatabase db( - vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, sqliteObserver); + SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, + /*sqliteMaxMemoryBytes=*/kj::maxValue, sqliteObserver); SqliteKv kv(db); kv.put("foo", "abc"_kj.asBytes()); @@ -168,8 +176,8 @@ KJ_TEST("SQLite-KV multi-put") { auto dir = kj::newInMemoryDirectory(kj::nullClock()); SqliteDatabase::Vfs vfs(*dir); TestSqliteObserver sqliteObserver; - SqliteDatabase db( - vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, sqliteObserver); + SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, + /*sqliteMaxMemoryBytes=*/kj::maxValue, sqliteObserver); SqliteKv kv(db); // Test basic multi-put with a simple struct diff --git a/src/workerd/util/sqlite-metadata-test.c++ b/src/workerd/util/sqlite-metadata-test.c++ index 911672b8742..b153a44e781 100644 --- a/src/workerd/util/sqlite-metadata-test.c++ +++ b/src/workerd/util/sqlite-metadata-test.c++ @@ -9,6 +9,14 @@ namespace workerd { namespace { +struct GlobalInit { + GlobalInit() { + installSqliteCustomAllocator(); + } +}; + +static GlobalInit init; + KJ_TEST("SQLite-METADATA") { auto dir = kj::newInMemoryDirectory(kj::nullClock()); SqliteDatabase::Vfs vfs(*dir); diff --git a/src/workerd/util/sqlite-metering-test.c++ b/src/workerd/util/sqlite-metering-test.c++ new file mode 100644 index 00000000000..b414ed1b26a --- /dev/null +++ b/src/workerd/util/sqlite-metering-test.c++ @@ -0,0 +1,193 @@ +// Copyright (c) 2026 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. 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 + 1); + 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 a different pointer if a larger size is requested. + ptr = sqliteMemMalloc(1024); + void* new_ptr = sqliteMemRealloc(ptr, 1024 * 1024); + KJ_EXPECT(ptr != new_ptr); + KJ_EXPECT(ptr != nullptr); + if (expectEnforcedLimits) { + // Realloc was rejected: new_ptr is nullptr and ptr still owns the original allocation. + KJ_EXPECT(new_ptr == nullptr); + sqliteMemFree(ptr); + } 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") { + // We expect a bunch of error logs, but we do not crash. + KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemRealloc did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope."); + KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope."); + test_sqlite3_mem_methods(/*expectEnforcedLimits*/ false); +} + +KJ_TEST("sqlite3_mem_methods work with SqliteMemoryScope and high limits") { + size_t memoryBytes = 0; + SqliteMemoryScope scope(memoryBytes, 2 * 1024 * 1024); + test_sqlite3_mem_methods(/*expectEnforcedLimits*/ false); +} + +KJ_TEST("sqlite3_mem_methods work with SqliteMemoryScope and low limits") { + size_t memoryBytes = 0; + SqliteMemoryScope scope(memoryBytes, 512 * 1024); + test_sqlite3_mem_methods(/*expectEnforcedLimits*/ true); +} + +KJ_TEST("sqlite3_mem_methods correctly modify the SqliteMemoryScope memoryBytes") { + auto runThread = [&]() { + size_t memoryBytes = 0; + SqliteMemoryScope scope(memoryBytes, 1024 * 1024); + + // sqliteMemMalloc followed by sqliteMemFree should result in a zero memoryBytes. + void* p1 = sqliteMemMalloc(1024); + KJ_EXPECT(p1 != nullptr); + KJ_EXPECT(memoryBytes >= 1024); + sqliteMemFree(p1); + KJ_EXPECT(memoryBytes == 0); + + // Multiple sqliteMemMalloc's should increase memoryBytes respectively. + void* p2 = sqliteMemMalloc(4096); + KJ_EXPECT(p2 != nullptr); + KJ_EXPECT(memoryBytes >= 4096); + void* p3 = sqliteMemMalloc(4096); + KJ_EXPECT(p3 != nullptr); + KJ_EXPECT(memoryBytes >= 8192); + size_t memoryBytesSnapshot = memoryBytes; + + // sqliteMemMalloc leaves memoryBytes unchanged when hardLimitBytes is exceeded. + KJ_EXPECT(sqliteMemMalloc(1024 * 1024) == nullptr); + KJ_EXPECT(memoryBytes == memoryBytesSnapshot); + + // sqliteMemRealloc with an increase in size should increase memoryBytes respectively, but less than + // a sqliteMemMalloc of equivalent size. + void* p4 = sqliteMemRealloc(p2, 8192); + KJ_EXPECT(p4 != nullptr); + KJ_EXPECT(memoryBytes >= 12288); + KJ_EXPECT(memoryBytes <= 16384); + + // sqliteMemRealloc with an decrease in size should decrease memoryBytes respectively. + void* p5 = sqliteMemRealloc(p3, 1024); + KJ_EXPECT(p5 != nullptr); + KJ_EXPECT(memoryBytes >= 9216); + KJ_EXPECT(memoryBytes <= 12288); + memoryBytesSnapshot = memoryBytes; + + // sqliteMemRealloc leaves memoryBytes unchanged when hardLimitBytes is exceeded. + KJ_EXPECT(sqliteMemRealloc(p5, 1024 * 1024) == nullptr); + KJ_EXPECT(memoryBytes == memoryBytesSnapshot); + + // sqliteMemFree on all active pointers should result in a zero memoryBytes. + sqliteMemFree(p4); + sqliteMemFree(p5); + KJ_EXPECT(memoryBytes == 0); + }; + + kj::Thread thread1(runThread); + kj::Thread thread2(runThread); +} + +KJ_TEST("SqliteMemoryScope is idempotent and handles nested scopes") { + size_t memoryBytes = 0; + SqliteMemoryScope outerScope(memoryBytes, 1024 * 1024); + + void* p1 = sqliteMemMalloc(1024); + KJ_EXPECT(p1 != nullptr); + size_t memoryBytesSnapshot = memoryBytes; + KJ_EXPECT(memoryBytesSnapshot >= 1024); + + { + size_t innerMemoryBytes = 0; + SqliteMemoryScope innerScope(innerMemoryBytes, 512); + + void* p2 = sqliteMemMalloc(2048); + KJ_EXPECT(p2 != nullptr); + KJ_EXPECT(memoryBytes > memoryBytesSnapshot); + KJ_EXPECT(innerMemoryBytes == 0); + + sqliteMemFree(p2); + } + + void* p3 = sqliteMemMalloc(512); + KJ_EXPECT(p3 != nullptr); + KJ_EXPECT(memoryBytes >= memoryBytesSnapshot); + + sqliteMemFree(p1); + sqliteMemFree(p3); + KJ_EXPECT(memoryBytes == 0); +} + +} // 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..dd14a796bed --- /dev/null +++ b/src/workerd/util/sqlite-metering.c++ @@ -0,0 +1,216 @@ +// Copyright (c) 2026 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) +#elifdef __APPLE__ +#include +#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) malloc_size(p) +#elifdef _WIN32 +#include +#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) _msize(p) +#else +#error "Platform not supported" +#endif + +namespace { + +inline size_t usableSize(void* ptr) { + if (ptr == nullptr) return 0; + return SQLITE_METERING_MALLOC_USABLE_SIZE(ptr); +} + +} // namespace + +namespace workerd { + +thread_local SqliteMemoryScope* SqliteMemoryScope::threadLocalScope = nullptr; + +void* sqliteMemMalloc(int size) { + if (size <= 0) return nullptr; + + SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope; + if (scope != nullptr) { + if (scope->memoryBytes + static_cast(size) > scope->maxMemoryBytes) { + return nullptr; + } + } else { + // This branch is a defensive measure, if we have custom sqlite3_mem_methods installed, and we + // must to hit this branch, then scope should not be a nullptr. If we see this, then something + // is wrong with our memory metering. + KJ_LOG(ERROR, "sqliteMemMalloc did not find a valid scope."); + } + + void* ptr = malloc(static_cast(size)); + if (ptr == nullptr) return nullptr; + if (scope != nullptr) { + size_t actual = usableSize(ptr); + if (actual > 0) { + // We have already checked that allocating the requested size will not exceed maxMemoryBytes, + // but we still have to check that the actual size allocated does not exceed maxMemoryBytes. + if (scope->memoryBytes + actual > scope->maxMemoryBytes) { + free(ptr); + return nullptr; + } + scope->memoryBytes += 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) { + if (scope->memoryBytes >= actual) { + scope->memoryBytes -= actual; + } else { + // This branch is a defensive measure, but it is important since scope->memoryBytes is a size_t + // and an underflow would result in all subsequent allocations failing with SQLITE_NOMEM. + // If we see this, then something is wrong with our memory metering. + KJ_LOG(ERROR, "sqliteMemFree would have triggered a memoryBytes underflow."); + scope->memoryBytes = 0; + } + } + } else { + // This branch is a defensive measure, if we have custom sqlite3_mem_methods installed, and we + // must to hit this branch, then scope should not be a nullptr. If we see this, then something + // is wrong with our memory metering. + KJ_LOG(ERROR, "sqliteMemFree did not find a valid scope."); + } + 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; + if (scope == nullptr) { + // This branch is a defensive measure, if we have custom sqlite3_mem_methods installed, and we + // must to hit this branch, then scope should not be a nullptr. If we see this, then something + // is wrong with our memory metering. + KJ_LOG(ERROR, "sqliteMemRealloc did not find a valid scope."); + } + + size_t oldActual = usableSize(ptr); + + // sqliteMemRealloc must leave the original buffer intact per SQLite's contract when returning + // a nullptr. Because realloc consumes the original pointer, we must check whether the requested + // size would exceed the limit before calling realloc. + if (scope != nullptr && oldActual > 0 && static_cast(newSize) > oldActual) { + size_t growth = static_cast(newSize) - oldActual; + if (scope->memoryBytes + growth >= scope->maxMemoryBytes) { + return nullptr; + } + } + + void* newPtr = realloc(ptr, static_cast(newSize)); + if (newPtr == nullptr) return nullptr; + + if (scope != nullptr && oldActual > 0) { + size_t newActual = usableSize(newPtr); + // Note that `scope->memoryBytes + (newActual - oldActual) >= scope->maxMemoryBytes` could + // still be true, but realloc has already freed the original memory block, so we cannot return + // a nullptr. + if (scope->memoryBytes + newActual >= oldActual) { + scope->memoryBytes = scope->memoryBytes + newActual - oldActual; + } else { + // This branch is a defensive measure, but it is important since scope->memoryBytes is a size_t + // and an underflow would result in all subsequent allocations failing with SQLITE_NOMEM. If + // we see this, then something is wrong with our memory metering. + KJ_LOG(ERROR, "sqliteMemRealloc would have triggered a memoryBytes underflow."); + scope->memoryBytes = 0; + } + } + return newPtr; +} + +namespace { + +int sqliteMemSize(void* ptr) { + return static_cast(usableSize(ptr)); +} + +// sqliteMemRoundup is used by xRoundup as a performance optimization to tell SQLite how the memory +// allocator rounds up allocation requests. This implementation rounds up to the next multiple of 8 +// bytes because we assume that most allocators round up to the next 8-bytes. TCMalloc rounds up to +// the next 8 bytes for small size classes. +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& memoryBytes, size_t maxMemoryBytes) + : memoryBytes(memoryBytes), + maxMemoryBytes(maxMemoryBytes) { + if (threadLocalScope == nullptr) { + threadLocalScope = this; + } +} + +SqliteMemoryScope::~SqliteMemoryScope() noexcept(false) { + if (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 + // 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); + + // Hack: Explicitly initialize SQLite now with a SqliteMemoryScope to avoid cases that + // implicitly call sqlite3_initialize() without an active memory scope. sqlite3_initialize() is + // idempotent per-process (https://sqlite.org/c3ref/initialize.html). + static size_t initMemoryBytes = 0; + SqliteMemoryScope scope(initMemoryBytes, kj::maxValue); + sqlite3_initialize(); + 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..948abbe8fda --- /dev/null +++ b/src/workerd/util/sqlite-metering.h @@ -0,0 +1,69 @@ +// Copyright (c) 2026 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-database SQLite memory metering. +// +// SQLite uses a single process-wide memory allocator, but we want to account allocations against +// the specific SqliteDatabase on whose behalf they are made, and enforce a hard limit for memory +// allocations per-database. We do this by: +// +// 1. Installing a custom sqlite3_mem_methods that wraps the system allocator. +// +// 2. SqliteDatabase wraps every sqlite3_* entry point that may allocate or deallocate memory with +// a stack-allocated SqliteMemoryScope. The scope points at a size_t byte counter owned by +// SqliteDatabase 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-database hard limit by returning nullptr, signalling +// SQLite to throw a SQLITE_NOMEM exception. +// +// SqliteMemoryScope is idempotent: if a scope is already active on the current thread, additional +// scopes are no-ops. This allows safe nesting (e.g., when a column accessor calls into SQLite while +// a Query::nextRow() scope is already active). +class SqliteMemoryScope { + public: + // memoryBytes: the per-database byte counter owned by SqliteDatabase for its lifetime. + // maxMemoryBytes: the per-database cap from WorkerLimits::sqliteMaxMemoryMb. + // + // If a scope is already active on this thread, this constructor is a no-op (idempotent). + explicit SqliteMemoryScope(size_t& memoryBytes, size_t maxMemoryBytes); + + ~SqliteMemoryScope() noexcept(false); + KJ_DISALLOW_COPY_AND_MOVE(SqliteMemoryScope); + + private: + // These are accessed by the allocator functions via threadLocalScope. + size_t& memoryBytes; + const size_t maxMemoryBytes; + + // Thread-local pointer to the active scope. Set to this on construction (if not already set), + // cleared on destruction (if we set it). + 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-test.c++ b/src/workerd/util/sqlite-test.c++ index b27094e2555..5d3fb4cea5a 100644 --- a/src/workerd/util/sqlite-test.c++ +++ b/src/workerd/util/sqlite-test.c++ @@ -24,6 +24,14 @@ namespace workerd { namespace { +struct GlobalInit { + GlobalInit() { + installSqliteCustomAllocator(); + } +}; + +static GlobalInit init; + // Initialize the database with some data. void setupSql(SqliteDatabase& db) { // TODO(sqlite): Do this automatically and don't permit it via run(). @@ -271,6 +279,12 @@ KJ_TEST("Read-only database picks up on changes from mutable database (in-memory } KJ_TEST("Read-only database picks up on changes from mutable database (on-disk)") { +#if !_WIN32 + // doReadOnlyUpdateTest triggers the subsequent error in our memory metering because it creates + // two databases in the same thread. This will not happen in production. The on-disk test variant + // is impacted because it uses SQLite's native OS VFS, which can enable shared memory. + KJ_EXPECT_LOG(ERROR, "sqliteMemFree would have triggered a memoryBytes underflow."); +#endif TempDirOnDisk dir; doReadOnlyUpdateTest(*dir); } @@ -847,8 +861,8 @@ KJ_TEST("SQLite observer addQueryStats") { SqliteDatabase::Vfs vfs(*dir); TestSqliteObserver sqliteObserver = TestSqliteObserver(); TestQueryStatsRegulator regulator; - SqliteDatabase db( - vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, sqliteObserver); + SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, + /*sqliteMaxMemoryBytes=*/kj::maxValue, sqliteObserver); db.run(R"( CREATE TABLE things ( @@ -938,8 +952,8 @@ KJ_TEST("SQLite observer reportQueryEvent") { auto dir = kj::newInMemoryDirectory(kj::nullClock()); SqliteDatabase::Vfs vfs(*dir); TestSqliteObserver sqliteObserver; - SqliteDatabase db( - vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, sqliteObserver); + SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, + /*sqliteMaxMemoryBytes=*/kj::maxValue, sqliteObserver); db.run("PRAGMA journal_mode=WAL;"); @@ -1587,6 +1601,59 @@ class ErrorInjectableDirectory final: public kj::Directory, public kj::AtomicRef } }; +KJ_TEST("SQLite memory metering enforces SQLITE_NOMEM when limit is exceeded") { + auto dir = kj::newInMemoryDirectory(kj::nullClock()); + SqliteDatabase::Vfs vfs(*dir); + // 64 KiB is large enough for the database to open and run simple queries, + // but zeroblob(65536) requires allocating a 64 KiB buffer which, on top of + // the memory already consumed by opening the database, pushes total usage + // past the limit. This exercises our sqliteMemMalloc limit enforcement + // rather than SQLite's own hard_heap_limit pragma. + static constexpr size_t kTestMemoryLimit = 64 * 1024; + SqliteDatabase db( + vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, kTestMemoryLimit); + KJ_EXPECT_THROW_MESSAGE("out of memory: SQLITE_NOMEM", db.run("SELECT hex(zeroblob(65536))")); +} + +KJ_TEST("SQLite memory metering tracks allocations correctly") { + auto dir = kj::newInMemoryDirectory(kj::nullClock()); + SqliteDatabase::Vfs vfs(*dir); + SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, + /*sqliteMaxMemoryBytes=*/1024 * 1024); + size_t memoryBytes = db.getSqliteMemoryBytes(); + KJ_EXPECT(memoryBytes > 0, "memory should be greater than zero after creating a database"); + size_t memoryBytesSnapshot = memoryBytes; + + db.run("CREATE TABLE test (id INTEGER PRIMARY KEY, data TEXT)"); + memoryBytes = db.getSqliteMemoryBytes(); + KJ_EXPECT(memoryBytes >= memoryBytesSnapshot, "memory should not decrease when creating a table"); + memoryBytesSnapshot = memoryBytes; + + db.run("INSERT INTO test VALUES (1, 'hello world')"); + memoryBytes = db.getSqliteMemoryBytes(); + KJ_EXPECT( + memoryBytes >= memoryBytesSnapshot, "memory should not decrease when inserting into a table"); + memoryBytesSnapshot = memoryBytes; + + { + // Note that query has a different scope so that it does not prevent `PRAGMA shrink_memory` + // from taking effect. + auto query = db.run("SELECT * FROM test WHERE id = 1"); + KJ_ASSERT(!query.isDone()); + KJ_EXPECT(query.getInt(0) == 1); + KJ_EXPECT(query.getText(1) == "hello world"); + memoryBytes = db.getSqliteMemoryBytes(); + KJ_EXPECT( + memoryBytes >= memoryBytesSnapshot, "memory should not decrease when querying a table"); + memoryBytesSnapshot = memoryBytes; + } + + db.run("PRAGMA shrink_memory"); + memoryBytes = db.getSqliteMemoryBytes(); + KJ_EXPECT(memoryBytes < memoryBytesSnapshot, + "memory should decrease when running `PRAGMA shrink_memory`"); +} + KJ_TEST("I/O exceptions pass through SQLite") { auto dir = kj::atomicRefcounted(); SqliteDatabase::Vfs vfs(*dir); diff --git a/src/workerd/util/sqlite.c++ b/src/workerd/util/sqlite.c++ index 5f3c472b456..045c6db4a8b 100644 --- a/src/workerd/util/sqlite.c++ +++ b/src/workerd/util/sqlite.c++ @@ -571,12 +571,14 @@ SqliteObserver SqliteObserver::DEFAULT = SqliteObserver{}; SqliteDatabase::SqliteDatabase(const Vfs& vfs, kj::Path path, kj::Maybe maybeMode, + size_t sqliteMaxMemoryBytes, SqliteObserver& sqliteObserver, kj::Maybe actorAccountLimits) : vfs(vfs), path(kj::mv(path)), readOnly(maybeMode == kj::none), sqliteObserver(sqliteObserver), + sqliteMaxMemoryBytes(sqliteMaxMemoryBytes), actorAccountLimits(actorAccountLimits) { init(maybeMode); } @@ -585,6 +587,8 @@ void SqliteDatabase::init(kj::Maybe maybeMode) { KJ_ASSERT(maybeDb == kj::none); sqlite3* db = nullptr; + auto memoryScope = enterMemoryScope(); + KJ_IF_SOME(mode, maybeMode) { int flags = SQLITE_OPEN_READWRITE; if (kj::has(mode, kj::WriteMode::CREATE)) { @@ -629,6 +633,8 @@ void SqliteDatabase::init(kj::Maybe maybeMode) { SqliteDatabase::~SqliteDatabase() noexcept(false) { sqlite3* db = &KJ_UNWRAP_OR(maybeDb, return); + auto memoryScope = enterMemoryScope(); + auto err = sqlite3_close(db); if (err == SQLITE_BUSY) { KJ_LOG(ERROR, "sqlite database destroyed while dependent objects still exist"); @@ -646,6 +652,10 @@ SqliteDatabase::operator sqlite3*() { return &KJ_ASSERT_NONNULL(maybeDb, "previous reset() failed"); } +SqliteMemoryScope SqliteDatabase::enterMemoryScope() { + return SqliteMemoryScope(sqliteMemoryBytes, sqliteMaxMemoryBytes); +} + bool SqliteDatabase::observedCriticalError() { return criticalErrorOccurred; } @@ -808,6 +818,8 @@ SqliteDatabase::StatementAndEffect SqliteDatabase::prepareSql(const Regulator& r } }); + auto memoryScope = enterMemoryScope(); + for (;;) { sqlite3_stmt* result; const char* tail; @@ -966,6 +978,8 @@ void SqliteDatabase::executeWithRegulator( currentRegulator = regulator; KJ_DEFER(currentRegulator = kj::none); + + auto memoryScope = enterMemoryScope(); func(); } @@ -976,6 +990,8 @@ void SqliteDatabase::reset() { // confused at best, or lose data at worst. Let's just not allow this. KJ_REQUIRE(!inTransaction && savepoints.empty(), "can't reset() a database during a transaction"); + auto memoryScope = enterMemoryScope(); + // Temporarily disable the on-write callback while resetting. auto writeCb = kj::mv(onWriteCallback); KJ_DEFER(onWriteCallback = kj::mv(writeCb)); @@ -1441,6 +1457,14 @@ void SqliteDatabase::Statement::beforeSqliteReset() { } } +SqliteDatabase::Statement::~Statement() noexcept(false) { + // Install memory scope for sqlite3_finalize called when stmt (containing StatementAndEffect + // with kj::Own) is destroyed. Also covers prelude destruction. + auto memoryScope = db.enterMemoryScope(); + auto stmtToDestroy = kj::mv(stmt); + auto preludeToDestroy = kj::mv(prelude); +} + SqliteDatabase::Query::Query(SqliteDatabase& db, QueryOptions options, Statement& statement, @@ -1473,8 +1497,13 @@ SqliteDatabase::Query::~Query() noexcept(false) { } void SqliteDatabase::Query::destroy() { + // Install memory scope for sqlite3_reset, sqlite3_clear_bindings, and sqlite3_finalize (via + // ownStatement destruction). The scope is idempotent, so this is safe even if a scope is already + // active from the caller. + auto memoryScope = db.enterMemoryScope(); + if (regulator.shouldAddQueryStats()) { - //Update the db stats that we have collected for the query + // Update the db stats that we have collected for the query. db.sqliteObserver.addQueryStats(rowsRead, rowsWritten); } @@ -1490,9 +1519,14 @@ void SqliteDatabase::Query::destroy() { errorDescription.slice(0, kj::min(RA_MAX_METRICS_QUERY_SIZE, errorDescription.size())))); } + // Move ownStatement to a local variable so that it goes out of scope while memoryScope is still + // active. ownStatement is used by getStatementAndEffect(), so we cannot move it until this + // point. + auto ownStatementToDestroy = kj::mv(ownStatement); + // We only need to reset the statement if we don't own it. If we own it, it's about to be // destroyed anyway. - if (ownStatement.statement.get() == nullptr) { + if (ownStatementToDestroy.statement.get() == nullptr) { KJ_IF_SOME(statement, maybeStatement) { // The error code returned by sqlite3_reset() actually represents the last error encountered // when stepping the statement. This doesn't mean that the reset failed. @@ -1574,26 +1608,31 @@ uint64_t SqliteDatabase::Query::getRowsWritten() { } void SqliteDatabase::Query::bind(uint i, kj::ArrayPtr value) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); SQLITE_CALL(sqlite3_bind_blob(statement, i + 1, value.begin(), value.size(), SQLITE_STATIC)); } void SqliteDatabase::Query::bind(uint i, kj::StringPtr value) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); SQLITE_CALL(sqlite3_bind_text(statement, i + 1, value.begin(), value.size(), SQLITE_STATIC)); } void SqliteDatabase::Query::bind(uint i, long long value) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); SQLITE_CALL(sqlite3_bind_int64(statement, i + 1, value)); } void SqliteDatabase::Query::bind(uint i, double value) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); SQLITE_CALL(sqlite3_bind_double(statement, i + 1, value)); } void SqliteDatabase::Query::bind(uint i, decltype(nullptr)) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); SQLITE_CALL(sqlite3_bind_null(statement, i + 1)); } @@ -1611,6 +1650,7 @@ void SqliteDatabase::Query::nextRow(bool first) { KJ_DEFER(db.currentRegulator = kj::none); db.currentRegulator = regulator; + auto memoryScope = db.enterMemoryScope(); SQLITE_CALL_SCOPE { int err = sqlite3_step(statement); queryEvent.setQueryResult(err); @@ -1641,11 +1681,13 @@ uint SqliteDatabase::Query::changeCount() { } uint SqliteDatabase::Query::columnCount() { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); return sqlite3_column_count(statement); } SqliteDatabase::Query::ValuePtr SqliteDatabase::Query::getValue(uint column) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); switch (sqlite3_column_type(statement, column)) { case SQLITE_INTEGER: @@ -1663,38 +1705,45 @@ SqliteDatabase::Query::ValuePtr SqliteDatabase::Query::getValue(uint column) { } kj::StringPtr SqliteDatabase::Query::getColumnName(uint column) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); return sqlite3_column_name(statement, column); } kj::ArrayPtr SqliteDatabase::Query::getBlob(uint column) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); const byte* ptr = reinterpret_cast(sqlite3_column_blob(statement, column)); return kj::arrayPtr(ptr, sqlite3_column_bytes(statement, column)); } kj::StringPtr SqliteDatabase::Query::getText(uint column) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); const char* ptr = reinterpret_cast(sqlite3_column_text(statement, column)); return kj::StringPtr(ptr, sqlite3_column_bytes(statement, column)); } int SqliteDatabase::Query::getInt(uint column) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); return sqlite3_column_int(statement, column); } int64_t SqliteDatabase::Query::getInt64(uint column) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); return sqlite3_column_int64(statement, column); } double SqliteDatabase::Query::getDouble(uint column) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); return sqlite3_column_double(statement, column); } bool SqliteDatabase::Query::isNull(uint column) { + auto memoryScope = db.enterMemoryScope(); sqlite3_stmt* statement = getStatement(); return sqlite3_column_type(statement, column) == SQLITE_NULL; } diff --git a/src/workerd/util/sqlite.h b/src/workerd/util/sqlite.h index bb2b5a4a4d4..618e0df9395 100644 --- a/src/workerd/util/sqlite.h +++ b/src/workerd/util/sqlite.h @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include @@ -89,8 +90,15 @@ class SqliteDatabase { SqliteDatabase(const Vfs& vfs, kj::Path path, kj::Maybe maybeMode = kj::none, + size_t sqliteMaxMemoryBytes = kj::maxValue, SqliteObserver& sqliteObserver = SqliteObserver::DEFAULT, kj::Maybe actorAccountLimits = kj::none); + + // 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); @@ -219,6 +227,8 @@ class SqliteDatabase { onCriticalErrorCallback = kj::mv(callback); } + SqliteMemoryScope enterMemoryScope(); + // Returns true if a transaction was automatically rolled due to a critical error. bool observedCriticalError(); @@ -315,6 +325,15 @@ class SqliteDatabase { kj::Path path; bool readOnly; SqliteObserver& sqliteObserver; + + // The amount of memory in bytes used by this database for use by sqlite3_mem_methods. + size_t sqliteMemoryBytes = 0; + + // The maximum amount of memory in bytes that can be used by this database for use by + // sqlite3_mem_methods (from WorkerLimits::sqliteMaxMemoryMb). This is set to kj::maxValue to + // when running in workerd local development mode. + size_t sqliteMaxMemoryBytes; + kj::Maybe actorAccountLimits; // This pointer can be left null if a call to reset() failed to re-open the database. @@ -433,6 +452,15 @@ class SqliteDatabase { // Represents a prepared SQL statement, which can be executed many times. class SqliteDatabase::Statement final: private ResetListener { public: + Statement(Statement&& other) + : ResetListener(kj::mv(other)), + regulator(other.regulator), + stmt(kj::mv(other.stmt)), + prelude(kj::mv(other.prelude)) {} + + // Destructor needs to install a memory scope for sqlite3_finalize. + ~Statement() noexcept(false); + // Convenience method to start a query. This is equivalent to: // // SqliteDatabase::Query(db, statement, bindings...);