Backout bug 1182681 for web-platform-test failures on a CLOSED TREE

testing/marionette/driver/marionette_driver/marionette.py

@@ -631,7 +631,8 @@ class Marionette(object):
         self.client = MarionetteTransport(
             self.host,
             self.port,
-            self.socket_timeout)
+            self.socket_timeout,
+            instance=self.instance)
 
         if emulator:
             if busybox:
@@ -689,12 +690,6 @@ class Marionette(object):
 
         try:
             response = self.client.send(message)
-        except IOError:
-            if self.instance and not hasattr(self.instance, 'detached'):
-                # If we've launched the binary we've connected to, wait
-                # for it to shut down.
-                returncode = self.instance.runner.wait()
-                raise IOError("process died with returncode %d" % returncode)
         except socket.timeout:
             self.session = None
             self.window = None

testing/marionette/transport/marionette_transport/transport.py

@@ -20,7 +20,7 @@ class MarionetteTransport(object):
     max_packet_length = 4096
     connection_lost_msg = "Connection to Marionette server is lost. Check gecko.log (desktop firefox) or logcat (b2g) for errors."
 
-    def __init__(self, addr, port, socket_timeout=360.0):
+    def __init__(self, addr, port, socket_timeout=360.0, instance=None):
         self.addr = addr
         self.port = port
         self.socket_timeout = socket_timeout
@@ -28,6 +28,7 @@ class MarionetteTransport(object):
         self.traits = None
         self.applicationType = None
         self.actor = 'root'
+        self.instance = instance
 
     def _recv_n_bytes(self, n):
         """ Convenience method for receiving exactly n bytes from
@@ -52,13 +53,16 @@ class MarionetteTransport(object):
         bytes_to_recv = 10
         while time.time() - now < self.socket_timeout:
             try:
-                data = self.sock.recv(bytes_to_recv)
-                response += data
+                response += self.sock.recv(bytes_to_recv)
             except socket.timeout:
                 pass
-            else:
-                if not data:
-                    raise IOError(self.connection_lost_msg)
+            if self.instance and not hasattr(self.instance, 'detached'):
+                # If we've launched the binary we've connected to, make
+                # sure it hasn't died.
+                poll = self.instance.runner.process_handler.proc.poll()
+                if poll is not None:
+                    # process isn't alive
+                    raise IOError("process has died with return code %d" % poll)
             sep = response.find(':')
             if sep > -1:
                 length = response[0:sep]
@@ -66,7 +70,7 @@ class MarionetteTransport(object):
                 if len(remaining) == int(length):
                     return json.loads(remaining)
                 bytes_to_recv = int(length) - len(remaining)
-        raise socket.timeout('connection timed out after %d s' % self.socket_timeout)
+        raise IOError(self.connection_lost_msg)
 
     def connect(self):
         """ Connect to the server and process the hello message we expect