Merge pull request #38421 from gottesmm/track-task-is-running-with-exclusivity-runtime-fixes

[exclusivity] Add support for Task local exclusivity access sets.

Author: gottesmm

include/swift/ABI/Task.h

@@ -279,6 +279,10 @@ class AsyncTask : public Job {
   void flagAsSuspended();
   void flagAsSuspended_slow();
 
+  /// Flag that this task is now completed. This normally does not do anything
+  /// but can be used to locally insert logging.
+  void flagAsCompleted();
+
   /// Check whether this task has been cancelled.
   /// Checking this is, of course, inherently race-prone on its own.
   bool isCancelled() const;

include/swift/Runtime/Exclusivity.h

@@ -91,6 +91,30 @@ void swift_dumpTrackedAccesses();
 
 #endif
 
+/// Called when a task inits, resumes and returns control to caller synchronous
+/// code to update any exclusivity specific state associated with the task.
+///
+/// State is assumed to point to a buffer of memory with
+/// swift_task_threadLocalContextSize bytes that was initialized with
+/// swift_task_initThreadLocalContext.
+///
+/// We describe the algorithm in detail on SwiftTaskThreadLocalContext in
+/// Exclusivity.cpp.
+SWIFT_RUNTIME_EXPORT
+void swift_task_enterThreadLocalContext(char *state);
+
+/// Called when a task suspends and returns control to caller synchronous code
+/// to update any exclusivity specific state associated with the task.
+///
+/// State is assumed to point to a buffer of memory with
+/// swift_task_threadLocalContextSize bytes that was initialized with
+/// swift_task_initThreadLocalContext.
+///
+/// We describe the algorithm in detail on SwiftTaskThreadLocalContext in
+/// Exclusivity.cpp.
+SWIFT_RUNTIME_EXPORT
+void swift_task_exitThreadLocalContext(char *state);
+
 } // end namespace swift
 
 #endif

stdlib/public/Concurrency/PartialAsyncTask.swift

@@ -13,11 +13,23 @@
 import Swift
 @_implementationOnly import _SwiftConcurrencyShims
 
+@available(SwiftStdlib 5.5, *)
+@_silgen_name("swift_job_run")
+@usableFromInline
+internal func _swiftJobRun(_ job: UnownedJob,
+                           _ executor: UnownedSerialExecutor) -> ()
+
 /// A job is a unit of scheduleable work.
 @available(SwiftStdlib 5.5, *)
 @frozen
 public struct UnownedJob {
   private var context: Builtin.Job
+
+  @_alwaysEmitIntoClient
+  @inlinable
+  public func _runSynchronously(on executor: UnownedSerialExecutor) {
+      _swiftJobRun(self, executor)
+  }
 }
 
 @available(SwiftStdlib 5.5, *)

stdlib/public/Concurrency/Task.cpp

