Merge pull request #73796 from mikeash/fix-resume-function-for-logging

mikeash · web-flow · commit 634295c85e99 · 2024-05-22T16:57:37.000-04:00
[Concurrency] Refine getResumeFunctionForLogging to avoid reading invalid future contexts.
diff --git a/include/swift/ABI/Task.h b/include/swift/ABI/Task.h
@@ -352,7 +352,11 @@ class AsyncTask : public Job {
   /// failing that will return ResumeTask. The returned function pointer may
   /// have a different signature than ResumeTask, and it's only for identifying
   /// code associated with the task.
-  const void *getResumeFunctionForLogging();
+  ///
+  /// If isStarting is true, look into the resume context when appropriate
+  /// to pull out a wrapped resume function. If isStarting is false, assume the
+  /// resume context may not be valid and just return the wrapper.
+  const void *getResumeFunctionForLogging(bool isStarting);
 
   /// Given that we've already fully established the job context
   /// in the current thread, start running this task.  To establish
diff --git a/stdlib/public/Concurrency/Task.cpp b/stdlib/public/Concurrency/Task.cpp
@@ -539,7 +539,7 @@ const void
     *const swift::_swift_concurrency_debug_task_future_wait_resume_adapter =
         reinterpret_cast<void *>(task_future_wait_resume_adapter);
 
-const void *AsyncTask::getResumeFunctionForLogging() {
+const void *AsyncTask::getResumeFunctionForLogging(bool isStarting) {
   const void *result = reinterpret_cast<const void *>(ResumeTask);
 
   if (ResumeTask == non_future_adapter) {
@@ -553,11 +553,16 @@ const void *AsyncTask::getResumeFunctionForLogging() {
         sizeof(FutureAsyncContextPrefix));
     result =
         reinterpret_cast<const void *>(asyncContextPrefix->asyncEntryPoint);
-  } else if (ResumeTask == task_wait_throwing_resume_adapter) {
-    auto context = static_cast<TaskFutureWaitAsyncContext *>(ResumeContext);
-    result = reinterpret_cast<const void *>(context->ResumeParent);
-  } else if (ResumeTask == task_future_wait_resume_adapter) {
-    result = reinterpret_cast<const void *>(ResumeContext->ResumeParent);
+  }
+
+  // Future contexts may not be valid if the task was already running before.
+  if (isStarting) {
+    if (ResumeTask == task_wait_throwing_resume_adapter) {
+      auto context = static_cast<TaskFutureWaitAsyncContext *>(ResumeContext);
+      result = reinterpret_cast<const void *>(context->ResumeParent);
+    } else if (ResumeTask == task_future_wait_resume_adapter) {
+      result = reinterpret_cast<const void *>(ResumeContext->ResumeParent);
+    }
   }
 
   return __ptrauth_swift_runtime_function_entry_strip(result);
diff --git a/stdlib/public/Concurrency/TaskPrivate.h b/stdlib/public/Concurrency/TaskPrivate.h
@@ -709,10 +709,10 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus {
     return record_iterator::rangeBeginning(getInnermostRecord());
   }
 
-  void traceStatusChanged(AsyncTask *task) {
+  void traceStatusChanged(AsyncTask *task, bool isStarting) {
     concurrency::trace::task_status_changed(
         task, static_cast<uint8_t>(getStoredPriority()), isCancelled(),
-        isStoredPriorityEscalated(), isRunning(), isEnqueued());
+        isStoredPriorityEscalated(), isStarting, isRunning(), isEnqueued());
   }
 };
 
@@ -938,7 +938,7 @@ inline void AsyncTask::flagAsRunning() {
       if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
                /* success */ std::memory_order_relaxed,
                /* failure */ std::memory_order_relaxed)) {
-        newStatus.traceStatusChanged(this);
+        newStatus.traceStatusChanged(this, true);
         adoptTaskVoucher(this);
         swift_task_enterThreadLocalContext(
             (char *)&_private().ExclusivityAccessSet[0]);
diff --git a/stdlib/public/Concurrency/TaskStatus.cpp b/stdlib/public/Concurrency/TaskStatus.cpp
@@ -233,7 +233,7 @@ static bool withStatusRecordLock(AsyncTask *task, ActiveTaskStatus status,
 
       status = newStatus;
 
-      status.traceStatusChanged(task);
+      status.traceStatusChanged(task, false);
       worker.flagQueueIsPublished(lockingRecord);
       installedLockRecord = true;
 
@@ -268,7 +268,7 @@ static bool withStatusRecordLock(AsyncTask *task, ActiveTaskStatus status,
     if (task->_private()._status().compare_exchange_weak(status, newStatus,
             /*success*/ std::memory_order_release,
             /*failure*/ std::memory_order_relaxed)) {
-      newStatus.traceStatusChanged(task);
+      newStatus.traceStatusChanged(task, false);
       break;
     }
   }
@@ -322,7 +322,7 @@ bool swift::addStatusRecord(AsyncTask *task, TaskStatusRecord *newRecord,
       if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
               /*success*/ std::memory_order_release,
               /*failure*/ std::memory_order_relaxed)) {
-        newStatus.traceStatusChanged(task);
+        newStatus.traceStatusChanged(task, false);
         return true;
       } else {
         // Retry
@@ -404,7 +404,7 @@ void swift::removeStatusRecord(AsyncTask *task, TaskStatusRecord *record,
         if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
                /*success*/ std::memory_order_relaxed,
                /*failure*/ std::memory_order_relaxed)) {
-          newStatus.traceStatusChanged(task);
+          newStatus.traceStatusChanged(task, false);
           return;
         }
       }
@@ -436,7 +436,7 @@ void swift::removeStatusRecord(AsyncTask *task, TaskStatusRecord *record,
       if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
              /*success*/ std::memory_order_relaxed,
              /*failure*/ std::memory_order_relaxed)) {
-        newStatus.traceStatusChanged(task);
+        newStatus.traceStatusChanged(task, false);
         return;
       }
       // Restart the loop again - someone else modified status concurrently
