[WIP, DNM] Track how much time a task spends scheduled.

include/swift/ABI/MetadataValues.h

@@ -2897,6 +2897,9 @@ enum class TaskStatusRecordKind : uint8_t {
   /// A human-readable task name.
   TaskName = 6,
 
+  // The total time the task has spent running.
+  TimeSpentRunning = 7,
+
   // Kinds >= 192 are private to the implementation.
   First_Reserved = 192,
   Private_RecordLock = 192

include/swift/ABI/Task.h

@@ -397,7 +397,14 @@ class AsyncTask : public Job {
   /// runInExecutorContext.
   SWIFT_CC(swiftasync)
   void runInFullyEstablishedContext() {
-    return ResumeTask(ResumeContext); // 'return' forces tail call
+    if (SWIFT_UNLIKELY(isTimeSpentRunningTracked())) {
+      auto begin = getNanosecondsOnSuspendingClock();
+      ResumeTask(ResumeContext);
+      auto end = getNanosecondsOnSuspendingClock();
+      ranForNanoseconds(end - begin);
+    } else {
+      return ResumeTask(ResumeContext); // 'return' forces tail call
+    }
   }
 
   /// A task can have the following states:
@@ -498,6 +505,28 @@ class AsyncTask : public Job {
   /// `swift_task_popTaskExecutorPreference(record)` method pair.
   void dropInitialTaskExecutorPreferenceRecord();
 
+  // ==== Tracking Time Spent Running ------------------------------------------
+
+  /// Whether or not the concurrency library is tracking the time spent running
+  /// tasks.
+  static inline bool isTimeSpentRunningTracked(void) {
+    return _isTimeSpentRunningTracked.load(std::memory_order_relaxed);
+  }
+
+  /// Set whether or not the concurrency library is tracking the time spent
+  /// running tasks. Returns the old value.
+  static inline bool setTimeSpentRunningTracked(bool isTracked) {
+    return _isTimeSpentRunningTracked.exchange(isTracked,
+                                               std::memory_order_relaxed);
+  }
+
+  /// Get the number of nanoseconds spent running this task so far, or `0` if
+  /// task duration tracking isn't enabled.
+  __attribute__((cold)) uint64_t getTimeSpentRunning(void);
+
+  void pushTimeSpentRunningRecord(void);
+  void popTimeSpentRunningRecord(void);
+
   // ==== Task Local Values ----------------------------------------------------
 
   void localValuePush(const HeapObject *key,
@@ -779,6 +808,17 @@ class AsyncTask : public Job {
     return reinterpret_cast<AsyncTask *&>(
         SchedulerPrivate[NextWaitingTaskIndex]);
   }
+
+  /// Whether or not the concurrency library is tracking the time spent running
+  /// tasks.
+  static std::atomic<bool> _isTimeSpentRunningTracked;
+
+  /// Record that the task spent an additional `ns` nanoseconds running.
+  void ranForNanoseconds(uint64_t ns);
+
+  /// Get the current instant on the system's suspending clock to use when
+  /// tracking the time spent running tasks.
+  static uint64_t getNanosecondsOnSuspendingClock(void);
 };
 
 // The compiler will eventually assume these.

include/swift/ABI/TaskStatus.h

@@ -454,6 +454,18 @@ class TaskDependencyStatusRecord : public TaskStatusRecord {
       JobPriority oldPriority, JobPriority newPriority);
 };
 
+class TimeSpentRunningStatusRecord : public TaskStatusRecord {
+public:
+  TimeSpentRunningStatusRecord()
+      : TaskStatusRecord(TaskStatusRecordKind::TimeSpentRunning) {}
+
+  uint64_t TimeSpentRunning = 0;
+
+  static bool classof(const TaskStatusRecord *record) {
+    return record->getKind() == TaskStatusRecordKind::TimeSpentRunning;
+  }
+};
+
 } // end namespace swift
 
 #endif

include/swift/Runtime/Concurrency.h

@@ -1072,6 +1072,20 @@ SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_donateThreadToGlobalExecutorUntil(bool (*condition)(void*),
                                                   void *context);
 
+/// Set whether or not the concurrency library is tracking the time spent
+/// running tasks. Returns the old value.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+__attribute__((__cold__)) bool
+_swift_task_setTimeSpentRunningTracked(bool isTracked);
+
+  /// Get the duration spent running the given task (so far) in nanoseconds.
+///
+/// If `AsyncTask::isTimeSpentRunningTracked()` is `false` (the common case),
+/// task duration isn't tracked and this function returns `false`.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+__attribute__((__cold__)) bool
+_swift_task_getTimeSpentRunning(AsyncTask *task, uint64_t *outNanoseconds);
+
 enum swift_clock_id : int {
   swift_clock_id_continuous = 1,
   swift_clock_id_suspending = 2,

include/swift/Runtime/ConcurrencyHooks.def

@@ -64,6 +64,12 @@ SWIFT_CONCURRENCY_HOOK(bool, swift_task_isMainExecutor, SerialExecutorRef);
 SWIFT_CONCURRENCY_HOOK(void, swift_task_donateThreadToGlobalExecutorUntil,
                        bool (*condition)(void *), void *context);
 
+SWIFT_CONCURRENCY_HOOK(bool, _swift_task_setTimeSpentRunningTracked,
+                       bool isTracked);
+
+SWIFT_CONCURRENCY_HOOK(bool, _swift_task_getTimeSpentRunning,
+                       AsyncTask *task, uint64_t *outNanoseconds);
+
 // .............................................................................
 
 #undef SWIFT_CONCURRENCY_HOOK

stdlib/public/Concurrency/ConcurrencyHooks.cpp

@@ -225,3 +225,17 @@ swift_task_donateThreadToGlobalExecutorUntil(bool (*condition)(void *),
   else
     return swift_task_donateThreadToGlobalExecutorUntilOrig(condition, context);
 }
+
+SWIFT_CC(swift)
+__attribute__((__cold__)) bool
+swift::_swift_task_setTimeSpentRunningTracked(bool isTracked) {
+  return AsyncTask::setTimeSpentRunningTracked(isTracked);
+}
+
+SWIFT_CC(swift)
+__attribute__((__cold__)) bool
+swift::_swift_task_getTimeSpentRunning(AsyncTask *task,
+                                       uint64_t *outNanoseconds) {
+  *outNanoseconds = task->getTimeSpentRunning();
+  return AsyncTask::isTimeSpentRunningTracked();
+}

stdlib/public/Concurrency/Task.cpp

@@ -646,6 +646,18 @@ const void *AsyncTask::getResumeFunctionForLogging(bool isStarting) {
   return __ptrauth_swift_runtime_function_entry_strip(result);
 }
 
+std::atomic<bool> AsyncTask::_isTimeSpentRunningTracked { false };
+
+uint64_t AsyncTask::getNanosecondsOnSuspendingClock(void) {
+  long long seconds = 0;
+  long long nanoseconds = 0;
+  swift_get_time(&seconds, &nanoseconds, swift_clock_id_suspending);
+
+  uint64_t result = static_cast<uint64_t>(seconds) * UINT64_C(1'000'000'000);
+  result += static_cast<uint64_t>(nanoseconds);
+  return result;
+}
+
 JobPriority swift::swift_task_currentPriority(AsyncTask *task) {
   // This is racey but this is to be used in an API is inherently racey anyways.
   auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
@@ -1191,6 +1203,10 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
     if (jobFlags.task_hasInitialTaskName()) {
       task->pushInitialTaskName(taskName);
     }
+
+    if (SWIFT_UNLIKELY(AsyncTask::isTimeSpentRunningTracked())) {
+      task->pushTimeSpentRunningRecord();
+    }
   }
 
   // If we're supposed to enqueue the task, do so now.

stdlib/public/Concurrency/Task.swift

@@ -991,6 +991,46 @@ internal func _getCurrentTaskNameString() -> String? {
   }
 }
 
+// MARK: - Time spent running (experimental)
+
+@available(SwiftStdlib 6.3, *)
+@_silgen_name("_swift_task_getTimeSpentRunning")
+internal func _getTimeSpentRunning(
+  _ task: Builtin.NativeObject,
+  _ outNanoseconds: UnsafeMutablePointer<UInt64>
+) -> Bool
+
+@available(SwiftStdlib 6.3, *)
+@unsafe
+private func _getTimeSpentRunning(_ task: Builtin.NativeObject) -> Duration? {
+  var result = UInt64(0)
+  guard unsafe _getTimeSpentRunning(task, &result) else {
+    return nil
+  }
+  return .nanoseconds(result)
+}
+
+@available(SwiftStdlib 6.3, *)
+extension Task {
+  /// Get the time this task has spent scheduled and running.
+  ///
+  /// This interface is experimental. It may be removed in a future release.
+  public var _timeSpentRunning: Duration? {
+    unsafe _getTimeSpentRunning(_task)
+  }
+}
+
+@available(SwiftStdlib 6.3, *)
+extension UnsafeCurrentTask {
+  /// Get the time this task has spent scheduled and running.
+  ///
+  /// This interface is experimental. It may be removed in a future release.
+  public var _timeSpentRunning: Duration? {
+    unsafe _getTimeSpentRunning(_task)
+  }
+}
+
+// MARK: -
 
 #if SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY
 @available(SwiftStdlib 5.8, *)

stdlib/public/Concurrency/TaskPrivate.h

@@ -826,6 +826,11 @@ struct AsyncTask::PrivateStorage {
       }
     }
 
+    // If we were tracking time spent running, clear that now too.
+    if (SWIFT_UNLIKELY(isTimeSpentRunningTracked())) {
+      task->popTimeSpentRunningRecord();
+    }
+
     // Drain unlock the task and remove any overrides on thread as a
     // result of the task
     auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);

