Skip to content
rachel/workerd
File

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

7.4 KB
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#include "sqlite-metering.h"
6 
7#include <sqlite3.h>
8#include <stdlib.h>
9 
10#include <kj/debug.h>
11 
12#if defined(__linux__) || defined(__GLIBC__)
13#include <malloc.h>
14#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) malloc_usable_size(p)
15#elifdef __APPLE__
16#include <malloc/malloc.h>
17#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) malloc_size(p)
18#elifdef _WIN32
19#include <malloc.h>
20#define SQLITE_METERING_MALLOC_USABLE_SIZE(p) _msize(p)
21#else
22#error "Platform not supported"
23#endif
24 
25namespace {
26 
27inline size_t usableSize(void* ptr) {
28  if (ptr == nullptr) return 0;
29  return SQLITE_METERING_MALLOC_USABLE_SIZE(ptr);
30}
31 
32}  // namespace
33 
34namespace workerd {
35 
36thread_local SqliteMemoryScope* SqliteMemoryScope::threadLocalScope = nullptr;
37 
38void* sqliteMemMalloc(int size) {
39  if (size <= 0) return nullptr;
40 
41  SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope;
42  if (scope != nullptr) {
43    if (scope->memoryBytes + static_cast<size_t>(size) > scope->maxMemoryBytes) {
44      return nullptr;
45    }
46  } else {
47    // This branch is a defensive measure, if we have custom sqlite3_mem_methods installed, and we
48    // must to hit this branch, then scope should not be a nullptr. If we see this, then something
49    // is wrong with our memory metering.
50    KJ_LOG(ERROR, "sqliteMemMalloc did not find a valid scope.");
51  }
52 
53  void* ptr = malloc(static_cast<size_t>(size));
54  if (ptr == nullptr) return nullptr;
55  if (scope != nullptr) {
56    size_t actual = usableSize(ptr);
57    if (actual > 0) {
58      // We have already checked that allocating the requested size will not exceed maxMemoryBytes,
59      // but we still have to check that the actual size allocated does not exceed maxMemoryBytes.
60      if (scope->memoryBytes + actual > scope->maxMemoryBytes) {
61        free(ptr);
62        return nullptr;
63      }
64      scope->memoryBytes += actual;
65    }
66  }
67  return ptr;
68}
69 
70void sqliteMemFree(void* ptr) {
71  if (ptr == nullptr) return;
72  SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope;
73  if (scope != nullptr) {
74    size_t actual = usableSize(ptr);
75    if (actual > 0) {
76      if (scope->memoryBytes >= actual) {
77        scope->memoryBytes -= actual;
78      } else {
79        // This branch is a defensive measure, but it is important since scope->memoryBytes is a size_t
80        // and an underflow would result in all subsequent allocations failing with SQLITE_NOMEM.
81        // If we see this, then something is wrong with our memory metering.
82        KJ_LOG(ERROR, "sqliteMemFree would have triggered a memoryBytes underflow.");
83        scope->memoryBytes = 0;
84      }
85    }
86  } else {
87    // This branch is a defensive measure, if we have custom sqlite3_mem_methods installed, and we
88    // must to hit this branch, then scope should not be a nullptr. If we see this, then something
89    // is wrong with our memory metering.
90    KJ_LOG(ERROR, "sqliteMemFree did not find a valid scope.");
91  }
92  free(ptr);
93}
94 
95void* sqliteMemRealloc(void* ptr, int newSize) {
96  // SQLite assumes that sqliteMemRealloc(ptr, newSize) where newSize <= 0 is equivalent to
97  // sqliteMemFree(ptr).
98  if (newSize <= 0) {
99    sqliteMemFree(ptr);
100    return nullptr;
101  }
102 
103  // SQLite assumes that sqliteMemRealloc(nullptr, newSize) is equivalent to
104  // sqliteMemMalloc(newSize).
105  if (ptr == nullptr) {
106    return sqliteMemMalloc(newSize);
107  }
108 
109  SqliteMemoryScope* scope = SqliteMemoryScope::threadLocalScope;
110  if (scope == nullptr) {
111    // This branch is a defensive measure, if we have custom sqlite3_mem_methods installed, and we
112    // must to hit this branch, then scope should not be a nullptr. If we see this, then something
113    // is wrong with our memory metering.
114    KJ_LOG(ERROR, "sqliteMemRealloc did not find a valid scope.");
115  }
116 
117  size_t oldActual = usableSize(ptr);
118 
119  // sqliteMemRealloc must leave the original buffer intact per SQLite's contract when returning
120  // a nullptr. Because realloc consumes the original pointer, we must check whether the requested
121  // size would exceed the limit before calling realloc.
122  if (scope != nullptr && oldActual > 0 && static_cast<size_t>(newSize) > oldActual) {
123    size_t growth = static_cast<size_t>(newSize) - oldActual;
124    if (scope->memoryBytes + growth >= scope->maxMemoryBytes) {
125      return nullptr;
126    }
127  }
128 
129  void* newPtr = realloc(ptr, static_cast<size_t>(newSize));
130  if (newPtr == nullptr) return nullptr;
131 
132  if (scope != nullptr && oldActual > 0) {
133    size_t newActual = usableSize(newPtr);
134    // Note that `scope->memoryBytes + (newActual - oldActual) >= scope->maxMemoryBytes` could
135    // still be true, but realloc has already freed the original memory block, so we cannot return
136    // a nullptr.
137    if (scope->memoryBytes + newActual >= oldActual) {
138      scope->memoryBytes = scope->memoryBytes + newActual - oldActual;
139    } else {
140      // This branch is a defensive measure, but it is important since scope->memoryBytes is a size_t
141      // and an underflow would result in all subsequent allocations failing with SQLITE_NOMEM. If
142      // we see this, then something is wrong with our memory metering.
143      KJ_LOG(ERROR, "sqliteMemRealloc would have triggered a memoryBytes underflow.");
144      scope->memoryBytes = 0;
145    }
146  }
147  return newPtr;
148}
149 
150namespace {
151 
152int sqliteMemSize(void* ptr) {
153  return static_cast<int>(usableSize(ptr));
154}
155 
156// sqliteMemRoundup is used by xRoundup as a performance optimization to tell SQLite how the memory
157// allocator rounds up allocation requests. This implementation rounds up to the next multiple of 8
158// bytes because we assume that most allocators round up to the next 8-bytes. TCMalloc rounds up to
159// the next 8 bytes for small size classes.
160int sqliteMemRoundup(int n) {
161  // Round up to the next multiple of 8 as a conservative estimate.
162  return (n + 7) & ~7;
163}
164 
165int sqliteMemInit(void* /*pAppData*/) {
166  return SQLITE_OK;
167}
168 
169void sqliteMemShutdown(void* /*pAppData*/) {}
170 
171static const sqlite3_mem_methods kSqliteMemMethods = {
172  sqliteMemMalloc,
173  sqliteMemFree,
174  sqliteMemRealloc,
175  sqliteMemSize,
176  sqliteMemRoundup,
177  sqliteMemInit,
178  sqliteMemShutdown,
179  /*pAppData=*/nullptr,
180};
181 
182}  // namespace
183 
184SqliteMemoryScope::SqliteMemoryScope(size_t& memoryBytes, size_t maxMemoryBytes)
185    : memoryBytes(memoryBytes),
186      maxMemoryBytes(maxMemoryBytes) {
187  if (threadLocalScope == nullptr) {
188    threadLocalScope = this;
189  }
190}
191 
192SqliteMemoryScope::~SqliteMemoryScope() noexcept(false) {
193  if (threadLocalScope == this) {
194    threadLocalScope = nullptr;
195  }
196}
197 
198void installSqliteCustomAllocator() {
199  // sqlite3_config() must be called before sqlite3_initialize(), which is itself invoked
200  // implicitly by the first sqlite3_vfs_register() or sqlite3_open_v2() call. We rely on callers
201  // to invoke this before any of those.
202  static bool installed KJ_UNUSED = []() {
203    int rc = sqlite3_config(SQLITE_CONFIG_MALLOC, &kSqliteMemMethods);
204    KJ_ASSERT(rc == SQLITE_OK, "sqlite3_config(SQLITE_CONFIG_MALLOC) failed", rc);
205 
206    // Hack: Explicitly initialize SQLite now with a SqliteMemoryScope to avoid cases that
207    // implicitly call sqlite3_initialize() without an active memory scope. sqlite3_initialize() is
208    // idempotent per-process (https://sqlite.org/c3ref/initialize.html).
209    static size_t initMemoryBytes = 0;
210    SqliteMemoryScope scope(initMemoryBytes, kj::maxValue);
211    sqlite3_initialize();
212    return true;
213  }();
214}
215 
216}  // namespace workerd