@@ -494,7 +494,7 @@ void swift::removeStatusRecordWhere(
         if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
                /*success*/ std::memory_order_relaxed,
                /*failure*/ std::memory_order_relaxed)) {
-          newStatus.traceStatusChanged(task);
+          newStatus.traceStatusChanged(task, false);
           return;
         }
       }
@@ -904,7 +904,7 @@ static void swift_task_cancelImpl(AsyncTask *task) {
     }
   }
 
-  newStatus.traceStatusChanged(task);
+  newStatus.traceStatusChanged(task, false);
   if (newStatus.getInnermostRecord() == NULL) {
      // No records, nothing to propagate
      return;
diff --git a/stdlib/public/Concurrency/Tracing.h b/stdlib/public/Concurrency/Tracing.h
@@ -64,7 +64,7 @@ void task_create(AsyncTask *task, AsyncTask *parent, TaskGroup *group,
 void task_destroy(AsyncTask *task);
 
 void task_status_changed(AsyncTask *task, uint8_t maxPriority, bool isCancelled,
-                         bool isEscalated, bool isRunning, bool isEnqueued);
+                         bool isEscalated, bool isStarting, bool isRunning, bool isEnqueued);
 
 void task_flags_changed(AsyncTask *task, uint8_t jobPriority, bool isChildTask,
                         bool isFuture, bool isGroupChildTask,
diff --git a/stdlib/public/Concurrency/TracingSignpost.h b/stdlib/public/Concurrency/TracingSignpost.h
@@ -193,7 +193,7 @@ inline void task_create(AsyncTask *task, AsyncTask *parent, TaskGroup *group,
       " resumefn=%p jobPriority=%u isChildTask=%{bool}d, isFuture=%{bool}d "
       "isGroupChildTask=%{bool}d isAsyncLetTask=%{bool}d parent=%" PRIx64
       " group=%p asyncLet=%p",
-      task->getTaskId(), task->getResumeFunctionForLogging(), jobPriority,
+      task->getTaskId(), task->getResumeFunctionForLogging(true), jobPriority,
       isChildTask, isFuture, isGroupChildTask, isAsyncLetTask, parentID, group,
       asyncLet);
 }
@@ -207,15 +207,15 @@ inline void task_destroy(AsyncTask *task) {
 
 inline void task_status_changed(AsyncTask *task, uint8_t maxPriority,
                                 bool isCancelled, bool isEscalated,
-                                bool isRunning, bool isEnqueued) {
+                                bool isStarting, bool isRunning, bool isEnqueued) {
   ENSURE_LOGS();
   auto id = os_signpost_id_make_with_pointer(TaskLog, task);
   os_signpost_event_emit(
       TaskLog, id, SWIFT_LOG_TASK_STATUS_CHANGED_NAME,
       "task=%" PRIx64 " resumefn=%p "
       "maxPriority=%u, isCancelled=%{bool}d "
       "isEscalated=%{bool}d, isRunning=%{bool}d, isEnqueued=%{bool}d",
-      task->getTaskId(), task->getResumeFunctionForLogging(), maxPriority,
+      task->getTaskId(), task->getResumeFunctionForLogging(isStarting), maxPriority,
       isCancelled, isEscalated, isRunning, isEnqueued);
 }
 
diff --git a/stdlib/public/Concurrency/TracingStubs.h b/stdlib/public/Concurrency/TracingStubs.h
@@ -56,7 +56,7 @@ inline void task_resume(AsyncTask *task) {}
 
 inline void task_status_changed(AsyncTask *task, uint8_t maxPriority,
                                 bool isCancelled, bool isEscalated,
-                                bool isRunning, bool isEnqueued) {}
+                                bool isStarting, bool isRunning, bool isEnqueued) {}
 
 inline void task_flags_changed(AsyncTask *task, uint8_t jobPriority,
                                bool isChildTask, bool isFuture,
