Bug 1790185 - [marionette] Use PollPromise when trying to connect to marionette server. r=webdriver-reviewers,whimboo

Differential Revision: https://phabricator.services.mozilla.com/D177134

remote/components/Marionette.sys.mjs

@@ -223,10 +223,9 @@ class MarionetteParentProcess {
 
     try {
       this.server = new lazy.TCPListener(lazy.MarionettePrefs.port);
-      this.server.start();
+      await this.server.start();
     } catch (e) {
       lazy.logger.fatal("Marionette server failed to start", e);
-      await this.uninit();
       Services.startup.quit(Ci.nsIAppStartup.eForceQuit);
       return;
     }
@@ -253,7 +252,7 @@ class MarionetteParentProcess {
 
   async uninit() {
     if (this.running) {
-      this.server.stop();
+      await this.server.stop();
       Services.obs.notifyObservers(this, NOTIFY_LISTENING);
       lazy.logger.debug("Marionette stopped listening");
 

remote/marionette/driver.sys.mjs

@@ -2857,9 +2857,9 @@ GeckoDriver.prototype._handleUserPrompts = async function() {
  * @param {boolean} cmd.parameters.value
  *     True if the server should accept new socket connections.
  */
-GeckoDriver.prototype.acceptConnections = function(cmd) {
+GeckoDriver.prototype.acceptConnections = async function(cmd) {
   lazy.assert.boolean(cmd.parameters.value);
-  this._server.acceptConnections = cmd.parameters.value;
+  await this._server.setAcceptConnections(cmd.parameters.value);
 };
 
 /**
@@ -2945,7 +2945,7 @@ GeckoDriver.prototype.quit = async function(cmd) {
     mode |= Ci.nsIAppStartup.eAttemptQuit;
   }
 
-  this._server.acceptConnections = false;
+  await this._server.setAcceptConnections(false);
   this.deleteSession();
 
   // Notify all windows that an application quit has been requested.

remote/marionette/server.sys.mjs

@@ -15,6 +15,7 @@ ChromeUtils.defineESModuleGetters(lazy, {
   Log: "chrome://remote/content/shared/Log.sys.mjs",
   MarionettePrefs: "chrome://remote/content/marionette/prefs.sys.mjs",
   Message: "chrome://remote/content/marionette/message.sys.mjs",
+  PollPromise: "chrome://remote/content/shared/Sync.sys.mjs",
   Response: "chrome://remote/content/marionette/message.sys.mjs",
 });
 
@@ -65,15 +66,30 @@ export class TCPListener {
     return new lazy.GeckoDriver(this);
   }
 
-  set acceptConnections(value) {
+  async setAcceptConnections(value) {
     if (value) {
       if (!this.socket) {
-        try {
-          const flags = KeepWhenOffline | LoopbackOnly;
-          const backlog = 1;
-          this.socket = new lazy.ServerSocket(this.port, flags, backlog);
-        } catch (e) {
-          throw new Error(`Could not bind to port ${this.port} (${e.name})`);
+        await lazy.PollPromise(
+          (resolve, reject) => {
+            try {
+              const flags = KeepWhenOffline | LoopbackOnly;
+              const backlog = 1;
+              this.socket = new lazy.ServerSocket(this.port, flags, backlog);
+              resolve();
+            } catch (e) {
+              lazy.logger.debug(
+                `Could not bind to port ${this.port} (${e.name})`
+              );
+              reject();
+            }
+          },
+          { interval: 250, timeout: 5000 }
+        );
+
+        // Since PollPromise doesn't throw when timeout expires,
+        // we can end up in the situation when the socket is undefined.
+        if (!this.socket) {
+          throw new Error(`Could not bind to port ${this.port}`);
         }
 
         this.port = this.socket.port;
@@ -97,24 +113,24 @@ export class TCPListener {
    * The marionette.port preference will be populated with the value
    * of {@link #port}.
    */
-  start() {
+  async start() {
     if (this.alive) {
       return;
     }
 
     // Start socket server and listening for connection attempts
-    this.acceptConnections = true;
+    await this.setAcceptConnections(true);
     lazy.MarionettePrefs.port = this.port;
     this.alive = true;
   }
 
-  stop() {
+  async stop() {
     if (!this.alive) {
       return;
     }
 
     // Shutdown server socket, and no longer listen for new connections
-    this.acceptConnections = false;
+    await this.setAcceptConnections(false);
     this.alive = false;
   }