Have DumpPool return futures

Have DumpPool return futures when tasks are enqueued, instead of
having users register their task with a string. This change has the
following benefits:

* Simplifies DumpPool since it no longer needs to keep the map of
  strings to futures.

* Allows adding a set of tasks to DumpPool that might not be known at
  build time. For example, we can add a bunch of tasks from a vector,
  without needing to generate arbitrary strings for them.

* Reduces the ability to misuse DumpPool. Before this change, the user
  had to come up with strings representing each task, and needed to
  know a task with said name existed. This change moves it from being a
  runtime error to being a compile time one.

Test: atest dumpstate_smoke_test

Change-Id: Ic6c1a983b049e5236ac596e8af200a21c07ac31b
diff --git a/cmds/dumpstate/DumpPool.cpp b/cmds/dumpstate/DumpPool.cpp
index c2c8a72..4d3a67b 100644
--- a/cmds/dumpstate/DumpPool.cpp
+++ b/cmds/dumpstate/DumpPool.cpp
@@ -33,6 +33,20 @@
 
 const std::string DumpPool::PREFIX_TMPFILE_NAME = "dump-tmp.";
 
+
+void WaitForTask(std::future<std::string> future, const std::string& title, int out_fd) {
+    DurationReporter duration_reporter("Wait for " + title, true);
+
+    std::string result = future.get();
+    if (result.empty()) {
+        return;
+    }
+    DumpFileToFd(out_fd, title, result);
+    if (unlink(result.c_str())) {
+        MYLOGE("Failed to unlink (%s): %s\n", result.c_str(), strerror(errno));
+    }
+}
+
 DumpPool::DumpPool(const std::string& tmp_root) : tmp_root_(tmp_root), shutdown_(false),
         log_duration_(true) {
     assert(!tmp_root.empty());
@@ -40,31 +54,10 @@
 }
 
 DumpPool::~DumpPool() {
-    shutdown();
-}
-
-void DumpPool::start(int thread_counts) {
-    assert(thread_counts > 0);
-    assert(threads_.empty());
-    if (thread_counts > MAX_THREAD_COUNT) {
-        thread_counts = MAX_THREAD_COUNT;
-    }
-    MYLOGI("Start thread pool:%d", thread_counts);
-    shutdown_ = false;
-    for (int i = 0; i < thread_counts; i++) {
-        threads_.emplace_back(std::thread([=]() {
-            setThreadName(pthread_self(), i + 1);
-            loop();
-        }));
-    }
-}
-
-void DumpPool::shutdown() {
     std::unique_lock lock(lock_);
     if (shutdown_ || threads_.empty()) {
         return;
     }
-    futures_map_.clear();
     while (!tasks_.empty()) tasks_.pop();
 
     shutdown_ = true;
@@ -76,27 +69,22 @@
     }
     threads_.clear();
     deleteTempFiles(tmp_root_);
-    MYLOGI("shutdown thread pool");
+    MYLOGI("shutdown thread pool\n");
 }
 
