Bug 1470151 - Make dumpCoverage and resetCoverage return a Promise that is resolved when the parent process and all content processes are done with dumping or resetting coverage counters. r=froydnj

tools/code-coverage/nsCodeCoverage.cpp

@@ -7,9 +7,11 @@
 #include "mozilla/CodeCoverageHandler.h"
 #include "mozilla/Unused.h"
 #include "mozilla/dom/ContentParent.h"
+#include "mozilla/dom/Promise.h"
 
 using namespace mozilla;
 using namespace mozilla::dom;
+using namespace mozilla::ipc;
 
 NS_IMPL_ISUPPORTS(nsCodeCoverage,
                   nsICodeCoverage)
@@ -22,28 +24,82 @@ nsCodeCoverage::~nsCodeCoverage()
 {
 }
 
-NS_IMETHODIMP nsCodeCoverage::DumpCounters()
+enum RequestType { Dump, Reset };
+
+class ProcessCount final {
+  NS_INLINE_DECL_REFCOUNTING(ProcessCount);
+
+public:
+  ProcessCount(uint32_t c) : mCount(c) {}
+  operator uint32_t() const { return mCount; }
+  ProcessCount& operator--() { mCount--; return *this; }
+
+private:
+  ~ProcessCount() {}
+  uint32_t mCount;
+};
+
+nsresult Request(JSContext* cx, Promise** aPromise, RequestType requestType)
 {
   MOZ_ASSERT(XRE_IsParentProcess());
   MOZ_ASSERT(NS_IsMainThread());
 
-  CodeCoverageHandler::DumpCounters(0);
-  for (auto* cp : ContentParent::AllProcesses(ContentParent::eLive)) {
-    Unused << cp->SendDumpCodeCoverageCounters();
+  nsIGlobalObject* global = xpc::NativeGlobal(JS::CurrentGlobalOrNull(cx));
+  if (NS_WARN_IF(!global)) {
+    return NS_ERROR_FAILURE;
   }
 
+  ErrorResult result;
+  RefPtr<Promise> promise = Promise::Create(global, result);
+  if (NS_WARN_IF(result.Failed())) {
+    return result.StealNSResult();
+  }
+
+  uint32_t processCount = 0;
+  for (auto* cp : ContentParent::AllProcesses(ContentParent::eLive)) {
+    ++processCount;
+  }
+
+  if (requestType == RequestType::Dump) {
+    CodeCoverageHandler::DumpCounters();
+  } else if (requestType == RequestType::Reset) {
+    CodeCoverageHandler::ResetCounters();
+  }
+
+  if (processCount == 0) {
+    promise->MaybeResolveWithUndefined();
+  } else {
+    RefPtr<ProcessCount> processCountHolder(new ProcessCount(processCount));
+
+    auto resolve = [processCountHolder, promise](bool unused) {
+      if (--(*processCountHolder) == 0) {
+        promise->MaybeResolveWithUndefined();
+      }
+    };
+
+    auto reject = [promise](ResponseRejectReason aReason) {
+      promise->MaybeReject(NS_ERROR_FAILURE);
+    };
+
+    for (auto* cp : ContentParent::AllProcesses(ContentParent::eLive)) {
+      if (requestType == RequestType::Dump) {
+        cp->SendDumpCodeCoverageCounters(resolve, reject);
+      } else if (requestType == RequestType::Reset) {
+        cp->SendResetCodeCoverageCounters(resolve, reject);
+      }
+    }
+  }
+
+  promise.forget(aPromise);
   return NS_OK;
 }
 
-NS_IMETHODIMP nsCodeCoverage::ResetCounters()
+NS_IMETHODIMP nsCodeCoverage::DumpCounters(JSContext *cx, Promise** aPromise)
 {
-  MOZ_ASSERT(XRE_IsParentProcess());
-  MOZ_ASSERT(NS_IsMainThread());
-
-  CodeCoverageHandler::ResetCounters(0);
-  for (auto* cp : ContentParent::AllProcesses(ContentParent::eLive)) {
-    Unused << cp->SendResetCodeCoverageCounters();
-  }
-
-  return NS_OK;
+  return Request(cx, aPromise, RequestType::Dump);
+}
+
+NS_IMETHODIMP nsCodeCoverage::ResetCounters(JSContext *cx, Promise** aPromise)
+{
+  return Request(cx, aPromise, RequestType::Reset);
 }