Group queries in SRS by a top-level queryId

src/workerd/api/sql.c++

@@ -7,6 +7,7 @@
 #include "actor-state.h"
 
 #include <workerd/io/io-context.h>
+#include <workerd/util/sqlite.h>
 
 #if _WIN32
 #define strncasecmp _strnicmp
@@ -72,6 +73,8 @@ jsg::Ref<SqlStorage::Cursor> SqlStorage::exec(
     };
   }
 
+  auto queryId = db.generateQueryId();
+
   if (slot->isShared()) {
     // Oops, this CachedStatement is currently in-use (presumably by a Cursor).
     //
@@ -81,10 +84,12 @@ jsg::Ref<SqlStorage::Cursor> SqlStorage::exec(
     // 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.
     return js.alloc<Cursor>(js, kj::mv(doneCallback), db,
-        SqliteDatabase::StaticRegulator(regulator), js.toString(querySql), kj::mv(bindings));
+        SqliteDatabase::StaticRegulator(regulator), js.toString(querySql), kj::mv(bindings),
+        queryId);
   }
 
-  auto result = js.alloc<Cursor>(js, kj::mv(doneCallback), slot.addRef(), kj::mv(bindings));
+  auto result =
+      js.alloc<Cursor>(js, kj::mv(doneCallback), slot.addRef(), kj::mv(bindings), queryId);
 
   // If the statement cache grew too big, drop the least-recently-used entry.
   while (statementCache.totalSize > SQL_STATEMENT_CACHE_MAX_SIZE) {
@@ -195,15 +200,19 @@ jsg::JsValue SqlStorage::wrapSqlValue(jsg::Lock& js, SqlValue value) {
 SqlStorage::Cursor::State::State(SqliteDatabase& db,
     SqliteDatabase::StaticRegulator regulator,
     kj::StringPtr sqlCode,
-    kj::Array<BindingValue> bindingsParam)
+    kj::Array<BindingValue> bindingsParam,
+    kj::Maybe<SqliteQueryId> queryId)
     : bindings(kj::mv(bindingsParam)),
-      query(db.run({.regulator = regulator}, sqlCode, mapBindings(bindings).asPtr())) {}
+      query(db.run(
+          {.regulator = regulator, .queryId = queryId}, sqlCode, mapBindings(bindings).asPtr())) {}
 
-SqlStorage::Cursor::State::State(
-    kj::Rc<CachedStatement> cachedStatementParam, kj::Array<BindingValue> bindingsParam)
+SqlStorage::Cursor::State::State(kj::Rc<CachedStatement> cachedStatementParam,
+    kj::Array<BindingValue> bindingsParam,
+    kj::Maybe<SqliteQueryId> queryId)
     : bindings(kj::mv(bindingsParam)),
       query(cachedStatement.emplace(kj::mv(cachedStatementParam))
-                ->statement.run(mapBindings(bindings).asPtr())) {}
+                ->statement.run(SqliteDatabase::Statement::StatementOptions{.queryId = queryId},
+                    mapBindings(bindings).asPtr())) {}
 
 SqlStorage::Cursor::~Cursor() noexcept(false) {
   // If this Cursor was created from a Statement, clear the Statement's currentCursor weak ref.

src/workerd/api/sql.h

@@ -257,9 +257,12 @@ class SqlStorage::Cursor final: public jsg::Object {
     State(SqliteDatabase& db,
         SqliteDatabase::StaticRegulator regulator,
         kj::StringPtr sqlCode,
-        kj::Array<BindingValue> bindings);
+        kj::Array<BindingValue> bindings,
+        kj::Maybe<SqliteQueryId> queryId);
 
-    State(kj::Rc<CachedStatement> cachedStatement, kj::Array<BindingValue> bindings);
+    State(kj::Rc<CachedStatement> cachedStatement,
+        kj::Array<BindingValue> bindings,
+        kj::Maybe<SqliteQueryId> queryId);
   };
 
   // Nulled out when query is done or canceled.

src/workerd/util/BUILD.bazel

