From 1937b2388638740a4e9909894cd14dbf83f8acc1 Mon Sep 17 00:00:00 2001 From: Joe Lee Date: Tue, 17 Sep 2024 12:23:25 -0700 Subject: [PATCH] ActorSqlite: Preserve alarm setting when running deleteAll() --- src/workerd/io/actor-sqlite-test.c++ | 79 ++++++++++++++++++++++++++++ src/workerd/io/actor-sqlite.c++ | 25 ++++++++- 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/src/workerd/io/actor-sqlite-test.c++ b/src/workerd/io/actor-sqlite-test.c++ index db38a3f7666..b96294afebc 100644 --- a/src/workerd/io/actor-sqlite-test.c++ +++ b/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 diff --git a/src/workerd/io/actor-sqlite.c++ b/src/workerd/io/actor-sqlite.c++ index 9f1dc46b191..a0b57fe153b 100644 --- a/src/workerd/io/actor-sqlite.c++ +++ b/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 { // 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,