Track running tasks

Author: Rafał Hibner

cpp/src/arrow/util/async_util.cc

@@ -170,6 +170,7 @@ class AsyncTaskSchedulerImpl : public AsyncTaskScheduler {
     if (IsAborted()) {
       return false;
     }
+    running_tasks_set_.insert(task.get());
     SubmitTaskUnlocked(std::move(task), std::move(lk));
     return true;
   }
@@ -210,6 +211,7 @@ class AsyncTaskSchedulerImpl : public AsyncTaskScheduler {
             // OnTaskFinished might trigger the scheduler to end.  We want to ensure that
             // is the very last thing that happens after all task destructors have run so
             // we eagerly destroy the task first.
+            running_tasks_set_.erase(task_inner2.get());
             task_inner2.reset();
             OnTaskFinished(st);
           };
@@ -267,6 +269,7 @@ class AsyncTaskSchedulerImpl : public AsyncTaskScheduler {
     return DoSubmitTask(std::move(task));
   }
 
+  std::unordered_set<Task*> running_tasks_set_;
   Future<> finished_ = Future<>::Make();
   // The initial task is our first task
   int running_tasks_ = 1;