Pass Spans to OutputGate lockWhile

Author: shrima-cf

src/workerd/api/actor-state.c++

@@ -499,8 +499,9 @@ jsg::Promise<void> DurableObjectStorageOperations::setAlarm(
   // they want immediate execution.
   kj::Date dateNowKjDate = static_cast<int64_t>(dateNow()) * kj::MILLISECONDS + kj::UNIX_EPOCH;
 
-  auto maybeBackpressure = transformMaybeBackpressure(
-      js, options, getCache(OP_PUT_ALARM).setAlarm(kj::max(scheduledTime, dateNowKjDate), options));
+  auto maybeBackpressure = transformMaybeBackpressure(js, options,
+      getCache(OP_PUT_ALARM)
+          .setAlarm(kj::max(scheduledTime, dateNowKjDate), options, context.getCurrentTraceSpan()));
 
   // setAlarm() is billed as a single write unit.
   context.addTask(updateStorageWriteUnit(context, currentActorMetrics(), 1));
@@ -515,10 +516,11 @@ jsg::Promise<void> DurableObjectStorageOperations::putOne(
 
   auto units = billingUnits(key.size() + buffer.size());
 
-  jsg::Promise<void> maybeBackpressure = transformMaybeBackpressure(
-      js, options, getCache(OP_PUT).put(kj::mv(key), kj::mv(buffer), options));
-
   auto& context = IoContext::current();
+
+  jsg::Promise<void> maybeBackpressure = transformMaybeBackpressure(js, options,
+      getCache(OP_PUT).put(kj::mv(key), kj::mv(buffer), options, context.getCurrentTraceSpan()));
+
   context.addTask(updateStorageWriteUnit(context, currentActorMetrics(), units));
   return maybeBackpressure;
 }
@@ -555,8 +557,8 @@ jsg::Promise<void> DurableObjectStorageOperations::deleteAlarm(
   }).orDefault(PutOptions{}));
 
   return context.attachSpans(js,
-      transformMaybeBackpressure(
-          js, options, getCache(OP_DELETE_ALARM).setAlarm(kj::none, options)),
+      transformMaybeBackpressure(js, options,
+          getCache(OP_DELETE_ALARM).setAlarm(kj::none, options, context.getCurrentTraceSpan())),
       kj::mv(userSpan));
 }
 
