Add pending count overflow protection to TaskGroup

Author: ktoso

stdlib/public/Concurrency/TaskGroup.cpp

@@ -467,6 +467,19 @@ struct TaskGroupStatus {
   static const uint64_t maskDiscardingPending   = 0b0011111111111111111111111111111111111111111111111111111111111111;
   static const uint64_t onePendingTask          = 0b0000000000000000000000000000000000000000000000000000000000000001;
 
+  /// Depending on kind of task group, we can either support 2^31 or 2^62 pending tasks.
+  ///
+  /// While a discarding task group's max pending count is unrealistic to be exceeded, the lower
+  /// maximum number used in an accumulating task group has potential to be exceeded, and thus we must crash
+  /// rather than start overflowing status if this were to happen.
+  static uint64_t maximumPendingTasks(TaskGroupBase* group) {
+    if (group->isAccumulatingResults()) {
+      return maskAccumulatingPending;
+    } else {
+      return maskDiscardingPending;
+    }
+  }
+
   uint64_t status;
 
   bool isCancelled() {
@@ -525,6 +538,39 @@ struct TaskGroupStatus {
     return TaskGroupStatus{status | (cancel ? cancelled : 0)};
   }
 
+  static void reportPendingTaskOverflow(TaskGroupBase* group, TaskGroupStatus status) {
+    char *message;
+    swift_asprintf(
+        &message,
+        "error: %sTaskGroup: detected pending task count overflow, in task group %p! Status: %s",
+        group->isDiscardingResults() ? "Discarding" : "", group, status.to_string(group).c_str());
+
+    if (_swift_shouldReportFatalErrorsToDebugger()) {
+      RuntimeErrorDetails details = {
+          .version = RuntimeErrorDetails::currentVersion,
+          .errorType = "task-group-violation",
+          .currentStackDescription = "TaskGroup exceeded supported pending task count",
+          .framesToSkip = 1,
+      };
+      _swift_reportToDebugger(RuntimeErrorFlagFatal, message, &details);
+    }
+
+#if defined(_WIN32)
+    #define STDERR_FILENO 2
+   _write(STDERR_FILENO, message, strlen(message));
+#else
+    write(STDERR_FILENO, message, strlen(message));
+#endif
+#if defined(__APPLE__)
+    asl_log(nullptr, nullptr, ASL_LEVEL_ERR, "%s", message);
+#elif defined(__ANDROID__)
+    __android_log_print(ANDROID_LOG_FATAL, "SwiftRuntime", "%s", message);
+#endif
+
+    free(message);
+    abort();
+  }
+
   /// Pretty prints the status, as follows:
   /// If accumulating results:
   ///     TaskGroupStatus{ C:{cancelled} W:{waiting task} R:{ready tasks} P:{pending tasks} {binary repr} }
@@ -620,6 +666,11 @@ TaskGroupStatus TaskGroupBase::statusAddPendingTaskAssumeRelaxed(bool unconditio
                               std::memory_order_relaxed);
   auto s = TaskGroupStatus{old + TaskGroupStatus::onePendingTask};
 
+  if (s.pendingTasks(this) == TaskGroupStatus::maximumPendingTasks(this)) {
+    TaskGroupStatus::reportPendingTaskOverflow(this, s);
+    llvm_unreachable(); // the above call will abort();
+  }
+
   if (!unconditionally && s.isCancelled()) {
     // revert that add, it was meaningless
     auto o = status.fetch_sub(TaskGroupStatus::onePendingTask,