From 8d6640a18fabae146f1b663ed19ac9b85cb83296 Mon Sep 17 00:00:00 2001
From: Kris Maglione <maglione.k@gmail.com>
Date: Wed, 12 Jul 2017 16:58:15 -0700
Subject: [PATCH] Bug 1380267: Fix shutdown blocker corner cases. r=aswan

MozReview-Commit-ID: 8Dgvnm96hgT

--HG--
extra : rebase_source : c5cd4f096bf85af2ff098182037fff36a479bada
---
 toolkit/components/extensions/Extension.jsm          | 10 +++++++++-
 .../extensions/extension-process-script.js           | 12 +++++++-----
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/toolkit/components/extensions/Extension.jsm b/toolkit/components/extensions/Extension.jsm
index de2a81ae8f82..ade0ee7c8822 100644
--- a/toolkit/components/extensions/Extension.jsm
+++ b/toolkit/components/extensions/Extension.jsm
@@ -982,7 +982,7 @@ this.Extension = class extends ExtensionData {
           resolve();
         }
       };
-      ppmm.addMessageListener(msg + "Complete", listener);
+      ppmm.addMessageListener(msg + "Complete", listener, true);
       Services.obs.addObserver(observer, "message-manager-close");
       Services.obs.addObserver(observer, "message-manager-disconnect");
 
@@ -1198,6 +1198,14 @@ this.Extension = class extends ExtensionData {
       `Extension Shutdown: ${this.id} (${this.manifest && this.name})`,
       promise.catch(() => {}));
 
+    // If we already have a shutdown promise for this extension, wait
+    // for it to complete before replacing it with a new one. This can
+    // sometimes happen during tests with rapid startup/shutdown cycles
+    // of multiple versions.
+    if (shutdownPromises.has(this.id)) {
+      await shutdownPromises.get(this.id);
+    }
+
     let cleanup = () => {
       shutdownPromises.delete(this.id);
     };
diff --git a/toolkit/components/extensions/extension-process-script.js b/toolkit/components/extensions/extension-process-script.js
index f4b81cb8e688..8976d7345480 100644
--- a/toolkit/components/extensions/extension-process-script.js
+++ b/toolkit/components/extensions/extension-process-script.js
@@ -343,12 +343,14 @@ ExtensionManager = {
       case "Extension:Shutdown": {
         let policy = WebExtensionPolicy.getByID(data.id);
 
-        if (extensions.has(policy)) {
-          extensions.get(policy).shutdown();
-        }
+        if (policy) {
+          if (extensions.has(policy)) {
+            extensions.get(policy).shutdown();
+          }
 
-        if (isContentProcess) {
-          policy.active = false;
+          if (isContentProcess) {
+            policy.active = false;
+          }
         }
         Services.cpmm.sendAsyncMessage("Extension:ShutdownComplete");
         break;