SQLite: Cache prepared statements behind exec().

src/workerd/api/sql-test.js

@@ -19,6 +19,10 @@ async function test(state) {
   assert.equal(resultNumberRaw[0].length, 1);
   assert.equal(resultNumberRaw[0][0], 123);
 
+  sql.exec('SELECT 123');
+  sql.exec('SELECT 123');
+  sql.exec('SELECT 123');
+
   // Test string results
   const resultStr = [...sql.exec("SELECT 'hello'")];
   assert.equal(resultStr.length, 1);

src/workerd/api/sql.c++

@@ -10,14 +10,60 @@
 
 namespace workerd::api {
 
-SqlStorage::SqlStorage(jsg::Ref<DurableObjectStorage> storage): storage(kj::mv(storage)) {}
+// Maximum total size of all cached statements (measured in size of the SQL code). If cached
+// statements exceed this, we remove the LRU statement(s).
+//
+// Hopefully most apps don't ever hit this, but it's important to have a limit in case of
+// queries containing dynamic content or excessively large one-off queries.
+static constexpr uint SQL_STATEMENT_CACHE_MAX_SIZE = 1024 * 1024;
+
+SqlStorage::SqlStorage(jsg::Ref<DurableObjectStorage> storage)
+    : storage(kj::mv(storage)),
+      statementCache(IoContext::current().addObject(kj::heap<StatementCache>())) {}
 
 SqlStorage::~SqlStorage() {}
 
 jsg::Ref<SqlStorage::Cursor> SqlStorage::exec(
     jsg::Lock& js, jsg::JsString querySql, jsg::Arguments<BindingValue> bindings) {
-  SqliteDatabase::Regulator& regulator = *this;
-  return jsg::alloc<Cursor>(getDb(js), regulator, js.toString(querySql), kj::mv(bindings));
+  auto& db = getDb(js);
+  auto& statementCache = *this->statementCache;
+
+  kj::Rc<CachedStatement>& slot = statementCache.map.findOrCreate(querySql, [&]() {
+    auto result = kj::rc<CachedStatement>(js, *this, db, querySql, js.toString(querySql));
+    statementCache.totalSize += result->statementSize;
+    return result;
+  });
+
+  // Move cached statement to end of LRU queue.
+  if (slot->lruLink.isLinked()) {
+    statementCache.lru.remove(*slot.get());
+  }
+  statementCache.lru.add(*slot.get());
+
+  if (slot->isShared()) {
+    // Oops, this CachedStatement is currently in-use (presumably by a Cursor).
+    //
+    // SQLite only allows one instance of a statement to run at a time, so we will have to compile
+    // the statement again as a one-off.
+    //
+    // In theory we could try to cache multiple copies of the statement, but as this is probably
+    // exceedingly rare, it is not worth the added code complexity.
+    SqliteDatabase::Regulator& regulator = *this;
+    return jsg::alloc<Cursor>(db, regulator, js.toString(querySql), kj::mv(bindings));
+  }
+
+  auto result = jsg::alloc<Cursor>(slot.addRef(), kj::mv(bindings));
+
+  // If the statement cache grew too big, drop the least-recently-used entry.
+  while (statementCache.totalSize > SQL_STATEMENT_CACHE_MAX_SIZE) {
+    auto& toRemove = *statementCache.lru.begin();
+    auto oldQuery = jsg::JsString(toRemove.query.getHandle(js));
+    statementCache.totalSize -= toRemove.statementSize;
+    statementCache.lru.remove(toRemove);
+    KJ_ASSERT(statementCache.map.eraseMatch(oldQuery));
+  }
+
+  return result;
 }
 
 SqlStorage::IngestResult SqlStorage::ingest(jsg::Lock& js, kj::String querySql) {
@@ -60,6 +106,12 @@ bool SqlStorage::allowTransactions() const {
       "write coalescing.");
 }
 
+SqlStorage::StatementCache::~StatementCache() noexcept(false) {
+  for (auto& entry: lru) {
+    lru.remove(entry);
+  }
+}
+
 jsg::JsValue SqlStorage::wrapSqlValue(jsg::Lock& js, SqlValue value) {
   KJ_IF_SOME(v, value) {
     KJ_SWITCH_ONEOF(v) {
@@ -106,6 +158,12 @@ SqlStorage::Cursor::State::State(SqliteDatabase& db,
     : bindings(kj::mv(bindingsParam)),
       query(db.run(regulator, sqlCode, mapBindings(bindings).asPtr())) {}
 
+SqlStorage::Cursor::State::State(
+    kj::Rc<CachedStatement> cachedStatementParam, kj::Array<BindingValue> bindingsParam)
+    : bindings(kj::mv(bindingsParam)),
+      query(cachedStatement.emplace(kj::mv(cachedStatementParam))
+                ->statement.run(mapBindings(bindings).asPtr())) {}
+
 SqlStorage::Cursor::~Cursor() noexcept(false) {
   // If this Cursor was created from a Statement, clear the Statement's currentCursor weak ref.
   KJ_IF_SOME(s, selfRef) {

src/workerd/api/sql.h

@@ -107,6 +107,58 @@ class SqlStorage final: public jsg::Object, private SqliteDatabase::Regulator {
     kj::Maybe<kj::Array<jsg::JsRef<jsg::JsString>>> names;
   };
 
+  // A statement in the statement cache.
+  struct CachedStatement: public kj::Refcounted {
+    jsg::HashableV8Ref<v8::String> query;
+    size_t statementSize;
+    SqliteDatabase::Statement statement;
+    CachedColumnNames cachedColumnNames;
+    kj::ListLink<CachedStatement> lruLink;
+
+    CachedStatement(jsg::Lock& js,
+        SqlStorage& sqlStorage,
+        SqliteDatabase& db,
+        jsg::JsString jsQuery,
+        kj::String kjQuery)
+        : query(js.v8Isolate, jsQuery),
+          statementSize(kjQuery.size()),
+          statement(db.prepareMulti(sqlStorage, kj::mv(kjQuery))) {}
+  };
+
+  class StatementCacheCallbacks {
+  public:
+    inline const jsg::HashableV8Ref<v8::String>& keyForRow(
+        const kj::Rc<CachedStatement>& entry) const {
+      return entry->query;
+    }
+
+    inline bool matches(const kj::Rc<CachedStatement>& entry, jsg::JsString key) const {
+      return entry->query == key;
+    }
+    inline bool matches(
+        const kj::Rc<CachedStatement>& entry, const jsg::HashableV8Ref<v8::String>& key) const {
+      return entry->query == key;
+    }
+
+    inline auto hashCode(jsg::JsString key) const {
+      return key.hashCode();
+    }
+    inline auto hashCode(const jsg::HashableV8Ref<v8::String>& key) const {
+      return key.hashCode();
+    }
+  };
+
+  using StatementMap = kj::Table<kj::Rc<CachedStatement>, kj::HashIndex<StatementCacheCallbacks>>;
+
+  struct StatementCache {
+    StatementMap map;
+    kj::List<CachedStatement, &CachedStatement::lruLink> lru;
+    size_t totalSize = 0;
+
+    ~StatementCache() noexcept(false);
+  };
+  IoOwn<StatementCache> statementCache;
+
   template <size_t size, typename... Params>
   SqliteDatabase::Query execMemoized(SqliteDatabase& db,
       kj::Maybe<IoOwn<SqliteDatabase::Statement>>& slot,
@@ -196,6 +248,8 @@ class SqlStorage::Cursor final: public jsg::Object {
 
 private:
   struct State {
+    kj::Maybe<kj::Rc<CachedStatement>> cachedStatement;
+
     // The bindings that were used to construct `query`. We have to keep these alive until the query
     // is done since it might contain pointers into strings and blobs.
     kj::Array<BindingValue> bindings;
@@ -208,6 +262,8 @@ class SqlStorage::Cursor final: public jsg::Object {
         SqliteDatabase::Regulator& regulator,
         kj::StringPtr sqlCode,
         kj::Array<BindingValue> bindings);
+
+    State(kj::Rc<CachedStatement> cachedStatement, kj::Array<BindingValue> bindings);
   };
 
   // Nulled out when query is done or canceled.