1
// Copyright (c) 2026 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
#pragma once
6
 
7
#include <stddef.h>
8
 
9
#include <kj/common.h>
10
 
11
namespace workerd {
12
 
13
// This module implements per-database SQLite memory metering.
14
//
15
// SQLite uses a single process-wide memory allocator, but we want to account allocations against
16
// the specific SqliteDatabase on whose behalf they are made, and enforce a hard limit for memory
17
// allocations per-database. We do this by:
18
//
19
// 1. Installing a custom sqlite3_mem_methods that wraps the system allocator.
20
//
21
// 2. SqliteDatabase wraps every sqlite3_* entry point that may allocate or deallocate memory with
22
//    a stack-allocated SqliteMemoryScope. The scope points at a size_t byte counter owned by
23
//    SqliteDatabase to meter memory allocations.
24
//
25
// 3. When a SqliteMemoryScope is active on the current thread, we count each memory allocation
26
//    against the byte counter and enforce the per-database hard limit by returning nullptr, signalling
27
//    SQLite to throw a SQLITE_NOMEM exception.
28
//
29
// SqliteMemoryScope is idempotent: if a scope is already active on the current thread, additional
30
// scopes are no-ops. This allows safe nesting (e.g., when a column accessor calls into SQLite while
31
// a Query::nextRow() scope is already active).
32
class SqliteMemoryScope {
33
 public:
34
  // memoryBytes: the per-database byte counter owned by SqliteDatabase for its lifetime.
35
  // maxMemoryBytes: the per-database cap from WorkerLimits::sqliteMaxMemoryMb.
36
  //
37
  // If a scope is already active on this thread, this constructor is a no-op (idempotent).
38
  explicit SqliteMemoryScope(size_t& memoryBytes, size_t maxMemoryBytes);
39
 
40
  ~SqliteMemoryScope() noexcept(false);
41
  KJ_DISALLOW_COPY_AND_MOVE(SqliteMemoryScope);
42
 
43
 private:
44
  // These are accessed by the allocator functions via threadLocalScope.
45
  size_t& memoryBytes;
46
  const size_t maxMemoryBytes;
47
 
48
  // Thread-local pointer to the active scope. Set to this on construction (if not already set),
49
  // cleared on destruction (if we set it).
50
  static thread_local SqliteMemoryScope* threadLocalScope;
51
 
52
  friend void* sqliteMemMalloc(int);
53
  friend void sqliteMemFree(void*);
54
  friend void* sqliteMemRealloc(void*, int);
55
};
56
 
57
// The custom sqlite3_mem_methods functions. Declared here so tests can call them directly to
58
// verify accounting behaviour without going through the SQLite query engine.
59
void* sqliteMemMalloc(int size);
60
void sqliteMemFree(void* ptr);
61
void* sqliteMemRealloc(void* ptr, int newSize);
62
 
63
// Install the custom sqlite3_mem_methods.
64
//
65
// This must be called before the first sqlite3_initialize(), sqlite3_open_v2(), or
66
// sqlite3_vfs_register() call in the process. Idempotent (uses a static once-flag).
67
void installSqliteCustomAllocator();
68
 
69
}  // namespace workerd