stdlib/public/Concurrency/TaskStatus.cpp

@@ -873,6 +873,10 @@ static void performCancellationAction(TaskStatusRecord *record) {
   case TaskStatusRecordKind::TaskName:
     break;
 
+  // Cancellation has no impact on the time a task spends running.
+  case TaskStatusRecordKind::TimeSpentRunning:
+    break;
+
   // This should never be found, but the compiler complains if we don't check.
   case TaskStatusRecordKind::First_Reserved:
     break;
@@ -972,6 +976,9 @@ static void performEscalationAction(TaskStatusRecord *record,
   /// Task names don't matter to priority escalation.
   case TaskStatusRecordKind::TaskName:
     return;
+  // Time spent running doesn't affect priority (outside the scheduler, anyway.)
+  case TaskStatusRecordKind::TimeSpentRunning:
+    return;
   // This should never be found, but the compiler complains if we don't check.
   case TaskStatusRecordKind::First_Reserved:
     break;
@@ -1100,5 +1107,58 @@ void TaskDependencyStatusRecord::performEscalationAction(
   }
 }
 
+/**************************************************************************/
+/*************************** TIME SPENT RUNNING ***************************/
+/**************************************************************************/
+
+void AsyncTask::pushTimeSpentRunningRecord(void) {
+  void *allocation = std::malloc(sizeof(class TimeSpentRunningStatusRecord));
+  auto record = ::new (allocation) TimeSpentRunningStatusRecord();
+
+  addStatusRecord(this, record,
+                  [&](ActiveTaskStatus oldStatus, ActiveTaskStatus &newStatus) {
+                    return true; // always add the record
+                  });
+}
+
+void AsyncTask::popTimeSpentRunningRecord(void) {
+  if (auto record = popStatusRecordOfType<TimeSpentRunningStatusRecord>(this)) {
+    record->~TimeSpentRunningStatusRecord();
+    std::free(record);
+  }
+}
+
+__attribute__((cold)) uint64_t AsyncTask::getTimeSpentRunning(void) {
+  uint64_t result = 0;
+
+  withStatusRecordLock(this, [&](ActiveTaskStatus status) {
+    for (auto record : status.records()) {
+      if (auto timeRecord = dyn_cast<TimeSpentRunningStatusRecord>(record)) {
+        result = timeRecord->TimeSpentRunning;
+        break;
+      }
+    }
+  });
+
+  return result;
+}
+
+void AsyncTask::ranForNanoseconds(uint64_t ns) {
+  withStatusRecordLock(this, [&](ActiveTaskStatus status) {
+    for (auto record : status.records()) {
+      if (auto timeRecord = dyn_cast<TimeSpentRunningStatusRecord>(record)) {
+        timeRecord->TimeSpentRunning += ns;
+        break;
+      }
+    }
+  });
+
+  if (hasChildFragment()) {
+    if (auto parent = childFragment()->getParent()) {
+      parent->ranForNanoseconds(ns);
+    }
+  }
+}
+
 #define OVERRIDE_TASK_STATUS COMPATIBILITY_OVERRIDE
 #include "../CompatibilityOverride/CompatibilityOverrideIncludePath.h"