ActorSqlite: Preserve alarm setting when running deleteAll()

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

@@ -738,5 +738,84 @@ KJ_TEST("getAlarm/setAlarm check for brokenness") {
   test.pollAndExpectCalls({});
 }
 
+KJ_TEST("calling deleteAll() preserves alarm state if alarm is set") {
+  ActorSqliteTest test;
+
+  // Initialize alarm state to 1ms.
+  test.setAlarm(oneMs);
+  test.pollAndExpectCalls({"scheduleRun(1ms)"})[0]->fulfill();
+  test.pollAndExpectCalls({"commit"})[0]->fulfill();
+  test.pollAndExpectCalls({});
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+
+  ActorCache::DeleteAllResults results = test.actor.deleteAll({});
+  KJ_ASSERT(results.backpressure == kj::none);
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+
+  auto commitFulfiller = kj::mv(test.pollAndExpectCalls({"commit"})[0]);
+  KJ_ASSERT(results.count.wait(test.ws) == 0);
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+
+  commitFulfiller->fulfill();
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+
+  test.pollAndExpectCalls({});
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+}
+
+KJ_TEST("calling deleteAll() preserves alarm state if alarm is not set") {
+  ActorSqliteTest test;
+
+  // Initialize alarm state to empty value in metadata table.
+  test.setAlarm(oneMs);
+  test.pollAndExpectCalls({"scheduleRun(1ms)"})[0]->fulfill();
+  test.pollAndExpectCalls({"commit"})[0]->fulfill();
+  test.pollAndExpectCalls({});
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+  test.setAlarm(kj::none);
+  test.pollAndExpectCalls({"commit"})[0]->fulfill();
+  test.pollAndExpectCalls({"scheduleRun(none)"})[0]->fulfill();
+  test.pollAndExpectCalls({});
+  KJ_ASSERT(expectSync(test.getAlarm()) == kj::none);
+
+  ActorCache::DeleteAllResults results = test.actor.deleteAll({});
+  KJ_ASSERT(results.backpressure == kj::none);
+  KJ_ASSERT(expectSync(test.getAlarm()) == kj::none);
+
+  auto commitFulfiller = kj::mv(test.pollAndExpectCalls({"commit"})[0]);
+  KJ_ASSERT(results.count.wait(test.ws) == 0);
+  KJ_ASSERT(expectSync(test.getAlarm()) == kj::none);
+
+  commitFulfiller->fulfill();
+  KJ_ASSERT(expectSync(test.getAlarm()) == kj::none);
+
+  // We can also assert that we leave the database empty, in case that turns out to be useful later:
+  auto q = test.db.run("SELECT name FROM sqlite_master WHERE type='table' AND name='_cf_METADATA'");
+  KJ_ASSERT(q.isDone());
+}
+
+KJ_TEST("calling deleteAll() during an implicit transaction preserves alarm state") {
+  ActorSqliteTest test;
+
+  // Initialize alarm state to 1ms.
+  test.setAlarm(oneMs);
+
+  ActorCache::DeleteAllResults results = test.actor.deleteAll({});
+  KJ_ASSERT(results.backpressure == kj::none);
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+
+  test.pollAndExpectCalls({"scheduleRun(1ms)"})[0]->fulfill();
+
+  auto commitFulfiller = kj::mv(test.pollAndExpectCalls({"commit"})[0]);
+  KJ_ASSERT(results.count.wait(test.ws) == 0);
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+
+  commitFulfiller->fulfill();
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+
+  test.pollAndExpectCalls({});
+  KJ_ASSERT(expectSync(test.getAlarm()) == oneMs);
+}
+
 }  // namespace
 }  // namespace workerd

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

@@ -423,6 +423,11 @@ ActorCacheInterface::DeleteAllResults ActorSqlite::deleteAll(WriteOptions option
   // TODO(soon): handle deleteAll and alarm state interaction.
   requireNotBroken();
 
+  // kv.deleteAll() clears the database, so we need to save and possibly restore alarm state in
+  // the metadata table, to try to match the behavior of ActorCache, which preserves the set alarm
+  // when running deleteAll().
+  auto localAlarmState = metadata.getAlarm();
+
   // deleteAll() cannot be part of a transaction because it deletes the database altogether. So,
   // we have to close our transactions or fail.
   KJ_SWITCH_ONEOF(currentTxn) {
@@ -451,8 +456,9 @@ ActorCacheInterface::DeleteAllResults ActorSqlite::deleteAll(WriteOptions option
     }
   }
 
-  if (!deleteAllCommitScheduled) {
-    // We'll want to make sure the commit callback is called for the deleteAll().
+  if (localAlarmState == kj::none && !deleteAllCommitScheduled) {
+    // If we're not going to perform a write to restore alarm state, we'll want to make sure the
+    // commit callback is called for the deleteAll().
     commitTasks.add(outputGate.lockWhile(kj::evalLater([this]() mutable -> kj::Promise<void> {
       // Don't commit if shutdown() has been called.
       requireNotBroken();
@@ -464,6 +470,21 @@ ActorCacheInterface::DeleteAllResults ActorSqlite::deleteAll(WriteOptions option
   }
 
   uint count = kv.deleteAll();
+
+  // TODO(correctness): Since workerd doesn't have a separate durability step, in the unlikely
+  // event of a failure here, between deleteAll() and setAlarm(), we could theoretically lose the
+  // current alarm state when running under workerd.  Not sure if there's a practical way to avoid
+  // this.
+
+  // Reinitialize metadata wrapper, to allow on-demand recreation of metadata table.
+  metadata = SqliteMetadata(*db);
+
+  // Reset alarm state, if necessary.  If no alarm is set, OK to just leave metadata table
+  // uninitialized.
+  if (localAlarmState != kj::none) {
+    metadata.setAlarm(localAlarmState);
+  }
+
   return {
     .backpressure = kj::none,
     .count = count,