[Concurrency] Add an environment variable to validate unchecked continuation usage.

When enabled, we track all active unchecked continuations in a global set, and fatal error if one is called twice.

rdar://97390481

Author: mikeash

include/swift/Runtime/EnvironmentVariables.h

@@ -40,6 +40,10 @@ extern swift::once_t initializeToken;
 // Concurrency library can call.
 SWIFT_RUNTIME_STDLIB_SPI bool concurrencyEnableJobDispatchIntegration();
 
+// Wrapper around SWIFT_DEBUG_VALIDATE_UNCHECKED_CONTINUATIONS that the
+// Concurrency library can call.
+SWIFT_RUNTIME_STDLIB_SPI bool concurrencyValidateUncheckedContinuations();
+
 } // end namespace environment
 } // end namespace runtime
 } // end namespace Swift

stdlib/public/Concurrency/Task.cpp

@@ -30,7 +30,9 @@
 #include "swift/ABI/Task.h"
 #include "swift/ABI/TaskLocal.h"
 #include "swift/ABI/TaskOptions.h"
+#include "swift/Basic/Lazy.h"
 #include "swift/Runtime/Concurrency.h"
+#include "swift/Runtime/EnvironmentVariables.h"
 #include "swift/Runtime/HeapObject.h"
 #include "swift/Threading/Mutex.h"
 #include <atomic>
@@ -1238,9 +1240,59 @@ swift_task_enqueueTaskOnExecutorImpl(AsyncTask *task, ExecutorRef executor)
   task->flagAsAndEnqueueOnExecutor(executor);
 }
 
+namespace continuationChecking {
+
+enum class State : uint8_t { Uninitialized, On, Off };
+
+static std::atomic<State> CurrentState;
+
+static LazyMutex ActiveContinuationsLock;
+static Lazy<std::unordered_set<ContinuationAsyncContext *>> ActiveContinuations;
+
+static bool isEnabled() {
+  auto state = CurrentState.load(std::memory_order_relaxed);
+  if (state == State::Uninitialized) {
+    bool enabled =
+        runtime::environment::concurrencyValidateUncheckedContinuations();
+    state = enabled ? State::On : State::Off;
+    CurrentState.store(state, std::memory_order_relaxed);
+  }
+  return state == State::On;
+}
+
+static void init(ContinuationAsyncContext *context) {
+  if (!isEnabled())
+    return;
+
+  LazyMutex::ScopedLock guard(ActiveContinuationsLock);
+  auto result = ActiveContinuations.get().insert(context);
+  auto inserted = std::get<1>(result);
+  if (!inserted)
+    swift_Concurrency_fatalError(
+        0,
+        "Initializing continuation context %p that was already initialized.\n",
+        context);
+}
+
+static void willResume(ContinuationAsyncContext *context) {
+  if (!isEnabled())
+    return;
+
+  LazyMutex::ScopedLock guard(ActiveContinuationsLock);
+  auto removed = ActiveContinuations.get().erase(context);
+  if (!removed)
+    swift_Concurrency_fatalError(0,
+                      "Resuming continuation context %p that was not awaited "
+                      "(may have already been resumed).\n",
+                      context);
+}
+
+} // namespace continuationChecking
+
 SWIFT_CC(swift)
 static AsyncTask *swift_continuation_initImpl(ContinuationAsyncContext *context,
                                               AsyncContinuationFlags flags) {
+  continuationChecking::init(context);
   context->Flags = ContinuationAsyncContext::FlagsType();
   if (flags.canThrow()) context->Flags.setCanThrow(true);
   if (flags.isExecutorSwitchForced())
@@ -1341,6 +1393,8 @@ static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
 
 static void resumeTaskAfterContinuation(AsyncTask *task,
                                         ContinuationAsyncContext *context) {
+  continuationChecking::willResume(context);
+
   auto &sync = context->AwaitSynchronization;
   auto status = sync.load(std::memory_order_acquire);
   assert(status != ContinuationStatus::Resumed &&

stdlib/public/Concurrency/TaskPrivate.h

@@ -40,8 +40,8 @@ namespace swift {
 // If this is enabled, tests with `swift_task_debug_log` requirement can run.
 #if 0
 #define SWIFT_TASK_DEBUG_LOG(fmt, ...)                                         \
-  fprintf(stderr, "[%#lx] [%s:%d](%s) " fmt "\n",                               \
-          (unsigned long)Thread::current()::platformThreadId(), __FILE__,      \
+  fprintf(stderr, "[%#lx] [%s:%d](%s) " fmt "\n",                              \
+          (unsigned long)Thread::current().platformThreadId(), __FILE__,       \
           __LINE__, __FUNCTION__, __VA_ARGS__)
 #else
 #define SWIFT_TASK_DEBUG_LOG(fmt, ...) (void)0

stdlib/public/runtime/EnvironmentVariables.cpp

@@ -246,3 +246,7 @@ SWIFT_RUNTIME_STDLIB_SPI bool concurrencyEnableJobDispatchIntegration() {
   return runtime::environment::
       SWIFT_ENABLE_ASYNC_JOB_DISPATCH_INTEGRATION();
 }
+
+SWIFT_RUNTIME_STDLIB_SPI bool concurrencyValidateUncheckedContinuations() {
+  return runtime::environment::SWIFT_DEBUG_VALIDATE_UNCHECKED_CONTINUATIONS();
+}

stdlib/public/runtime/EnvironmentVariables.def

@@ -50,6 +50,9 @@ VARIABLE(SWIFT_DEBUG_ENABLE_COW_CHECKS, bool, false,
 VARIABLE(SWIFT_ENABLE_ASYNC_JOB_DISPATCH_INTEGRATION, bool, true,
          "Enable use of dispatch_async_swift_job when available.")
 
+VARIABLE(SWIFT_DEBUG_VALIDATE_UNCHECKED_CONTINUATIONS, bool, false,
+         "Check for and error on double-calls of unchecked continuations.")
+
 #if defined(__APPLE__) && defined(__MACH__)
 
 VARIABLE(SWIFT_DEBUG_VALIDATE_SHARED_CACHE_PROTOCOL_CONFORMANCES, bool, false,

test/Concurrency/Runtime/continuation_validation.swift

@@ -0,0 +1,31 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -Xfrontend -disable-availability-checking -parse-as-library %s -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: env %env-SWIFT_DEBUG_VALIDATE_UNCHECKED_CONTINUATIONS=1 %target-run %t/a.out
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+// UNSUPPORTED: use_os_stdlib
+
+import StdlibUnittest
+
+@main struct Main {
+  static func main() async {
+    let tests = TestSuite("ContinuationValidation")
+
+    if #available(SwiftStdlib 5.1, *) {
+      tests.test("trap on double resume of unchecked continuation") {
+        expectCrashLater(withMessage: "may have already been resumed")
+
+        await withUnsafeContinuation { c in
+          c.resume(returning: ())
+          c.resume(returning: ())
+        }
+      }
+    }
+
+    await runAllTestsAsync()
+  }
+}