@@ -244,6 +244,7 @@ wd_cc_library(
     ],
     implementation_deps = [
         ":autogate",
+        ":entropy",
         ":strings",
         "//src/workerd/jsg:exception",
         "@sqlite3",

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

@@ -95,37 +95,43 @@ kj::Own<SqliteKv::ListCursor> SqliteKv::list(
   if (!tableCreated) return kj::heap<ListCursor>(nullptr);
   auto& stmts = KJ_UNWRAP_OR(state.tryGet<Initialized>(), return kj::heap<ListCursor>(nullptr));
 
+  // Mint one id for this list so every statement it runs shares it.
+  SqliteDatabase::Statement::StatementOptions runOptions{.queryId = db.generateQueryId()};
+
   if (order == Order::FORWARD) {
     KJ_IF_SOME(e, end) {
       KJ_IF_SOME(l, limit) {
-        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListEndLimit, begin, e,
-            static_cast<int64_t>(l));
+        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListEndLimit,
+            runOptions, begin, e, static_cast<int64_t>(l));
       } else {
-        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListEnd, begin, e);
+        return kj::heap<ListCursor>(
+            kj::Badge<SqliteKv>(), *this, stmts.stmtListEnd, runOptions, begin, e);
       }
     } else {
       KJ_IF_SOME(l, limit) {
-        return kj::heap<ListCursor>(
-            kj::Badge<SqliteKv>(), *this, stmts.stmtListLimit, begin, static_cast<int64_t>(l));
+        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListLimit, runOptions,
+            begin, static_cast<int64_t>(l));
       } else {
-        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtList, begin);
+        return kj::heap<ListCursor>(
+            kj::Badge<SqliteKv>(), *this, stmts.stmtList, runOptions, begin);
       }
     }
   } else {
     KJ_IF_SOME(e, end) {
       KJ_IF_SOME(l, limit) {
         return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListEndLimitReverse,
-            begin, e, static_cast<int64_t>(l));
+            runOptions, begin, e, static_cast<int64_t>(l));
       } else {
         return kj::heap<ListCursor>(
-            kj::Badge<SqliteKv>(), *this, stmts.stmtListEndReverse, begin, e);
+            kj::Badge<SqliteKv>(), *this, stmts.stmtListEndReverse, runOptions, begin, e);
       }
     } else {
       KJ_IF_SOME(l, limit) {
-        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListLimitReverse, begin,
-            static_cast<int64_t>(l));
+        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListLimitReverse,
+            runOptions, begin, static_cast<int64_t>(l));
       } else {
-        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListReverse, begin);
+        return kj::heap<ListCursor>(
+            kj::Badge<SqliteKv>(), *this, stmts.stmtListReverse, runOptions, begin);
       }
     }
   }
@@ -151,28 +157,31 @@ void SqliteKv::put(KeyPtr key, ValuePtr value) {
 }
 
 void SqliteKv::put(KeyPtr key, ValuePtr value, WriteOptions options) {
+  auto queryId = resolveQueryId(options.queryId);
   ensureInitialized(options.allowUnconfirmed)
-      .stmtPut.run({.allowUnconfirmed = options.allowUnconfirmed}, key, value);
-  clearExternalsIfPresent(key);
+      .stmtPut.run({.allowUnconfirmed = options.allowUnconfirmed, .queryId = queryId}, key, value);
+  clearExternalsIfPresent(key, queryId);
 }
 
 bool SqliteKv::delete_(KeyPtr key) {
   return delete_(key, {});
 }
 
 bool SqliteKv::delete_(KeyPtr key, WriteOptions options) {
+  auto queryId = resolveQueryId(options.queryId);
   bool result;
   {
     auto query = ensureInitialized(options.allowUnconfirmed)
-                     .stmtDelete.run({.allowUnconfirmed = options.allowUnconfirmed}, key);
+                     .stmtDelete.run(
+                         {.allowUnconfirmed = options.allowUnconfirmed, .queryId = queryId}, key);
     result = query.changeCount() > 0;
   }
 
-  clearExternalsIfPresent(key);
+  clearExternalsIfPresent(key, queryId);
   return result;
 }
 