@@ -225,6 +225,8 @@ static void destroyJob(SWIFT_CONTEXT HeapObject *obj) {
 }
 
 AsyncTask::~AsyncTask() {
+  flagAsCompleted();
+
   // For a future, destroy the result.
   if (isFuture()) {
     futureFragment()->destroy();

stdlib/public/Concurrency/TaskPrivate.h

@@ -17,13 +17,14 @@
 #ifndef SWIFT_CONCURRENCY_TASKPRIVATE_H
 #define SWIFT_CONCURRENCY_TASKPRIVATE_H
 
-#include "swift/Runtime/Concurrency.h"
-#include "swift/ABI/Task.h"
+#include "Error.h"
 #include "swift/ABI/Metadata.h"
+#include "swift/ABI/Task.h"
 #include "swift/Runtime/Atomic.h"
-#include "swift/Runtime/HeapObject.h"
+#include "swift/Runtime/Concurrency.h"
 #include "swift/Runtime/Error.h"
-#include "Error.h"
+#include "swift/Runtime/Exclusivity.h"
+#include "swift/Runtime/HeapObject.h"
 
 #define SWIFT_FATAL_ERROR swift_Concurrency_fatalError
 #include "../runtime/StackAllocator.h"
@@ -282,14 +283,19 @@ struct AsyncTask::PrivateStorage {
   /// Currently one word.
   TaskLocal::Storage Local;
 
+  /// State inside the AsyncTask whose state is only managed by the exclusivity
+  /// runtime in stdlibCore. We zero initialize to provide a safe initial value,
+  /// but actually initialize its bit state to a const global provided by
+  /// libswiftCore so that libswiftCore can control the layout of our initial
+  /// state.
+  uintptr_t ExclusivityAccessSet[2] = {0, 0};
+
   PrivateStorage(JobFlags flags)
-    : Status(ActiveTaskStatus(flags)),
-      Local(TaskLocal::Storage()) {}
+      : Status(ActiveTaskStatus(flags)), Local(TaskLocal::Storage()) {}
 
   PrivateStorage(JobFlags flags, void *slab, size_t slabCapacity)
-    : Status(ActiveTaskStatus(flags)),
-      Allocator(slab, slabCapacity),
-      Local(TaskLocal::Storage()) {}
+      : Status(ActiveTaskStatus(flags)), Allocator(slab, slabCapacity),
+        Local(TaskLocal::Storage()) {}
 
   void complete(AsyncTask *task) {
     // Destroy and deallocate any remaining task local items.
@@ -347,7 +353,10 @@ inline void AsyncTask::flagAsRunning() {
   while (true) {
     assert(!oldStatus.isRunning());
     if (oldStatus.isLocked()) {
-      return flagAsRunning_slow();
+      flagAsRunning_slow();
+      swift_task_enterThreadLocalContext(
+          (char *)&_private().ExclusivityAccessSet[0]);
+      return;
     }
 
     auto newStatus = oldStatus.withRunning(true);
@@ -358,8 +367,11 @@ inline void AsyncTask::flagAsRunning() {
 
     if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
                                                 std::memory_order_relaxed,
-                                                std::memory_order_relaxed))
+                                                std::memory_order_relaxed)) {
+      swift_task_enterThreadLocalContext(
+          (char *)&_private().ExclusivityAccessSet[0]);
       return;
+    }
   }
 }
 
@@ -368,7 +380,10 @@ inline void AsyncTask::flagAsSuspended() {
   while (true) {
     assert(oldStatus.isRunning());
     if (oldStatus.isLocked()) {
-      return flagAsSuspended_slow();
+      flagAsSuspended_slow();
+      swift_task_exitThreadLocalContext(
+          (char *)&_private().ExclusivityAccessSet[0]);
+      return;
     }
 
     auto newStatus = oldStatus.withRunning(false);
@@ -379,11 +394,24 @@ inline void AsyncTask::flagAsSuspended() {
 
     if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
                                                 std::memory_order_relaxed,
-                                                std::memory_order_relaxed))
+                                                std::memory_order_relaxed)) {
+      swift_task_exitThreadLocalContext(
+          (char *)&_private().ExclusivityAccessSet[0]);
       return;
+    }
   }
 }
 
+// READ ME: This is not a dead function! Do not remove it! This is a function
+// that can be used when debugging locally to instrument when a task actually is
+// dealloced.
+inline void AsyncTask::flagAsCompleted() {
+#if SWIFT_TASK_PRINTF_DEBUG
+  fprintf(stderr, "[%lu] task completed %p\n",
+          _swift_get_thread_id(), this);
+#endif
+}
+
 inline void AsyncTask::localValuePush(const HeapObject *key,
                                       /* +1 */ OpaqueValue *value,
                                       const Metadata *valueType) {

stdlib/public/runtime/EnvironmentVariables.def

@@ -62,4 +62,11 @@ VARIABLE(SWIFT_DEBUG_ENABLE_SHARED_CACHE_PROTOCOL_CONFORMANCES, bool, true,
 
 #endif
 
+#ifndef NDEBUG
+
+VARIABLE(SWIFT_DEBUG_RUNTIME_EXCLUSIVITY_LOGGING, bool, false,
+         "Enable the an asserts runtime to emit logging as it works.")
+
+#endif
+
 #undef VARIABLE