@@ -566,7 +568,7 @@ jsg::Promise<void> DurableObjectStorage::deleteAll(
   auto userSpan = context.makeUserTraceSpan("durable_object_storage_deleteAll"_kjc);
   auto options = configureOptions(kj::mv(maybeOptions).orDefault(PutOptions{}));
 
-  auto deleteAll = cache->deleteAll(options);
+  auto deleteAll = cache->deleteAll(options, context.getCurrentTraceSpan());
 
   context.addTask(updateStorageDeletes(context, currentActorMetrics(), kj::mv(deleteAll.count)));
 
@@ -580,8 +582,11 @@ void DurableObjectTransaction::deleteAll() {
 
 jsg::Promise<bool> DurableObjectStorageOperations::deleteOne(
     jsg::Lock& js, kj::String key, const PutOptions& options) {
-  return transformCacheResult(
-      js, getCache(OP_DELETE).delete_(kj::mv(key), options), options, [](jsg::Lock&, bool value) {
+  auto& context = IoContext::current();
+
+  return transformCacheResult(js,
+      getCache(OP_DELETE).delete_(kj::mv(key), options, context.getCurrentTraceSpan()), options,
+      [](jsg::Lock&, bool value) {
     currentActorMetrics().addStorageDeletes(1);
     return value;
   });
@@ -613,10 +618,11 @@ jsg::Promise<void> DurableObjectStorageOperations::putMultiple(
     kvs.add(ActorCacheOps::KeyValuePair{kj::mv(field.name), kj::mv(buffer)});
   }
 
-  jsg::Promise<void> maybeBackpressure =
-      transformMaybeBackpressure(js, options, getCache(OP_PUT).put(kvs.releaseAsArray(), options));
-
   auto& context = IoContext::current();
+
+  jsg::Promise<void> maybeBackpressure = transformMaybeBackpressure(js, options,
+      getCache(OP_PUT).put(kvs.releaseAsArray(), options, context.getCurrentTraceSpan()));
+
   context.addTask(updateStorageWriteUnit(context, currentActorMetrics(), units));
 
   return maybeBackpressure;
@@ -626,7 +632,10 @@ jsg::Promise<int> DurableObjectStorageOperations::deleteMultiple(
     jsg::Lock& js, kj::Array<kj::String> keys, const PutOptions& options) {
   auto numKeys = keys.size();
 
-  return transformCacheResult(js, getCache(OP_DELETE).delete_(kj::mv(keys), options), options,
+  auto& context = IoContext::current();
+
+  return transformCacheResult(js,
+      getCache(OP_DELETE).delete_(kj::mv(keys), options, context.getCurrentTraceSpan()), options,
       [numKeys](jsg::Lock&, uint count) -> int {
     currentActorMetrics().addStorageDeletes(numKeys);
     return count;

src/workerd/api/global-scope.c++

@@ -438,7 +438,8 @@ kj::Promise<WorkerInterface::AlarmResult> ServiceWorkerGlobalScope::runAlarm(kj:
     }
   }
 
-  KJ_SWITCH_ONEOF(persistent.armAlarmHandler(scheduledTime, false, actorId)) {
+  KJ_SWITCH_ONEOF(persistent.armAlarmHandler(
+                      scheduledTime, context.getCurrentTraceSpan(), false, actorId)) {
     KJ_CASE_ONEOF(armResult, ActorCacheInterface::RunAlarmHandler) {
       auto& handler = KJ_REQUIRE_NONNULL(exportedHandler);
       if (handler.alarm == kj::none) {

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

@@ -179,22 +179,22 @@ struct ActorCacheConvenienceWrappers {
   }
 
   auto put(kj::StringPtr key, kj::StringPtr value, ActorCache::WriteOptions options = {}) {
-    return target.put(kj::str(key), kj::heapArray(value.asBytes()), options);
+    return target.put(kj::str(key), kj::heapArray(value.asBytes()), options, nullptr);
   }
   auto put(kj::ArrayPtr<const KeyValuePtr> kvs, ActorCache::WriteOptions options = {}) {
     return target.put(KJ_MAP(kv, kvs) {
       return ActorCache::KeyValuePair{kj::str(kv.key), kj::heapArray(kv.value.asBytes())};
-    }, options);
+    }, options, nullptr);
   }
   auto setAlarm(kj::Maybe<kj::Date> newTime, ActorCache::WriteOptions options = {}) {
-    return target.setAlarm(newTime, options);
+    return target.setAlarm(newTime, options, nullptr);
   }
 
   auto delete_(kj::StringPtr key, ActorCache::WriteOptions options = {}) {
-    return target.delete_(kj::str(key), options);
+    return target.delete_(kj::str(key), options, nullptr);
   }
   auto delete_(kj::ArrayPtr<const kj::StringPtr> keys, ActorCache::WriteOptions options = {}) {
-    return target.delete_(KJ_MAP(k, keys) { return kj::str(k); }, options);
+    return target.delete_(KJ_MAP(k, keys) { return kj::str(k); }, options, nullptr);
   }
 
  private:
@@ -700,7 +700,7 @@ KJ_TEST("ActorCache batching due to max storage RPC words") {
   for (int i = 0; i < 128; ++i) {
     test.cache.put(kj::str(i),
         kj::Array<const byte>(bigVal.begin(), bigVal.size(), kj::NullArrayDisposer::instance),
-        ActorCache::WriteOptions());
+        ActorCache::WriteOptions(), nullptr);
   }
 
   auto mockTxn = mockStorage->expectCall("txn", ws).returnMock("transaction");
@@ -736,7 +736,7 @@ KJ_TEST("ActorCache deleteAll()") {
   test.put("baz", "789");  // overwrites a counted delete
   test.delete_("garply");  // uncounted delete
 
-  auto deleteAll = test.cache.deleteAll({});
+  auto deleteAll = test.cache.deleteAll({}, nullptr);
 
   // Post-deleteAll writes.
   test.put("grault", "12345");
@@ -819,7 +819,7 @@ KJ_TEST("ActorCache deleteAll() during transaction commit") {
 
     // Issue a put and a deleteAll() here!
     test.put("bar", "456");
-    test.cache.deleteAll({});
+    test.cache.deleteAll({}, nullptr);
 
     kj::mv(inProgressFlush).thenReturn(CAPNP());
   }
@@ -861,7 +861,7 @@ KJ_TEST("ActorCache deleteAll() again when previous one isn't done yet") {
   test.put("baz", "789");  // overwrites a counted delete
   test.delete_("garply");  // uncounted delete
 
-  auto deleteAllA = test.cache.deleteAll({});
+  auto deleteAllA = test.cache.deleteAll({}, nullptr);
 
   // Post-deleteAll writes.
   test.put("grault", "12345");
@@ -884,7 +884,7 @@ KJ_TEST("ActorCache deleteAll() again when previous one isn't done yet") {
   }
 
   // Do another deleteAll() before the first one is done.
-  auto deleteAllB = test.cache.deleteAll({});
+  auto deleteAllB = test.cache.deleteAll({}, nullptr);
 
   // And a write after that.
   test.put("fred", "2323");
@@ -4979,7 +4979,7 @@ KJ_TEST("ActorCache alarm get/put") {
 
   {
     // we have a cached time == nullptr, so we should not attempt to run an alarm
-    auto armResult = test.cache.armAlarmHandler(10 * kj::SECONDS + kj::UNIX_EPOCH, false);
+    auto armResult = test.cache.armAlarmHandler(10 * kj::SECONDS + kj::UNIX_EPOCH, nullptr);
     KJ_ASSERT(armResult.is<ActorCache::CancelAlarmHandler>());
     auto cancelResult = kj::mv(armResult.get<ActorCache::CancelAlarmHandler>());
     KJ_ASSERT(cancelResult.waitBeforeCancel.poll(ws));
@@ -4997,7 +4997,7 @@ KJ_TEST("ActorCache alarm get/put") {
   {
     // Test that alarm handler handle clears alarm when dropped with no writes
     {
-      auto armResult = test.cache.armAlarmHandler(oneMs, false);
+      auto armResult = test.cache.armAlarmHandler(oneMs, nullptr);
       KJ_ASSERT(armResult.is<ActorCache::RunAlarmHandler>());
     }
     mockStorage->expectCall("deleteAlarm", ws)
@@ -5010,7 +5010,7 @@ KJ_TEST("ActorCache alarm get/put") {
 
     // Test that alarm handler handle does not clear alarm when dropped with writes
     {
-      auto armResult = test.cache.armAlarmHandler(oneMs, false);
+      auto armResult = test.cache.armAlarmHandler(oneMs, nullptr);
       KJ_ASSERT(armResult.is<ActorCache::RunAlarmHandler>());
       test.setAlarm(twoMs);
     }
@@ -5024,7 +5024,7 @@ KJ_TEST("ActorCache alarm get/put") {
 
     // Test that alarm handler handle does not cache delete when it fails
     {
-      auto armResult = test.cache.armAlarmHandler(oneMs, false);
+      auto armResult = test.cache.armAlarmHandler(oneMs, nullptr);
       KJ_ASSERT(armResult.is<ActorCache::RunAlarmHandler>());
     }
     mockStorage->expectCall("deleteAlarm", ws)
@@ -5036,7 +5036,7 @@ KJ_TEST("ActorCache alarm get/put") {
   {
     // Test that alarm handler handle does not cache alarm delete when noCache == true
     {
-      auto armResult = test.cache.armAlarmHandler(twoMs, true);
+      auto armResult = test.cache.armAlarmHandler(twoMs, nullptr, true);
       KJ_ASSERT(armResult.is<ActorCache::RunAlarmHandler>());
     }
     mockStorage->expectCall("deleteAlarm", ws)
@@ -5090,7 +5090,7 @@ KJ_TEST("ActorCache alarm delete when flush fails") {
   // we want to test that even if a flush is retried
   // that the post-delete actions for a checked delete happen.
   {
-    auto handle = test.cache.armAlarmHandler(oneMs, false);
+    auto handle = test.cache.armAlarmHandler(oneMs, nullptr);
 
     auto time = expectCached(test.getAlarm());
     KJ_ASSERT(time == kj::none);
@@ -5283,7 +5283,7 @@ KJ_TEST("ActorCache can wait for flush") {
 
   {
     // Join in on a scheduled deleteAll.
-    test.cache.deleteAll(ActorCacheWriteOptions{.allowUnconfirmed = false});
+    test.cache.deleteAll(ActorCacheWriteOptions{.allowUnconfirmed = false}, nullptr);
 
     verify(
         [&]() {
@@ -5301,7 +5301,7 @@ KJ_TEST("ActorCache can wait for flush") {
 
   {
     // Join in on a scheduled deleteAll with allowUnconfirmed.
-    test.cache.deleteAll(ActorCacheWriteOptions{.allowUnconfirmed = true});
+    test.cache.deleteAll(ActorCacheWriteOptions{.allowUnconfirmed = true}, nullptr);
 
     verify(
         [&]() {