-void SqliteKv::clearExternalsIfPresent(KeyPtr key) {
+void SqliteKv::clearExternalsIfPresent(KeyPtr key, SqliteQueryId queryId) {
   // If the externals table hasn't been created yet, there's nothing to clear. We deliberately
   // avoid creating it here -- it should only be created when externals are actually being
   // written.
@@ -184,14 +193,19 @@ void SqliteKv::clearExternalsIfPresent(KeyPtr key) {
 
   // Externals writes are always paired with a regular KV write, and that paired write decides
   // whether the operation is allowed to be unconfirmed. So we always pass allowUnconfirmed here.
-  stmts.stmtDeleteExternals.run({.allowUnconfirmed = true}, key);
+  stmts.stmtDeleteExternals.run({.allowUnconfirmed = true, .queryId = queryId}, key);
 }
 
 uint SqliteKv::deleteAll() {
   // TODO(perf): Consider introducing a compatibility flag that causes deleteAll() to always return
   //   1. Apps almost certainly don't care about the return value but historically we returned the
   //   count of keys deleted, so now we're stuck counting the table size for no good reason.
-  uint count = tableCreated ? ensureInitialized(false).stmtCountKeys.run().getInt(0) : 0;
+  uint count = tableCreated
+      ? ensureInitialized(false)
+            .stmtCountKeys
+            .run(SqliteDatabase::Statement::StatementOptions{.queryId = db.generateQueryId()})
+            .getInt(0)
+      : 0;
   db.reset();
   return count;
 }
@@ -233,7 +247,8 @@ kj::Array<kj::Array<byte>> SqliteKv::getExternals(kj::StringPtr key) {
   if (!externalsTableCreated) return nullptr;
   auto& stmts = KJ_UNWRAP_OR(externalsState.tryGet<ExternalsInitialized>(), return nullptr);
 
-  auto query = stmts.stmtGetExternals.run(key);
+  auto query = stmts.stmtGetExternals.run(
+      SqliteDatabase::Statement::StatementOptions{.queryId = db.generateQueryId()}, key);
 
   kj::Vector<kj::Array<byte>> result;
   while (!query.isDone()) {
@@ -249,12 +264,15 @@ void SqliteKv::putExternals(kj::StringPtr key, kj::Array<kj::Array<byte>> tokens
   // Externals writes are always paired with a regular KV write, and that paired write decides
   // whether the operation is allowed to be unconfirmed. So we always pass allowUnconfirmed here.
 
+  // Group every statement of this putExternals under one id.
+  auto queryId = db.generateQueryId();
+
   // Replace any existing tokens for this key with the new set.
-  stmts.stmtDeleteExternals.run({.allowUnconfirmed = true}, key);
+  stmts.stmtDeleteExternals.run({.allowUnconfirmed = true, .queryId = queryId}, key);
 
   for (auto i: kj::indices(tokens)) {
-    stmts.stmtPutExternal.run(
-        {.allowUnconfirmed = true}, key, static_cast<int64_t>(i), tokens[i].asPtr());
+    stmts.stmtPutExternal.run({.allowUnconfirmed = true, .queryId = queryId}, key,
+        static_cast<int64_t>(i), tokens[i].asPtr());
   }
 }
 
@@ -267,13 +285,14 @@ void SqliteKv::beforeSqliteReset() {
 void SqliteKv::rollbackMultiPut(Initialized& stmts, WriteOptions options) {
   KJ_IF_SOME(e, kj::runCatchingExceptions([&]() {
     // This should be rare, so we don't prepare a statement for it.
-    stmts.db.run(
-        {.regulator = Initialized::regulator, .allowUnconfirmed = options.allowUnconfirmed},
+    stmts.db.run({.regulator = Initialized::regulator,
+                   .allowUnconfirmed = options.allowUnconfirmed,
+                   .queryId = options.queryId},
         kj::str("ROLLBACK TO _cf_put_multiple_savepoint"));
-    stmts.stmtMultiPutRelease.run({.allowUnconfirmed = options.allowUnconfirmed});
+    stmts.stmtMultiPutRelease.run(
+        {.allowUnconfirmed = options.allowUnconfirmed, .queryId = options.queryId});
   })) {
     KJ_LOG(WARNING, "silencing exception encountered while rolling back multi-put", e);
   }
 }
-
 }  // namespace workerd

src/workerd/util/sqlite-kv.h

@@ -59,6 +59,10 @@ class SqliteKv: private SqliteDatabase::ResetListener {
 
   struct WriteOptions {
     bool allowUnconfirmed = false;
+
+    // Groups all statements of one top-level KV operation under a single id. Left null by callers;
+    // minted by the operation and threaded into any statements/sub-operations it spawns.
+    kj::Maybe<SqliteQueryId> queryId = kj::none;
   };
 
   // Store a value into the table.
@@ -220,10 +224,20 @@ class SqliteKv: private SqliteDatabase::ResetListener {
   // Delete all externals associated with `key`, if the _cf_EXTERNALS table exists. No-op if the
   // table has never been created. Called automatically from put() and delete_(). Always runs
   // with `allowUnconfirmed = true` since the paired KV write decides confirmation semantics.
-  void clearExternalsIfPresent(KeyPtr key);
+  void clearExternalsIfPresent(KeyPtr key, SqliteQueryId queryId);
 
   void beforeSqliteReset() override;
 
+  // Returns `queryId` if set, otherwise mints a fresh one. Used so every statement of a single
+  // top-level KV operation shares one grouping id, while recursive operations (e.g. multi-put)
+  // reuse the id minted by the outer operation.
+  SqliteQueryId resolveQueryId(kj::Maybe<SqliteQueryId> queryId) {
+    KJ_IF_SOME(id, queryId) {
+      return id;
+    }
+    return db.generateQueryId();
+  }
+
   // Helper function that rolls back a multi-put statement and swallows any exceptions that may
   // occur during the rollback.
   void rollbackMultiPut(Initialized& stmts, WriteOptions options);
@@ -300,7 +314,8 @@ bool SqliteKv::get(KeyPtr key, Func&& callback) {
   if (!tableCreated) return false;
   auto& stmts = KJ_UNWRAP_OR(state.tryGet<Initialized>(), return false);
 
-  auto query = stmts.stmtGet.run(key);
+  auto query = stmts.stmtGet.run(
+      SqliteDatabase::Statement::StatementOptions{.queryId = db.generateQueryId()}, key);
 
   if (query.isDone()) {
     return false;
@@ -320,18 +335,24 @@ template <typename ArrayOfKeyValuePair>
 void SqliteKv::put(ArrayOfKeyValuePair& pairs, WriteOptions options) {
   // TODO(cleanup): This code is very similar to DurableObjectStorage::transactionSync.  Perhaps the
   // general structure can be shared somehow?
+  // Mint one id for the whole multi-put so every statement it issues (savepoint, each put, and
+  // release) shares it.
+  options.queryId = resolveQueryId(options.queryId);
   auto& stmts = ensureInitialized(options.allowUnconfirmed);
-  stmts.stmtMultiPutSavepoint.run({.allowUnconfirmed = options.allowUnconfirmed});
+  stmts.stmtMultiPutSavepoint.run(
+      {.allowUnconfirmed = options.allowUnconfirmed, .queryId = options.queryId});
 
   {
     // If any of the puts throw an exception, rollback the transaction and re-throw the exception
     // from the put that failed.
     KJ_ON_SCOPE_FAILURE(rollbackMultiPut(stmts, options));
     for (const auto& pair: pairs) {
-      put(pair.key, pair.value, {.allowUnconfirmed = options.allowUnconfirmed});
+      put(pair.key, pair.value,
+          {.allowUnconfirmed = options.allowUnconfirmed, .queryId = options.queryId});
     }
   }
-  stmts.stmtMultiPutRelease.run({.allowUnconfirmed = options.allowUnconfirmed});
+  stmts.stmtMultiPutRelease.run(
+      {.allowUnconfirmed = options.allowUnconfirmed, .queryId = options.queryId});
 }
 
 }  // namespace workerd