-void DumpPool::waitForTask(const std::string& task_name, const std::string& title,
-        int out_fd) {
-    DurationReporter duration_reporter("Wait for " + task_name, true);
-    auto iterator = futures_map_.find(task_name);
-    if (iterator == futures_map_.end()) {
-        MYLOGW("Task %s does not exist", task_name.c_str());
-        return;
+void DumpPool::start(int thread_counts) {
+    assert(thread_counts > 0);
+    assert(threads_.empty());
+    if (thread_counts > MAX_THREAD_COUNT) {
+        thread_counts = MAX_THREAD_COUNT;
     }
-    Future future = iterator->second;
-    futures_map_.erase(iterator);
-
-    std::string result = future.get();
-    if (result.empty()) {
-        return;
-    }
-    DumpFileToFd(out_fd, title, result);
-    if (unlink(result.c_str())) {
-        MYLOGE("Failed to unlink (%s): %s\n", result.c_str(), strerror(errno));
+    MYLOGI("Start thread pool:%d\n", thread_counts);
+    shutdown_ = false;
+    for (int i = 0; i < thread_counts; i++) {
+        threads_.emplace_back(std::thread([=]() {
+            setThreadName(pthread_self(), i + 1);
+            loop();
+        }));
     }
 }
 
diff --git a/cmds/dumpstate/DumpPool.h b/cmds/dumpstate/DumpPool.h
index 0c3c2cc..9fb0fcc 100644
--- a/cmds/dumpstate/DumpPool.h
+++ b/cmds/dumpstate/DumpPool.h
@@ -18,7 +18,6 @@
 #define FRAMEWORK_NATIVE_CMD_DUMPPOOL_H_
 
 #include <future>
-#include <map>
 #include <queue>
 #include <string>
 
@@ -32,8 +31,26 @@
 class DumpPoolTest;
 
 /*
+ * Waits until the task is finished. Dumps the task results to the specified
+ * out_fd.
+ *
+ * |future| The task future.
+ * |title| Dump title string to the out_fd, an empty string for nothing.
+ * |out_fd| The target file to dump the result from the task.
+ */
+void WaitForTask(std::future<std::string> future, const std::string& title, int out_fd);
+
+/*
+ * Waits until the task is finished. Dumps the task results to the STDOUT_FILENO.
+ */
+
+inline void WaitForTask(std::future<std::string> future) {
+    WaitForTask(std::move(future), "", STDOUT_FILENO);
+}
+
+/*
  * A thread pool with the fixed number of threads to execute multiple dump tasks
- * simultaneously for the dumpstate. The dump task is a callable function. It
+ * simultaneously for dumpstate. The dump task is a callable function. It
  * could include a file descriptor as a parameter to redirect dump results, if
  * it needs to output results to the bugreport. This can avoid messing up
  * bugreport's results when multiple dump tasks are running at the same time.
@@ -44,13 +61,16 @@
  * }
  * ...
  * DumpPool pool(tmp_root);
- * pool.enqueueTaskWithFd("TaskName", &DumpFoo, std::placeholders::_1);
+ * auto task = pool.enqueueTaskWithFd("TaskName", &DumpFoo, std::placeholders::_1);
  * ...
- * pool.waitForTask("TaskName");
+ * WaitForTask(task);
  *
  * DumpFoo is a callable function included a out_fd parameter. Using the
  * enqueueTaskWithFd method in DumpPool to enqueue the task to the pool. The
  * std::placeholders::_1 is a placeholder for DumpPool to pass a fd argument.
+ *
+ * std::futures returned by `enqueueTask*()` must all have their `get` methods
+ * called, or have been destroyed before the DumpPool itself is destroyed.
  */
 class DumpPool {
   friend class android::os::dumpstate::DumpPoolTest;
@@ -63,6 +83,12 @@
      * files.
      */
     explicit DumpPool(const std::string& tmp_root);
+
+    /*
+     * Will waits until all threads exit the loop. Destroying DumpPool before destroying the
+     * associated std::futures created by `enqueueTask*` will cause an abort on Android because
+     * Android is built with `-fno-exceptions`.
+     */
     ~DumpPool();
 
     /*
@@ -73,68 +99,47 @@
     void start(int thread_counts = MAX_THREAD_COUNT);
 
     /*
-     * Requests to shutdown the pool and waits until all threads exit the loop.
-     */
-    void shutdown();
-
-    /*
      * Adds a task into the queue of the thread pool.
      *
-     * |task_name| The name of the task. It's also the title of the
+     * |duration_title| The name of the task. It's also the title of the
      * DurationReporter log.
      * |f| Callable function to execute the task.
      * |args| A list of arguments.
      *
      * TODO(b/164369078): remove this api to have just one enqueueTask for consistency.
      */
-    template<class F, class... Args> void enqueueTask(const std::string& task_name, F&& f,
-            Args&&... args) {
+    template<class F, class... Args>
+    std::future<std::string> enqueueTask(const std::string& duration_title, F&& f, Args&&... args) {
         std::function<void(void)> func = std::bind(std::forward<F>(f),
                 std::forward<Args>(args)...);
-        futures_map_[task_name] = post(task_name, func);
+        auto future = post(duration_title, func);
         if (threads_.empty()) {
             start();
         }
+        return future;
     }
 
     /*
      * Adds a task into the queue of the thread pool. The task takes a file
      * descriptor as a parameter to redirect dump results to a temporary file.
      *
-     * |task_name| The name of the task. It's also the title of the
-     * DurationReporter log.
+     * |duration_title| The title of the DurationReporter log.
      * |f| Callable function to execute the task.
      * |args| A list of arguments. A placeholder std::placeholders::_1 as a fd
      * argument needs to be included here.
      */
-    template<class F, class... Args> void enqueueTaskWithFd(const std::string& task_name, F&& f,
-            Args&&... args) {
+    template<class F, class... Args> std::future<std::string> enqueueTaskWithFd(
+            const std::string& duration_title, F&& f, Args&&... args) {
         std::function<void(int)> func = std::bind(std::forward<F>(f),
                 std::forward<Args>(args)...);
-        futures_map_[task_name] = post(task_name, func);
+        auto future = post(duration_title, func);
         if (threads_.empty()) {
             start();
         }
+        return future;
     }
 
     /*
-     * Waits until the task is finished. Dumps the task results to the STDOUT_FILENO.
-     */
-    void waitForTask(const std::string& task_name) {
-        waitForTask(task_name, "", STDOUT_FILENO);
-    }
-
-    /*
-     * Waits until the task is finished. Dumps the task results to the specified
-     * out_fd.
-     *
-     * |task_name| The name of the task.
-     * |title| Dump title string to the out_fd, an empty string for nothing.
-     * |out_fd| The target file to dump the result from the task.
-     */
-    void waitForTask(const std::string& task_name, const std::string& title, int out_fd);
-
-    /*
      * Deletes temporary files created by DumpPool.
      */
     void deleteTempFiles();
@@ -143,22 +148,22 @@
 
   private:
     using Task = std::packaged_task<std::string()>;
-    using Future = std::shared_future<std::string>;
 
     template<class T> void invokeTask(T dump_func, const std::string& duration_title, int out_fd);
 
-    template<class T> Future post(const std::string& task_name, T dump_func) {
+    template<class T>
+    std::future<std::string> post(const std::string& duration_title, T dump_func) {
         Task packaged_task([=]() {
             std::unique_ptr<TmpFile> tmp_file_ptr = createTempFile();
             if (!tmp_file_ptr) {
                 return std::string("");
             }
-            invokeTask(dump_func, task_name, tmp_file_ptr->fd.get());
+            invokeTask(dump_func, duration_title, tmp_file_ptr->fd.get());
             fsync(tmp_file_ptr->fd.get());
             return std::string(tmp_file_ptr->path);
         });
         std::unique_lock lock(lock_);
-        auto future = packaged_task.get_future().share();
+        auto future = packaged_task.get_future();
         tasks_.push(std::move(packaged_task));
         condition_variable_.notify_one();
         return future;
@@ -194,7 +199,6 @@
 
     std::vector<std::thread> threads_;
     std::queue<Task> tasks_;
-    std::map<std::string, Future> futures_map_;
 
     DISALLOW_COPY_AND_ASSIGN(DumpPool);
 };
diff --git a/cmds/dumpstate/dumpstate.cpp b/cmds/dumpstate/dumpstate.cpp
index 8ccd940..17f52ce 100644
--- a/cmds/dumpstate/dumpstate.cpp
+++ b/cmds/dumpstate/dumpstate.cpp
@@ -120,6 +120,7 @@
 using android::os::dumpstate::DumpPool;
 using android::os::dumpstate::PropertiesHelper;
 using android::os::dumpstate::TaskQueue;
+using android::os::dumpstate::WaitForTask;
 
 // Keep in sync with
 // frameworks/base/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -218,9 +219,9 @@
     RUN_SLOW_FUNCTION_AND_LOG(log_title, func_ptr, __VA_ARGS__);               \
     RETURN_IF_USER_DENIED_CONSENT();
 
-#define WAIT_TASK_WITH_CONSENT_CHECK(task_name, pool_ptr) \
+#define WAIT_TASK_WITH_CONSENT_CHECK(future) \
     RETURN_IF_USER_DENIED_CONSENT();                      \
-    pool_ptr->waitForTask(task_name);                     \
+    WaitForTask(future);                     \
     RETURN_IF_USER_DENIED_CONSENT();
 
 static const char* WAKE_LOCK_NAME = "dumpstate_wakelock";
@@ -1549,15 +1550,18 @@
     DurationReporter duration_reporter("DUMPSTATE");
 
     // Enqueue slow functions into the thread pool, if the parallel run is enabled.
+    std::future<std::string> dump_hals, dump_incident_report, dump_board, dump_checkins;
     if (ds.dump_pool_) {
         // Pool was shutdown in DumpstateDefaultAfterCritical method in order to
         // drop root user. Restarts it with two threads for the parallel run.
         ds.dump_pool_->start(/* thread_counts = */2);
 
-        ds.dump_pool_->enqueueTaskWithFd(DUMP_HALS_TASK, &DumpHals, _1);
-        ds.dump_pool_->enqueueTask(DUMP_INCIDENT_REPORT_TASK, &DumpIncidentReport);
-        ds.dump_pool_->enqueueTaskWithFd(DUMP_BOARD_TASK, &Dumpstate::DumpstateBoard, &ds, _1);
-        ds.dump_pool_->enqueueTaskWithFd(DUMP_CHECKINS_TASK, &DumpCheckins, _1);
+        dump_hals = ds.dump_pool_->enqueueTaskWithFd(DUMP_HALS_TASK, &DumpHals, _1);
+        dump_incident_report = ds.dump_pool_->enqueueTask(
+            DUMP_INCIDENT_REPORT_TASK, &DumpIncidentReport);
+        dump_board = ds.dump_pool_->enqueueTaskWithFd(
+            DUMP_BOARD_TASK, &Dumpstate::DumpstateBoard, &ds, _1);
+        dump_checkins = ds.dump_pool_->enqueueTaskWithFd(DUMP_CHECKINS_TASK, &DumpCheckins, _1);
     }
 
     // Dump various things. Note that anything that takes "long" (i.e. several seconds) should
@@ -1592,7 +1596,7 @@
                                          CommandOptions::AS_ROOT);
 
     if (ds.dump_pool_) {
-        WAIT_TASK_WITH_CONSENT_CHECK(DUMP_HALS_TASK, ds.dump_pool_);
+        WAIT_TASK_WITH_CONSENT_CHECK(std::move(dump_hals));
     } else {
         RUN_SLOW_FUNCTION_WITH_CONSENT_CHECK_AND_LOG(DUMP_HALS_TASK, DumpHals);
     }
@@ -1689,7 +1693,7 @@
     ds.AddDir(SNAPSHOTCTL_LOG_DIR, false);
 
     if (ds.dump_pool_) {
-        WAIT_TASK_WITH_CONSENT_CHECK(DUMP_BOARD_TASK, ds.dump_pool_);
+        WAIT_TASK_WITH_CONSENT_CHECK(std::move(dump_board));
     } else {
         RUN_SLOW_FUNCTION_WITH_CONSENT_CHECK_AND_LOG(DUMP_BOARD_TASK, ds.DumpstateBoard);
     }
@@ -1718,7 +1722,7 @@
     ds.AddDir("/data/misc/bluetooth/logs", true);
 
     if (ds.dump_pool_) {
-        WAIT_TASK_WITH_CONSENT_CHECK(DUMP_CHECKINS_TASK, ds.dump_pool_);
+        WAIT_TASK_WITH_CONSENT_CHECK(std::move(dump_checkins));
     } else {
         RUN_SLOW_FUNCTION_WITH_CONSENT_CHECK_AND_LOG(DUMP_CHECKINS_TASK, DumpCheckins);
     }
@@ -1752,7 +1756,7 @@
     dump_frozen_cgroupfs();
 
     if (ds.dump_pool_) {
-        WAIT_TASK_WITH_CONSENT_CHECK(DUMP_INCIDENT_REPORT_TASK, ds.dump_pool_);
+        WAIT_TASK_WITH_CONSENT_CHECK(std::move(dump_incident_report));
     } else {
         RUN_SLOW_FUNCTION_WITH_CONSENT_CHECK_AND_LOG(DUMP_INCIDENT_REPORT_TASK,
                 DumpIncidentReport);
@@ -1777,6 +1781,7 @@
     time_t logcat_ts = time(nullptr);
 
     /* collect stack traces from Dalvik and native processes (needs root) */
+    std::future<std::string> dump_traces;
     if (dump_pool_) {
         RETURN_IF_USER_DENIED_CONSENT();
         // One thread is enough since we only need to enqueue DumpTraces here.
@@ -1784,7 +1789,8 @@
 
         // DumpTraces takes long time, post it to the another thread in the
         // pool, if pool is available
-        dump_pool_->enqueueTask(DUMP_TRACES_TASK, &Dumpstate::DumpTraces, &ds, &dump_traces_path);
+        dump_traces = dump_pool_->enqueueTask(
+            DUMP_TRACES_TASK, &Dumpstate::DumpTraces, &ds, &dump_traces_path);
     } else {
         RUN_SLOW_FUNCTION_WITH_CONSENT_CHECK_AND_LOG(DUMP_TRACES_TASK, ds.DumpTraces,
                 &dump_traces_path);
@@ -1833,12 +1839,11 @@
 
     if (dump_pool_) {
         RETURN_IF_USER_DENIED_CONSENT();
-        dump_pool_->waitForTask(DUMP_TRACES_TASK);
+        WaitForTask(std::move(dump_traces));
 
-        // Current running thread in the pool is the root user also. Shutdown
-        // the pool and restart later to ensure all threads in the pool could
-        // drop the root user.
-        dump_pool_->shutdown();
+        // Current running thread in the pool is the root user also. Delete
+        // the pool and make a new one later to ensure none of threads in the pool are root.
+        dump_pool_ = std::make_unique<DumpPool>(bugreport_internal_dir_);
     }
     if (!DropRootUser()) {
         return Dumpstate::RunStatus::ERROR;
@@ -1869,8 +1874,9 @@
     } else {
         // DumpHals takes long time, post it to the another thread in the pool,
         // if pool is available.
+        std::future<std::string> dump_hals;
         if (ds.dump_pool_) {
-            ds.dump_pool_->enqueueTaskWithFd(DUMP_HALS_TASK, &DumpHals, _1);
+            dump_hals = ds.dump_pool_->enqueueTaskWithFd(DUMP_HALS_TASK, &DumpHals, _1);
         }
         // Contains various system properties and process startup info.
         do_dmesg();
@@ -1880,7 +1886,7 @@
         DoKmsg();
         // DumpHals contains unrelated hardware info (camera, NFC, biometrics, ...).
         if (ds.dump_pool_) {
-            ds.dump_pool_->waitForTask(DUMP_HALS_TASK);
+            WaitForTask(std::move(dump_hals));
         } else {
             RUN_SLOW_FUNCTION_AND_LOG(DUMP_HALS_TASK, DumpHals);
         }
@@ -1914,12 +1920,14 @@
 
     // Starts thread pool after the root user is dropped, and two additional threads
     // are created for DumpHals in the DumpstateRadioCommon and DumpstateBoard.
+    std::future<std::string> dump_board;
     if (ds.dump_pool_) {
         ds.dump_pool_->start(/*thread_counts =*/2);
 
         // DumpstateBoard takes long time, post it to the another thread in the pool,
         // if pool is available.
-        ds.dump_pool_->enqueueTaskWithFd(DUMP_BOARD_TASK, &Dumpstate::DumpstateBoard, &ds, _1);
+        dump_board = ds.dump_pool_->enqueueTaskWithFd(
+            DUMP_BOARD_TASK, &Dumpstate::DumpstateBoard, &ds, _1);
     }
 
     DumpstateRadioCommon(include_sensitive_info);
@@ -2002,7 +2010,7 @@
     printf("========================================================\n");
 
     if (ds.dump_pool_) {
-        ds.dump_pool_->waitForTask(DUMP_BOARD_TASK);
+        WaitForTask(std::move(dump_board));
     } else {
         RUN_SLOW_FUNCTION_AND_LOG(DUMP_BOARD_TASK, ds.DumpstateBoard);
     }
@@ -3226,8 +3234,7 @@
 
 void Dumpstate::ShutdownDumpPool() {
     if (dump_pool_) {
-        dump_pool_->shutdown();
-        dump_pool_ = nullptr;
+        dump_pool_.reset();
     }
     if (zip_entry_tasks_) {
         zip_entry_tasks_->run(/* do_cancel = */true);
diff --git a/cmds/dumpstate/dumpstate.h b/cmds/dumpstate/dumpstate.h
index 4a99cd8..ee6b1ae 100644
--- a/cmds/dumpstate/dumpstate.h
+++ b/cmds/dumpstate/dumpstate.h
@@ -478,7 +478,7 @@
     // This is useful for debugging.
     std::string log_path_;
 
-    // Full path of the bugreport file, be it zip or text, inside bugreport_internal_dir_.
+    // Full path of the bugreport zip file inside bugreport_internal_dir_.
     std::string path_;
 
     // Full path of the file containing the screenshot (when requested).
diff --git a/cmds/dumpstate/tests/dumpstate_test.cpp b/cmds/dumpstate/tests/dumpstate_test.cpp
index 42beb2b..70b4e5c 100644
--- a/cmds/dumpstate/tests/dumpstate_test.cpp
+++ b/cmds/dumpstate/tests/dumpstate_test.cpp
@@ -1720,14 +1720,13 @@
         dprintf(out_fd, "C");
     };
     setLogDuration(/* log_duration = */false);
-    dump_pool_->enqueueTaskWithFd(/* task_name = */"1", dump_func_1, std::placeholders::_1);
-    dump_pool_->enqueueTaskWithFd(/* task_name = */"2", dump_func_2, std::placeholders::_1);
-    dump_pool_->enqueueTaskWithFd(/* task_name = */"3", dump_func_3, std::placeholders::_1);
+    auto t1 = dump_pool_->enqueueTaskWithFd("", dump_func_1, std::placeholders::_1);
+    auto t2 = dump_pool_->enqueueTaskWithFd("", dump_func_2, std::placeholders::_1);
+    auto t3 = dump_pool_->enqueueTaskWithFd("", dump_func_3, std::placeholders::_1);
 
-    dump_pool_->waitForTask("1", "", out_fd_.get());
-    dump_pool_->waitForTask("2", "", out_fd_.get());
-    dump_pool_->waitForTask("3", "", out_fd_.get());
-    dump_pool_->shutdown();
+    WaitForTask(std::move(t1), "", out_fd_.get());
+    WaitForTask(std::move(t2), "", out_fd_.get());
+    WaitForTask(std::move(t3), "", out_fd_.get());
 
     std::string result;
     ReadFileToString(out_path_, &result);
@@ -1741,9 +1740,8 @@
         run_1 = true;
     };
 
-    dump_pool_->enqueueTask(/* task_name = */"1", dump_func_1);
-    dump_pool_->waitForTask("1", "", out_fd_.get());
-    dump_pool_->shutdown();
+    auto t1 = dump_pool_->enqueueTask(/* duration_title = */"1", dump_func_1);
+    WaitForTask(std::move(t1), "", out_fd_.get());
 
     std::string result;
     ReadFileToString(out_path_, &result);
@@ -1752,27 +1750,6 @@
     EXPECT_THAT(getTempFileCounts(kTestDataPath), Eq(0));
 }
 
-TEST_F(DumpPoolTest, Shutdown_withoutCrash) {
-    bool run_1 = false;
-    auto dump_func_1 = [&]() {
-        run_1 = true;
-    };
-    auto dump_func = []() {
-        sleep(1);
-    };
-
-    dump_pool_->start(/* thread_counts = */1);
-    dump_pool_->enqueueTask(/* task_name = */"1", dump_func_1);
-    dump_pool_->enqueueTask(/* task_name = */"2", dump_func);
-    dump_pool_->enqueueTask(/* task_name = */"3", dump_func);
-    dump_pool_->enqueueTask(/* task_name = */"4", dump_func);
-    dump_pool_->waitForTask("1", "", out_fd_.get());
-    dump_pool_->shutdown();
-
-    EXPECT_TRUE(run_1);
-    EXPECT_THAT(getTempFileCounts(kTestDataPath), Eq(0));
-}
-
 class TaskQueueTest : public DumpstateBaseTest {
 public:
     void SetUp() {