Bug 1290372 - _wait_for_connection() has to return early if instance is not running. r=ato

In case when the instance is not running, the method would check for the port again
and again until the timeout is reached. This extra time can be spent if the process
status is checked.

MozReview-Commit-ID: C2WAWNC5CWE

--HG--
extra : rebase_source : 2a8ab59fa60ee07b34048e95453d27aa8a2b0f3b

testing/marionette/client/marionette_driver/marionette.py

@@ -2,12 +2,14 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-import base64
 import ConfigParser
+import base64
+import datetime
 import json
 import os
 import socket
 import sys
+import time
 import traceback
 import warnings
 
@@ -571,11 +573,10 @@ class Marionette(object):
         self.timeout = timeout
         self.socket_timeout = socket_timeout
 
-        startup_timeout = startup_timeout or self.DEFAULT_STARTUP_TIMEOUT
         if self.bin:
             self.instance = self._create_instance(app, instance_args)
             self.instance.start()
-            self.raise_for_port(self.wait_for_port(timeout=startup_timeout))
+            self.raise_for_connection(timeout=startup_timeout)
 
     def _create_instance(self, app, instance_args):
         if not Marionette.is_port_available(self.port, host=self.host):
@@ -639,14 +640,55 @@ class Marionette(object):
         finally:
             s.close()
 
-    def wait_for_port(self, timeout=None):
-        timeout = timeout or self.DEFAULT_STARTUP_TIMEOUT
-        return transport.wait_for_port(self.host, self.port, timeout=timeout)
+    def _wait_for_connection(self, timeout=None):
+        """Wait until Marionette server has been created the communication socket.
+
+        :param timeout: Timeout in seconds for the server to be ready.
+
+        """
+        if timeout is None:
+            timeout = self.DEFAULT_STARTUP_TIMEOUT
+
+        runner = None
+        if self.instance is not None:
+            runner = self.instance.runner
+
+        poll_interval = 0.1
+        starttime = datetime.datetime.now()
+
+        while datetime.datetime.now() - starttime < datetime.timedelta(seconds=timeout):
+            # If the instance we want to connect to is not running return immediately
+            if runner is not None and not runner.is_running():
+                return False
+
+            sock = None
+            try:
+                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                sock.settimeout(0.5)
+                sock.connect((self.host, self.port))
+                data = sock.recv(16)
+                if ":" in data:
+                    return True
+            except socket.error:
+                pass
+            finally:
+                if sock is not None:
+                    sock.close()
+
+            time.sleep(poll_interval)
+
+        return False
 
     @do_process_check
-    def raise_for_port(self, port_obtained):
-        if not port_obtained:
-            raise socket.timeout("Timed out waiting for port {}!".format(self.port))
+    def raise_for_connection(self, timeout=None):
+        """Raise socket.timeout if no connection can be established.
+
+        :param timeout: Timeout in seconds for the server to be ready.
+
+        """
+        if not self._wait_for_connection(timeout):
+            raise socket.timeout("Timed out waiting for connection on {0}:{1}!".format(
+                self.host, self.port))
 
     @do_process_check
     def _send_message(self, name, params=None, key=None):
@@ -1032,7 +1074,7 @@ class Marionette(object):
         if not pref_exists:
             self.delete_session()
             self.instance.restart(prefs)
-            self.raise_for_port(self.wait_for_port())
+            self.raise_for_connection()
             self.start_session()
             self.reset_timeouts()
 
@@ -1125,7 +1167,7 @@ class Marionette(object):
                 self._request_in_app_shutdown("eRestart")
 
             try:
-                self.raise_for_port(self.wait_for_port())
+                self.raise_for_connection()
             except socket.timeout:
                 if self.instance.runner.returncode is not None:
                     exc, val, tb = sys.exc_info()
@@ -1135,7 +1177,7 @@ class Marionette(object):
         else:
             self.delete_session()
             self.instance.restart(clean=clean)
-            self.raise_for_port(self.wait_for_port())
+            self.raise_for_connection()
 
         self.start_session(session_id=session_id)
         self.reset_timeouts()
@@ -1178,9 +1220,9 @@ class Marionette(object):
             self.port,
             self.socket_timeout)
 
-        # Call wait_for_port() before attempting to connect in
+        # Call _wait_for_connection() before attempting to connect in
         # the event gecko hasn't started yet.
-        self.wait_for_port(timeout=timeout)
+        self._wait_for_connection(timeout=timeout)
         self.protocol, _ = self.client.connect()
 
         body = {"capabilities": desired_capabilities, "sessionId": session_id}

testing/marionette/client/marionette_driver/transport.py

@@ -2,7 +2,6 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-import datetime
 import json
 import socket
 import time
@@ -283,27 +282,3 @@ class TcpTransport(object):
 
     def __del__(self):
         self.close()
-
-
-def wait_for_port(host, port, timeout=60):
-    """Wait for the specified host/port to become available."""
-    starttime = datetime.datetime.now()
-    poll_interval = 0.1
-    while datetime.datetime.now() - starttime < datetime.timedelta(seconds=timeout):
-        sock = None
-        try:
-            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            sock.settimeout(0.5)
-            sock.connect((host, port))
-            data = sock.recv(16)
-            sock.shutdown(socket.SHUT_RDWR)
-            sock.close()
-            if ":" in data:
-                return True
-        except socket.error:
-            pass
-        finally:
-            if sock is not None:
-                sock.close()
-        time.sleep(poll_interval)
-    return False

testing/marionette/harness/marionette/tests/unit/test_marionette.py

@@ -3,12 +3,14 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import itertools
+import time
 
-from marionette_driver import errors
 from marionette.marionette_test import MarionetteTestCase
 
+from marionette_driver import errors
 
-class TestMarionetteProperties(MarionetteTestCase):
+
+class TestMarionette(MarionetteTestCase):
 
     def test_correct_test_name(self):
         """Test that the correct test name gets set."""
@@ -20,6 +22,14 @@ class TestMarionetteProperties(MarionetteTestCase):
 
         self.assertEqual(self.marionette.test_name, expected_test_name)
 
+    def test_wait_for_connection_non_existing_process(self):
+        """Test that wait_for_port doesn't run into a timeout if instance is not running."""
+        self.marionette.quit()
+        self.assertIsNotNone(self.marionette.instance.runner.returncode)
+        start_time = time.time()
+        self.assertFalse(self.marionette._wait_for_connection(timeout=5))
+        self.assertLess(time.time() - start_time, 5)
+
 
 class TestProtocol1Errors(MarionetteTestCase):
     def setUp(self):