Implement per-actor SQLite memory metering.

Author: joshthoward

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",
     ],

src/workerd/io/actor-sqlite.h

@@ -9,6 +9,7 @@
 #include <workerd/io/trace.h>
 #include <workerd/util/sqlite-kv.h>
 #include <workerd/util/sqlite-metadata.h>
+#include <workerd/util/sqlite-metering.h>
 
 namespace workerd {
 

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

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

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 <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 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, <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 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