Meter memory for SQLite on a per-DO basis

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

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 {

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 {

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 = [

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

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);

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 <kj/test.h>
+#include <kj/thread.h>
+
+#include <atomic>
+
+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, <n>) should behave as sqliteMemMalloc(<n>).
+  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