Plumb allowUnconfirmed through SQL queries to the onWrite callback

It’s possible to mix writes that are confirmed with writes that are
unconfirmed.  We need to be able to distinguish which writes have what
unconfirmed status in the `onWrite` call.  The simplest way I could
figure out how to do this was to plumb whether a certain query is
confirmed or not all the way from where the queries are created
down into the SqliteDatabase layer, which can call back into the
ActorSqlite layer via the `onWrite` call.

Notably, this allows us to mix APIs that do support unconfirmed
writes (so far, just the async K/V API) with APIs that support
confirmed writes only (the `storage.sql` and `storage.kv` APIs) since
we default allowConfirmed to be false.

Author: justin-mp

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

@@ -220,7 +220,7 @@ void ActorSqlite::onCriticalError(
   }
 }
 
-void ActorSqlite::onWrite() {
+void ActorSqlite::onWrite(bool allowUnconfirmed) {
   requireNotBroken();
   if (currentTxn.is<NoTxn>()) {
     auto txn = kj::heap<ImplicitTxn>(*this);
@@ -471,7 +471,7 @@ kj::OneOf<ActorCacheOps::GetResultList, kj::Promise<ActorCacheOps::GetResultList
 
 kj::Maybe<kj::Promise<void>> ActorSqlite::put(Key key, Value value, WriteOptions options) {
   requireNotBroken();
-  kv.put(key, value);
+  kv.put(key, value, {.allowUnconfirmed = options.allowUnconfirmed});
   return kj::none;
 }
 
@@ -491,7 +491,7 @@ kj::Maybe<kj::Promise<void>> ActorSqlite::put(kj::Array<KeyValuePair> pairs, Wri
           {.regulator = SqliteDatabase::TRUSTED}, kj::str("SAVEPOINT _cf_put_multiple_savepoint"));
       for (const auto& pair: pairs) {
         try {
-          kv.put(pair.key, pair.value);
+          kv.put(pair.key, pair.value, {.allowUnconfirmed = options.allowUnconfirmed});
         } catch (kj::Exception& e) {
           // We need to rollback to the putMultiple SAVEPOINT. Do it, and then release the SAVEPOINT.
           db->run({.regulator = SqliteDatabase::TRUSTED},
@@ -507,7 +507,7 @@ kj::Maybe<kj::Promise<void>> ActorSqlite::put(kj::Array<KeyValuePair> pairs, Wri
     }
   } else {
     for (auto& pair: pairs) {
-      kv.put(pair.key, pair.value);
+      kv.put(pair.key, pair.value, {.allowUnconfirmed = options.allowUnconfirmed});
     }
   }
   return kj::none;
@@ -516,15 +516,15 @@ kj::Maybe<kj::Promise<void>> ActorSqlite::put(kj::Array<KeyValuePair> pairs, Wri
 kj::OneOf<bool, kj::Promise<bool>> ActorSqlite::delete_(Key key, WriteOptions options) {
   requireNotBroken();
 
-  return kv.delete_(key);
+  return kv.delete_(key, {.allowUnconfirmed = options.allowUnconfirmed});
 }
 
 kj::OneOf<uint, kj::Promise<uint>> ActorSqlite::delete_(kj::Array<Key> keys, WriteOptions options) {
   requireNotBroken();
 
   uint count = 0;
   for (auto& key: keys) {
-    count += kv.delete_(key);
+    count += kv.delete_(key, {.allowUnconfirmed = options.allowUnconfirmed});
   }
   return count;
 }

src/workerd/io/actor-sqlite.h

@@ -227,7 +227,7 @@ class ActorSqlite final: public ActorCacheInterface, private kj::TaskSet::ErrorH
   AlarmLaterErrorHandler alarmLaterErrorHandler;
   kj::TaskSet alarmLaterTasks;
 
-  void onWrite();
+  void onWrite(bool allowUnconfirmed);
 
   void onCriticalError(kj::StringPtr errorMessage, kj::Maybe<kj::Exception> maybeException);
 

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

@@ -53,9 +53,11 @@ void SqliteKv::cancelCurrentCursor() {
   }
 }
 
-SqliteKv::Initialized& SqliteKv::ensureInitialized() {
+SqliteKv::Initialized& SqliteKv::ensureInitialized(bool allowUnconfirmed) {
   if (!tableCreated) {
-    db.run(R"(
+    db.run(SqliteDatabase::QueryOptions{.regulator = SqliteDatabase::TRUSTED,
+             .allowUnconfirmed = allowUnconfirmed},
+        R"(
       CREATE TABLE IF NOT EXISTS _cf_KV (
         key TEXT PRIMARY KEY,
         value BLOB
@@ -137,19 +139,29 @@ kj::Maybe<SqliteKv::ListCursor::KeyValuePair> SqliteKv::ListCursor::next() {
 }
 
 void SqliteKv::put(KeyPtr key, ValuePtr value) {
-  ensureInitialized().stmtPut.run(key, value);
+  return put(key, value, {});
+}
+
+void SqliteKv::put(KeyPtr key, ValuePtr value, WriteOptions options) {
+  ensureInitialized(options.allowUnconfirmed)
+      .stmtPut.run({.allowUnconfirmed = options.allowUnconfirmed}, key, value);
 }
 
 bool SqliteKv::delete_(KeyPtr key) {
-  auto query = ensureInitialized().stmtDelete.run(key);
+  return delete_(key, {});
+}
+
+bool SqliteKv::delete_(KeyPtr key, WriteOptions options) {
+  auto query = ensureInitialized(options.allowUnconfirmed)
+                   .stmtDelete.run({.allowUnconfirmed = options.allowUnconfirmed}, key);
   return query.changeCount() > 0;
 }
 
 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().stmtCountKeys.run().getInt(0) : 0;
+  uint count = tableCreated ? ensureInitialized(false).stmtCountKeys.run().getInt(0) : 0;
   db.reset();
   return count;
 }

src/workerd/util/sqlite-kv.h

@@ -54,11 +54,17 @@ class SqliteKv: private SqliteDatabase::ResetListener {
   class ListCursor;
   kj::Own<ListCursor> list(KeyPtr begin, kj::Maybe<KeyPtr> end, kj::Maybe<uint> limit, Order order);
 
+  struct WriteOptions {
+    bool allowUnconfirmed = false;
+  };
+
   // Store a value into the table.
   void put(KeyPtr key, ValuePtr value);
+  void put(KeyPtr key, ValuePtr value, WriteOptions options);
 
   // Delete the key and return whether it was matched.
   bool delete_(KeyPtr key);
+  bool delete_(KeyPtr key, WriteOptions options);
 
   uint deleteAll();
 
@@ -148,7 +154,7 @@ class SqliteKv: private SqliteDatabase::ResetListener {
 
   void cancelCurrentCursor();
 
-  Initialized& ensureInitialized();
+  Initialized& ensureInitialized(bool allowUnconfirmed);
   // Make sure the KV table is created and prepared statements are ready. Not called until the
   // first write.
 

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

@@ -449,7 +449,7 @@ KJ_TEST("SQLite onWrite callback") {
   SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY);
 
   bool sawWrite = false;
-  db.onWrite([&]() { sawWrite = true; });
+  db.onWrite([&](bool allowUnconfirmed) { sawWrite = true; });
 
   setupSql(db);
   KJ_EXPECT(sawWrite);

src/workerd/util/sqlite.c++

@@ -557,9 +557,9 @@ bool SqliteDatabase::observedCriticalError() {
   return criticalErrorOccurred;
 }
 
-void SqliteDatabase::notifyWrite() {
+void SqliteDatabase::notifyWrite(bool allowUnconfirmed) {
   KJ_IF_SOME(cb, onWriteCallback) {
-    cb();
+    cb(allowUnconfirmed);
   }
 }
 
@@ -766,7 +766,7 @@ SqliteDatabase::StatementAndEffect SqliteDatabase::prepareSql(const Regulator& r
               KJ_DEFER(currentRegulator = regulator);
               currentParseContext = kj::none;
               KJ_DEFER(currentParseContext = parseContext);
-              cb();
+              cb(false);  // prepareSql doesn't have access to allowUnconfirmed, use safe default
             }
           }
 
@@ -1423,7 +1423,7 @@ void SqliteDatabase::Query::checkRequirements(size_t size) {
 
   KJ_IF_SOME(cb, db.onWriteCallback) {
     if (!sqlite3_stmt_readonly(statement)) {
-      cb();
+      cb(allowUnconfirmed);
     }
   }
 }

src/workerd/util/sqlite.h

@@ -76,6 +76,7 @@ class SqliteDatabase {
 
   struct QueryOptions {
     const Regulator& regulator;
+    bool allowUnconfirmed = false;
   };
 
   struct IngestResult {
@@ -204,7 +205,7 @@ class SqliteDatabase {
   //
   // Note that the write callback is NOT called before (or at any point during) a reset(). Use the
   // `ResetListener` mechanism or `afterReset()` instead for that case.
-  void onWrite(kj::Function<void()> callback) {
+  void onWrite(kj::Function<void(bool allowUnconfirmed)> callback) {
     onWriteCallback = kj::mv(callback);
   }
 
@@ -229,7 +230,7 @@ class SqliteDatabase {
   // implement explicit transactions. For synchronous transactions, the explicit transaction needs
   // to be nested inside the automatic transaction, so we need to force an auto-transaction to
   // start before the SAVEPOINT.
-  void notifyWrite();
+  void notifyWrite(bool allowUnconfirmed = false);
 
   // Get the currently-executing SQL query for debug purposes. The query is normalized to hide
   // any literal values that might contain sensitive information. This is intended to be safe for
@@ -332,7 +333,7 @@ class SqliteDatabase {
   kj::Maybe<sqlite3_stmt&> currentStatement;
 
   bool criticalErrorOccurred = false;
-  kj::Maybe<kj::Function<void()>> onWriteCallback;
+  kj::Maybe<kj::Function<void(bool allowUnconfirmed)>> onWriteCallback;
   kj::Maybe<kj::Function<void(kj::StringPtr errorMessage, kj::Maybe<kj::Exception> maybeException)>>
       onCriticalErrorCallback;
   kj::Maybe<kj::Function<void(SqliteDatabase&)>> afterResetCallback;
@@ -446,6 +447,13 @@ class SqliteDatabase::Statement final: private ResetListener {
   template <typename... Params>
   Query run(Params&&... bindings);
 
+  struct StatementOptions {
+    bool allowUnconfirmed = false;
+  };
+
+  template <typename... Params>
+  Query run(StatementOptions options, Params&&... bindings);
+
  private:
   const Regulator& regulator;
   kj::OneOf<kj::String, StatementAndEffect> stmt;
@@ -632,6 +640,9 @@ class SqliteDatabase::Query final: private ResetListener {
   uint64_t rowsRead = 0;
   uint64_t rowsWritten = 0;
 
+  // Whether this query allows unconfirmed writes.
+  bool allowUnconfirmed = false;
+
   friend class SqliteDatabase;
 
   Query(SqliteDatabase& db,
@@ -647,7 +658,8 @@ class SqliteDatabase::Query final: private ResetListener {
       : ResetListener(db),
         regulator(options.regulator),
         maybeStatement(statement.prepareForExecution()),
-        queryEvent(this->db.sqliteObserver) {
+        queryEvent(this->db.sqliteObserver),
+        allowUnconfirmed(options.allowUnconfirmed) {
     // If we throw from the constructor, the destructor won't run. Need to call destroy()
     // explicitly.
     KJ_ON_SCOPE_FAILURE(destroy());
@@ -659,7 +671,8 @@ class SqliteDatabase::Query final: private ResetListener {
         regulator(options.regulator),
         ownStatement(db.prepareSql(regulator, sqlCode, 0, MULTI)),
         maybeStatement(ownStatement),
-        queryEvent(this->db.sqliteObserver) {
+        queryEvent(this->db.sqliteObserver),
+        allowUnconfirmed(options.allowUnconfirmed) {
     // If we throw from the constructor, the destructor won't run. Need to call destroy()
     // explicitly.
     KJ_ON_SCOPE_FAILURE(destroy());
@@ -966,6 +979,12 @@ SqliteDatabase::Query SqliteDatabase::Statement::run(Params&&... params) {
   return Query(db, QueryOptions{.regulator = regulator}, *this, kj::fwd<Params>(params)...);
 }
 
+template <typename... Params>
+SqliteDatabase::Query SqliteDatabase::Statement::run(StatementOptions options, Params&&... params) {
+  return Query(db, {.regulator = regulator, .allowUnconfirmed = options.allowUnconfirmed}, *this,
+      kj::fwd<Params>(params)...);
+}
+
 template <size_t size, typename... Params>
 SqliteDatabase::Query SqliteDatabase::run(const char (&sqlCode)[size], Params&&... params) {
   return Query(*this, QueryOptions{.regulator = TRUSTED}, sqlCode, kj::fwd<Params>(params)...);