Bug 1023787 - Make Task.jsm stack rewriting play nicely with xpcshell and mochi tests. r=paolo, r=mikedeboer

testing/mochitest/browser-test.js

@@ -713,11 +713,17 @@ function testResult(aCondition, aName, aDiag, aIsTodo, aStack) {
     }
     if (aStack) {
       this.msg += "\nStack trace:\n";
-      var frame = aStack;
-      while (frame) {
-        this.msg += "    " + frame + "\n";
-        frame = frame.caller;
+      let normalized;
+      if (aStack instanceof Components.interfaces.nsIStackFrame) {
+        let frames = [];
+        for (let frame = aStack; frame; frame = frame.caller) {
+          frames.push(frame.filename + ":" + frame.name + ":" + frame.lineNumber);
+        }
+        normalized = frames.join("\n");
+      } else {
+        normalized = "" + aStack;
       }
+      this.msg += Task.Debugging.generateReadableStack(normalized, "    ");
     }
     if (aIsTodo)
       this.result = "TEST-UNEXPECTED-PASS";

testing/xpcshell/head.js

@@ -622,16 +622,17 @@ function do_throw(error, stack) {
 }
 
 function _format_stack(stack) {
+  let normalized;
   if (stack instanceof Components.interfaces.nsIStackFrame) {
-    let stack_msg = "";
-    let frame = stack;
-    while (frame != null) {
-      stack_msg += frame + "\n";
-      frame = frame.caller;
+    let frames = [];
+    for (let frame = stack; frame; frame = frame.caller) {
+      frames.push(frame.filename + ":" + frame.name + ":" + frame.lineNumber);
     }
-    return stack_msg;
+    normalized = frames.join("\n");
+  } else {
+    normalized = "" + stack;
   }
-  return "" + stack;
+  return _Task.Debugging.generateReadableStack(normalized, "    ");
 }
 
 function do_throw_todo(text, stack) {

testing/xpcshell/selftest.py

@@ -132,6 +132,32 @@ add_task(function () {
 });
 '''
 
+ADD_TASK_STACK_TRACE = '''
+Components.utils.import("resource://gre/modules/Promise.jsm", this);
+
+function run_test() { run_next_test(); }
+
+add_task(function* this_test_will_fail() {
+  for (let i = 0; i < 10; ++i) {
+    yield Promise.resolve();
+  }
+  Assert.ok(false);
+});
+'''
+
+ADD_TASK_STACK_TRACE_WITHOUT_STAR = '''
+Components.utils.import("resource://gre/modules/Promise.jsm", this);
+
+function run_test() { run_next_test(); }
+
+add_task(function this_test_will_fail() {
+  for (let i = 0; i < 10; ++i) {
+    yield Promise.resolve();
+  }
+  Assert.ok(false);
+});
+'''
+
 ADD_TEST_THROW_STRING = '''
 function run_test() {do_throw("Passing a string to do_throw")};
 '''
@@ -586,6 +612,37 @@ tail =
         self.assertEquals(0, self.x.passCount)
         self.assertEquals(1, self.x.failCount)
 
+    def testAddTaskStackTrace(self):
+        """
+        Ensuring that calling Assert.ok(false) from inside add_task()
+        results in a human-readable stack trace.
+        """
+        self.writeFile("test_add_task_stack_trace.js",
+            ADD_TASK_STACK_TRACE)
+        self.writeManifest(["test_add_task_stack_trace.js"])
+
+        self.assertTestResult(False)
+        self.assertInLog("this_test_will_fail")
+        self.assertInLog("run_next_test")
+        self.assertInLog("run_test")
+        self.assertNotInLog("Task.jsm")
+
+    def testAddTaskStackTraceWithoutStar(self):
+        """
+        Ensuring that calling Assert.ok(false) from inside add_task()
+        results in a human-readable stack trace. This variant uses deprecated
+        `function()` syntax instead of now standard `function*()`.
+        """
+        self.writeFile("test_add_task_stack_trace_without_star.js",
+            ADD_TASK_STACK_TRACE)
+        self.writeManifest(["test_add_task_stack_trace_without_star.js"])
+
+        self.assertTestResult(False)
+        self.assertInLog("this_test_will_fail")
+        self.assertInLog("run_next_test")
+        self.assertInLog("run_test")
+        self.assertNotInLog("Task.jsm")
+
     def testMissingHeadFile(self):
         """
         Ensure that missing head file results in fatal error.

toolkit/modules/Task.jsm

@@ -97,6 +97,31 @@ Cu.import("resource://gre/modules/Promise.jsm");
 // reported (possibly redundantly) so as to let programmers fix their code.
 const ERRORS_TO_REPORT = ["EvalError", "RangeError", "ReferenceError", "TypeError"];
 
+/**
+ * The Task currently being executed
+ */
+let gCurrentTask = null;
+
+/**
+ * If `true`, capture stacks whenever entering a Task and rewrite the
+ * stack any exception thrown through a Task.
+ */
+let gMaintainStack = false;
+
+
+/**
+ * Iterate through the lines of a string.
+ *
+ * @return Iterator<string>
+ */
+function* linesOf(string) {
+  let reLine = /([^\r\n])+/g;
+  let match;
+  while ((match = reLine.exec(string))) {
+    yield [match[0], match.index];
+  }
+};
+
 /**
  * Detect whether a value is a generator.
  *
@@ -241,7 +266,7 @@ function createAsyncFunction(aTask) {
  * that is fulfilled when the task terminates.
  */
 function TaskImpl(iterator) {
-  if (Task.Debugging.maintainStack) {
+  if (gMaintainStack) {
     this._stack = (new Error()).stack;
   }
   this.deferred = Promise.defer();
@@ -280,37 +305,64 @@ TaskImpl.prototype = {
    *        Resolution result or rejection exception, if any.
    */
   _run: function TaskImpl_run(aSendResolved, aSendValue) {
-    if (this._isStarGenerator) {
-      try {
-        let result = aSendResolved ? this._iterator.next(aSendValue)
-                                   : this._iterator.throw(aSendValue);
 
-        if (result.done) {
-          // The generator function returned.
-          this.deferred.resolve(result.value);
-        } else {
-          // The generator function yielded.
-          this._handleResultValue(result.value);
+    try {
+      gCurrentTask = this;
+
+      if (this._isStarGenerator) {
+        try {
+          let result = aSendResolved ? this._iterator.next(aSendValue)
+                                     : this._iterator.throw(aSendValue);
+
+          if (result.done) {
+            // The generator function returned.
+            this.deferred.resolve(result.value);
+          } else {
+            // The generator function yielded.
+            this._handleResultValue(result.value);
+          }
+        } catch (ex) {
+          // The generator function failed with an uncaught exception.
+          this._handleException(ex);
+        }
+      } else {
+        try {
+          let yielded = aSendResolved ? this._iterator.send(aSendValue)
+                                      : this._iterator.throw(aSendValue);
+          this._handleResultValue(yielded);
+        } catch (ex if ex instanceof Task.Result) {
+          // The generator function threw the special exception that allows it to
+          // return a specific value on resolution.
+          this.deferred.resolve(ex.value);
+        } catch (ex if ex instanceof StopIteration) {
+          // The generator function terminated with no specific result.
+          this.deferred.resolve(undefined);
+        } catch (ex) {
+          // The generator function failed with an uncaught exception.
+          this._handleException(ex);
         }
-      } catch (ex) {
-        // The generator function failed with an uncaught exception.
-        this._handleException(ex);
       }
-    } else {
-      try {
-        let yielded = aSendResolved ? this._iterator.send(aSendValue)
-                                    : this._iterator.throw(aSendValue);
-        this._handleResultValue(yielded);
-      } catch (ex if ex instanceof Task.Result) {
-        // The generator function threw the special exception that allows it to
-        // return a specific value on resolution.
-        this.deferred.resolve(ex.value);
-      } catch (ex if ex instanceof StopIteration) {
-        // The generator function terminated with no specific result.
-        this.deferred.resolve();
-      } catch (ex) {
-        // The generator function failed with an uncaught exception.
-        this._handleException(ex);
+    } finally {
+      //
+      // At this stage, the Task may have finished executing, or have
+      // walked through a `yield` or passed control to a sub-Task.
+      // Regardless, if we still own `gCurrentTask`, reset it. If we
+      // have not finished execution of this Task, re-entering `_run`
+      // will set `gCurrentTask` to `this` as needed.
+      //
+      // We just need to be careful here in case we hit the following
+      // pattern:
+      //
+      //   Task.spawn(foo);
+      //   Task.spawn(bar);
+      //
+      // Here, `foo` and `bar` may be interleaved, so when we finish
+      // executing `foo`, `gCurrentTask` may actually either `foo` or
+      // `bar`. If `gCurrentTask` has already been set to `bar`, leave
+      // it be and it will be reset to `null` once `bar` is complete.
+      //
+      if (gCurrentTask == this) {
+        gCurrentTask = null;
       }
     }
   },
@@ -349,43 +401,23 @@ TaskImpl.prototype = {
    *        The uncaught exception to handle.
    */
   _handleException: function TaskImpl_handleException(aException) {
+
+    gCurrentTask = this;
+
     if (aException && typeof aException == "object" && "stack" in aException) {
 
       let stack = aException.stack;
 
-      if (Task.Debugging.maintainStack &&
+      if (gMaintainStack &&
           aException._capturedTaskStack != this._stack &&
           typeof stack == "string") {
 
         // Rewrite the stack for more readability.
 
         let bottomStack = this._stack;
-        let topStack = aException.stack;
+        let topStack = stack;
 
-        // Cut `topStack` at the first line that contains Task.jsm, keep the head.
-        let reLine = /([^\r\n])+/g;
-        let match;
-        let lines = [];
-        while ((match = reLine.exec(topStack))) {
-          let line = match[0];
-          if (line.indexOf("/Task.jsm:") != -1) {
-            break;
-          }
-          lines.push(line);
-        }
-
-        // Cut `bottomStack` at the last line of the first block that contains Task.jsm
-        reLine = /([^\r\n])+/g;
-        while ((match = reLine.exec(bottomStack))) {
-          let line = match[0];
-          if (line.indexOf("/Task.jsm:") == -1) {
-            let tail = bottomStack.substring(match.index);
-            lines.push(tail);
-            break;
-          }
-        }
-
-        stack = lines.join("\n");
+        stack = Task.Debugging.generateReadableStack(stack);
 
         aException.stack = stack;
 
@@ -414,9 +446,74 @@ TaskImpl.prototype = {
     }
 
     this.deferred.reject(aException);
+  },
+
+  get callerStack() {
+    // Cut `this._stack` at the last line of the first block that
+    // contains Task.jsm, keep the tail.
+    for (let [line, index] of linesOf(this._stack || "")) {
+      if (line.indexOf("/Task.jsm:") == -1) {
+        return this._stack.substring(index);
+      }
+    }
+    return "";
   }
 };
 
+
 Task.Debugging = {
-  maintainStack: false
+
+  /**
+   * Control stack rewriting.
+   *
+   * If `true`, any exception thrown from a Task will be rewritten to
+   * provide a human-readable stack trace. Otherwise, stack traces will
+   * be left unchanged.
+   *
+   * There is a (small but existing) runtime cost associated to stack
+   * rewriting, so you should probably not activate this in production
+   * code.
+   *
+   * @type {bool}
+   */
+  get maintainStack() {
+    return gMaintainStack;
+  },
+  set maintainStack(x) {
+    if (!x) {
+      gCurrentTask = null;
+    }
+    return gMaintainStack = x;
+  },
+
+  /**
+   * Generate a human-readable stack for an error raised in
+   * a Task.
+   *
+   * @param {string} topStack The stack provided by the error.
+   * @param {string=} prefix Optionally, a prefix for each line.
+   */
+  generateReadableStack: function(topStack, prefix = "") {
+    if (!gCurrentTask) {
+      return topStack;
+    }
+
+    // Cut `topStack` at the first line that contains Task.jsm, keep the head.
+    let lines = [];
+    for (let [line] of linesOf(topStack)) {
+      if (line.indexOf("/Task.jsm:") != -1) {
+        break;
+      }
+      lines.push(prefix + line);
+    }
+    if (!prefix) {
+      lines.push(gCurrentTask.callerStack);
+    } else {
+      for (let [line] of linesOf(gCurrentTask.callerStack)) {
+        lines.push(prefix + line);
+      }
+    }
+
+    return lines.join("\n");
+  }
 };

toolkit/modules/tests/xpcshell/test_task.js

@@ -539,6 +539,31 @@ add_test(function test_async_method_yield_reject_stack() {
   });
 });
 
+// Test that two tasks whose execution takes place interleaved do not capture each other's stack.
+add_test(function test_throw_stack_do_not_capture_the_wrong_task() {
+  for (let iter_a of [3, 4, 5]) { // Vary the interleaving
+    for (let iter_b of [3, 4, 5]) {
+      Task.spawn(function* task_a() {
+        for (let i = 0; i < iter_a; ++i) {
+          yield Promise.resolve();
+        }
+        throw new Error("BOOM");
+      }).then(do_throw, function(ex) {
+        do_check_rewritten_stack(["task_a",
+                                  "test_throw_stack_do_not_capture_the_wrong_task"],
+                                  ex);
+        do_check_true(!ex.stack.contains("task_b"));
+        run_next_test();
+      });
+      Task.spawn(function* task_b() {
+        for (let i = 0; i < iter_b; ++i) {
+          yield Promise.resolve();
+        }
+      });
+    }
+  }
+});
+
 // Put things together
 add_test(function test_throw_complex_stack()
 {
@@ -594,7 +619,24 @@ add_test(function test_throw_complex_stack()
     });
 });
 
-add_test(function exit_stack_tests() {
+add_test(function test_without_maintainStack() {
+  do_print("Calling generateReadableStack without a Task");
+  Task.Debugging.generateReadableStack(new Error("Not a real error"));
+
   Task.Debugging.maintainStack = false;
+
+  do_print("Calling generateReadableStack with neither a Task nor maintainStack");
+  Task.Debugging.generateReadableStack(new Error("Not a real error"));
+
+  do_print("Calling generateReadableStack without maintainStack");
+  Task.spawn(function*() {
+    Task.Debugging.generateReadableStack(new Error("Not a real error"));
+    run_next_test();
+  });
+});
+
+add_test(function exit_stack_tests() {
+  Task.Debugging.maintainStack = maintainStack;
   run_next_test();
 });
+