Attribute SQLite memory consumed to isolate memory consumed

src/workerd/io/BUILD.bazel

@@ -202,10 +202,12 @@ wd_cc_library(
         ":actor-storage_capnp",
         ":io-gate",
         ":trace",
+        "//src/workerd/jsg",
         "//src/workerd/jsg:exception",
         "//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-cache.h

@@ -7,6 +7,7 @@
 #include <workerd/io/actor-storage.capnp.h>
 #include <workerd/io/trace.h>
 #include <workerd/jsg/exception.h>
+#include <workerd/jsg/jsg.h>
 
 #include <kj/async.h>
 #include <kj/debug.h>
@@ -191,6 +192,18 @@ class ActorCacheInterface: public ActorCacheOps {
   // old-style DOs have asyncronous storage.
   virtual kj::Maybe<SqliteKv&> getSqliteKv() = 0;
 
+  // If the actor's storage is backed by SQLite, return a reference to the ExternalMemoryAdjustment
+  // that tracks bytes currently allocated by SQLite on behalf of this actor, once it has been
+  // initialized.
+  virtual kj::Maybe<jsg::ExternalMemoryAdjustment&> getSqliteMemoryAdjustment() {
+    return kj::none;
+  }
+
+  // Called by LimitEnforcerImpl::enterJs() on each JS turn to ensure SQLite memory accounting is
+  // initialized for this actor. Idempotent: only initializes on the first call. The default
+  // implementation is a no-op for non-SQLite-backed actors.
+  virtual void setSqliteMemoryAdjustment(kj::Arc<const jsg::ExternalMemoryTarget> /*target*/) {}
+
   class Transaction: public ActorCacheOps {
    public:
     // Write all changes to the underlying ActorCache.

src/workerd/io/actor-sqlite.c++

@@ -7,6 +7,7 @@
 #include "io-gate.h"
 
 #include <workerd/jsg/exception.h>
+#include <workerd/jsg/setup.h>
 #include <workerd/util/autogate.h>
 #include <workerd/util/sentry.h>
 
@@ -64,6 +65,16 @@ ActorSqlite::ActorSqlite(kj::Own<SqliteDatabase> dbParam,
   alarmScheduledNoLaterThan = metadata.getAlarm();
 }
 
+kj::Maybe<jsg::ExternalMemoryAdjustment&> ActorSqlite::getSqliteMemoryAdjustment() {
+  return sqliteMemoryAdjustment;
+}
+
+void ActorSqlite::setSqliteMemoryAdjustment(kj::Arc<const jsg::ExternalMemoryTarget> target) {
+  if (sqliteMemoryAdjustment == kj::none) {
+    sqliteMemoryAdjustment = target->getAdjustment(0);
+  }
+}
+
 ActorSqlite::ImplicitTxn::ImplicitTxn(ActorSqlite& parent): parent(parent) {
   KJ_REQUIRE(parent.currentTxn.is<NoTxn>());
   parent.beginTxn.run();

src/workerd/io/actor-sqlite.h

@@ -7,8 +7,10 @@
 #include "actor-cache.h"
 
 #include <workerd/io/trace.h>
+#include <workerd/jsg/jsg.h>
 #include <workerd/util/sqlite-kv.h>
 #include <workerd/util/sqlite-metadata.h>
+#include <workerd/util/sqlite-metering.h>
 
 namespace workerd {
 
@@ -68,6 +70,15 @@ class ActorSqlite final: public ActorCacheInterface, private kj::TaskSet::ErrorH
     return kv;
   }
 
+  // Return a reference to the ExternalMemoryAdjustment that tracks bytes currently allocated by
+  // SQLite on behalf of this actor, once it has been initialized.
+  kj::Maybe<jsg::ExternalMemoryAdjustment&> getSqliteMemoryAdjustment() override;
+
+  // Called by LimitEnforcerImpl::enterJs() on each JS turn to ensure SQLite memory accounting is
+  // initialized for this actor. Guards against being called more than once to avoid resetting
+  // sqliteMemoryAdjustment, which would cause the running byte total to be lost.
+  void setSqliteMemoryAdjustment(kj::Arc<const jsg::ExternalMemoryTarget> target) override;
+
   kj::OneOf<kj::Maybe<Value>, kj::Promise<kj::Maybe<Value>>> get(
       Key key, ReadOptions options) override;
   kj::OneOf<GetResultList, kj::Promise<GetResultList>> get(
@@ -114,6 +125,10 @@ class ActorSqlite final: public ActorCacheInterface, private kj::TaskSet::ErrorH
   SqliteKv kv;
   SqliteMetadata metadata;
 
+  // Per-actor SQLite memory metering state. Populated lazily by setSqliteMemoryAdjustment() the
+  // first time the actor's isolate lock is taken via LimitEnforcerImpl::enterJs().
+  kj::Maybe<jsg::ExternalMemoryAdjustment> sqliteMemoryAdjustment;
+
   // Define a SqliteDatabase::Regulator that is similar to TRUSTED but turns certain SQLite errors
   // into application errors as appropriate when committing an implicit transaction.
   class TxnCommitRegulator: public SqliteDatabase::Regulator {

src/workerd/util/BUILD.bazel

@@ -191,6 +191,18 @@ wd_cc_library(
     ],
 )
 
+wd_cc_library(
+    name = "sqlite-metering",
+    srcs = ["sqlite-metering.c++"],
+    hdrs = ["sqlite-metering.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//src/workerd/jsg",
+        "@capnp-cpp//src/kj",
+        "@sqlite3",
+    ],
+)
+
 wd_cc_library(
     name = "sqlite",
     srcs = [
@@ -211,6 +223,7 @@ wd_cc_library(
     deps = [
         ":account-limits",
         ":sentry",
+        ":sqlite-metering",
         "@capnp-cpp//src/kj:kj-async",
     ],
 )

src/workerd/util/sqlite-metering.c++

@@ -0,0 +1,153 @@
+// 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 <sqlite3.h>
+
+#include <kj/debug.h>
+
+#include <cstdlib>
+
+#if defined(__linux__) || defined(__GLIBC__)
+#include <malloc.h>
+#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) malloc_usable_size(p)
+#elif defined(__APPLE__)
+#include <malloc/malloc.h>
+#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) malloc_size(p)
+#elif defined(_WIN32)
+#include <malloc.h>
+#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;
+
+// ---------------------------------------------------------------------------
+// Custom sqlite3_mem_methods implementation
+//
+// All actual allocation is delegated to the system allocator (malloc/free/realloc).
+// We layer per-actor accounting on top by adjusting the ExternalMemoryAdjustment when
+// threadLocalScope is non-null.
+//
+// Thread-safety: since threadLocalScope is thread-local, each thread sees its own actor's
+// scope independently with no locking needed.
+// ---------------------------------------------------------------------------
+
+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_t>(size));
+  if (ptr == nullptr) return nullptr;
+  SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope;
+  if (scope != nullptr) {
+    scope->adjustment.adjust(static_cast<ssize_t>(usableSize(ptr)));
+  }
+  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->adjustment.adjust(-static_cast<ssize_t>(actual));
+    }
+  }
+  std::free(ptr);
+}
+
+void* sqliteMemRealloc(void* ptr, int newSize) {
+  if (newSize <= 0) {
+    sqliteMemFree(ptr);
+    return nullptr;
+  }
+  SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope;
+  size_t oldActual = (scope != nullptr && ptr != nullptr) ? usableSize(ptr) : 0;
+
+  void* newPtr = std::realloc(ptr, static_cast<size_t>(newSize));
+  if (newPtr == nullptr) return nullptr;
+
+  if (scope != nullptr && oldActual > 0) {
+    size_t newActual = usableSize(newPtr);
+    if (newActual > 0) {
+      ssize_t delta = static_cast<ssize_t>(newActual) - static_cast<ssize_t>(oldActual);
+      scope->adjustment.adjust(delta);
+    }
+  }
+  return newPtr;
+}
+
+namespace {
+
+int sqliteMemSize(void* ptr) {
+  return static_cast<int>(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(jsg::ExternalMemoryAdjustment& adjustment)
+    : adjustment(adjustment) {
+  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() {
+  // Use a static once-flag. 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 call 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 ExternalMemoryAdjustment.
+    rc = sqlite3_config(SQLITE_CONFIG_MEMSTATUS, 0);
+    KJ_ASSERT(rc == SQLITE_OK, "sqlite3_config(SQLITE_CONFIG_MEMSTATUS) failed", rc);
+
+    return true;
+  }();
+}
+
+}  // namespace workerd

src/workerd/util/sqlite-metering.h

@@ -0,0 +1,57 @@
+// 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
+
+// Per-isolate 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. We do this by:
+//
+// 1. Installing a custom sqlite3_mem_methods that wraps the system allocator and checks a
+//    thread-local pointer on every allocation.
+//
+// 2. Installing a SqliteMemoryScope (RAII) on the current thread for the duration of each JS
+//    turn that may invoke SQLite, pointing at the actor's ExternalMemoryAdjustment.
+//
+// When the thread-local pointer is null (outside of an actor JS turn, or when the
+// SQLITE_MEMORY_METERING autogate is disabled), allocations pass through to the underlying
+// allocator with no accounting.
+
+#include <workerd/jsg/jsg.h>
+
+#include <kj/common.h>
+
+namespace workerd {
+
+// RAII scope: installs this instance as the active per-actor SQLite memory context on the current
+// thread for the duration of one JS turn, then clears it on destruction.
+//
+// Only one SqliteMemoryScope may be active per thread at a time. Construction asserts that no
+// other scope is active, mirroring the CPU limit enforcement in LimitEnforcerImpl.
+class SqliteMemoryScope {
+ public:
+  explicit SqliteMemoryScope(jsg::ExternalMemoryAdjustment& adjustment);
+  ~SqliteMemoryScope() noexcept(false);
+  KJ_DISALLOW_COPY_AND_MOVE(SqliteMemoryScope);
+
+  jsg::ExternalMemoryAdjustment& adjustment;
+
+ 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);
+};
+
+// 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. The implementation of this function must be
+// idempotent.
+void installSqliteCustomAllocator();
+
+}  // namespace workerd

src/workerd/util/sqlite.c++

@@ -4,6 +4,8 @@
 
 #include "sqlite.h"
 
+#include "sqlite-metering.h"
+
 #include <workerd/util/sentry.h>
 
 #include <kj/debug.h>
@@ -2588,6 +2590,7 @@ SqliteDatabase::Vfs::Vfs(const kj::Directory& directory, Options options)
       lockManager(*ownLockManager),
       options(kj::mv(options)),
       native(*sqlite3_vfs_find(nullptr)) {
+  installSqliteCustomAllocator();
 #if _WIN32
   vfs = kj::heap(makeKjVfs());
 #else
@@ -2609,6 +2612,7 @@ SqliteDatabase::Vfs::Vfs(
       native(*sqlite3_vfs_find(nullptr)),
       // Always use KJ VFS when using a custom LockManager.
       vfs(kj::heap(makeKjVfs())) {
+  installSqliteCustomAllocator();
   sqlite3_vfs_register(vfs, false);
 }