From b336d5e13d5a8b0065c8624a5f908172e09f12d3 Mon Sep 17 00:00:00 2001
From: Paolo Amadini <paolo.mozmail@amadzone.org>
Date: Mon, 12 Jun 2017 16:42:47 +0100
Subject: [PATCH] Bug 1371895 - Support legacy Task.jsm generators in
 DeferredTask.jsm. r=florian

This ensures legacy add-on compatibility by loading Task.jsm lazily when required. Once Task.jsm is removed, this code can be changed back to support only async functions.

MozReview-Commit-ID: 15nY8yArNlZ

--HG--
extra : rebase_source : 129fc4958730fb85176dbb26131cce79e531d41f
---
 toolkit/modules/DeferredTask.jsm              | 22 +++++++++++++--
 .../tests/xpcshell/test_DeferredTask.js       | 27 +++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/toolkit/modules/DeferredTask.jsm b/toolkit/modules/DeferredTask.jsm
index 3888f585b6ab..0c47ebae608b 100644
--- a/toolkit/modules/DeferredTask.jsm
+++ b/toolkit/modules/DeferredTask.jsm
@@ -90,6 +90,8 @@ Cu.import("resource://gre/modules/XPCOMUtils.jsm");
 
 XPCOMUtils.defineLazyModuleGetter(this, "Promise",
                                   "resource://gre/modules/Promise.jsm");
+XPCOMUtils.defineLazyModuleGetter(this, "Task",
+                                  "resource://gre/modules/Task.jsm");
 
 const Timer = Components.Constructor("@mozilla.org/timer;1", "nsITimer",
                                      "initWithCallback");
@@ -272,7 +274,7 @@ this.DeferredTask.prototype = {
 
     runningDeferred.resolve((async () => {
       // Execute the provided function asynchronously.
-      await (this._taskFn() || Promise.resolve()).then(null, Cu.reportError);
+      await this._runTask();
 
       // Now that the task has finished, we check the state of the object to
       // determine if we should restart the task again.
@@ -284,7 +286,7 @@ this.DeferredTask.prototype = {
           // property should return false while the task is running, and should
           // remain false after the last execution terminates.
           this._armed = false;
-          await (this._taskFn() || Promise.resolve()).then(null, Cu.reportError);
+          await this._runTask();
         }
       }
 
@@ -293,4 +295,20 @@ this.DeferredTask.prototype = {
       this._runningPromise = null;
     })().then(null, Cu.reportError));
   },
+
+  /**
+   * Executes the associated task and catches exceptions.
+   */
+  async _runTask() {
+    try {
+      let result = this._taskFn();
+      if (Object.prototype.toString.call(result) == "[object Generator]") {
+        await Task.spawn(result); // eslint-disable-line mozilla/no-task
+      } else {
+        await result;
+      }
+    } catch (ex) {
+      Cu.reportError(ex);
+    }
+  },
 };
diff --git a/toolkit/modules/tests/xpcshell/test_DeferredTask.js b/toolkit/modules/tests/xpcshell/test_DeferredTask.js
index 52d7bee6bd36..fac2230e56f1 100644
--- a/toolkit/modules/tests/xpcshell/test_DeferredTask.js
+++ b/toolkit/modules/tests/xpcshell/test_DeferredTask.js
@@ -185,6 +185,33 @@ add_test(function test_arm_async() {
   });
 });
 
+/**
+ * Checks that "arm" accepts a Task.jsm generator function.
+ */
+add_test(function test_arm_async_generator() {
+  let deferredTask = new DeferredTask(function*() {
+    yield Promise.resolve();
+    run_next_test();
+  }, 50);
+
+  deferredTask.arm();
+});
+
+/**
+ * Checks that "arm" accepts a Task.jsm legacy generator function.
+ */
+add_test(function test_arm_async_legacy_generator() {
+  // ESLint cannot parse legacy generator functions, so we need an eval block.
+  /* eslint-disable no-eval */
+  let deferredTask = new DeferredTask(eval(`(function() {
+    yield Promise.resolve();
+    run_next_test();
+  })`), 50);
+  /* eslint-enable no-eval */
+
+  deferredTask.arm();
+});
+
 /**
  * Checks that an armed task can be disarmed.
  */