Repository: google/mobly

Issue:
Snippet client's _is_app_running() doesn't work properly, leading to race conditions
Slow loading snippets fail because _is_app_running() returns True even when nothing is on the other side of the socket.
This leads to race conditions between the snippet app really starting and the connection attempt.

```
>>> sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
>>> print self._port
8822
>>> result = sock.connect_ex(('127.0.0.1', self._port))
>>> print result
0
And yet...
$ nc 127.0.0.1 8822
$
```

Discussion:

Patch:
diff --git a/mobly/controllers/android_device.py b/mobly/controllers/android_device.py
index 0d108a8..de89289 100644
--- a/mobly/controllers/android_device.py
+++ b/mobly/controllers/android_device.py
@@ -558,12 +558,9 @@ class AndroidDevice(object):
                     'Snippet package "%s" has already been loaded under name'
                     ' "%s".' % (package, client_name))
         host_port = utils.get_available_host_port()
-        # TODO(adorokhine): Don't assume that a free host-side port is free on
-        # the device as well. Both sides should allocate a unique port.
-        device_port = host_port
         client = snippet_client.SnippetClient(
-            package=package, port=host_port, adb_proxy=self.adb)
-        self._start_jsonrpc_client(client, host_port, device_port)
+            package=package, host_port=host_port, adb_proxy=self.adb)
+        self._start_jsonrpc_client(client)
         self._snippet_clients[name] = client
         setattr(self, name, client)
 
@@ -577,41 +574,41 @@ class AndroidDevice(object):
         If sl4a server is not started on the device, tries to start it.
         """
         host_port = utils.get_available_host_port()
-        device_port = sl4a_client.DEVICE_SIDE_PORT
-        self.sl4a = sl4a_client.Sl4aClient(self.adb)
-        self._start_jsonrpc_client(self.sl4a, host_port, device_port)
+        self.sl4a = sl4a_client.Sl4aClient(
+            host_port=host_port, adb_proxy=self.adb)
+        self._start_jsonrpc_client(self.sl4a)
 
         # Start an EventDispatcher for the current sl4a session
-        event_client = sl4a_client.Sl4aClient(self.adb)
+        event_client = sl4a_client.Sl4aClient(
+            host_port=host_port, adb_proxy=self.adb)
         event_client.connect(
-            port=host_port,
             uid=self.sl4a.uid,
             cmd=jsonrpc_client_base.JsonRpcCommand.CONTINUE)
         self.ed = event_dispatcher.EventDispatcher(event_client)
         self.ed.start()
 
-    def _start_jsonrpc_client(self, client, host_port, device_port):
+    def _start_jsonrpc_client(self, client):
         """Create a connection to a jsonrpc server running on the device.
 
         If the connection cannot be made, tries to restart it.
         """
         client.check_app_installed()
-        self.adb.tcp_forward(host_port, device_port)
+        self.adb.tcp_forward(client.host_port, client.device_port)
         try:
-            client.connect(port=host_port)
+            client.connect()
         except:
             try:
                 client.stop_app()
             except Exception as e:
                 self.log.warning(e)
             client.start_app()
-            client.connect(port=host_port)
+            client.connect()
 
     def _terminate_jsonrpc_client(self, client):
         client.closeSl4aSession()
         client.close()
         client.stop_app()
-        self.adb.forward('--remove tcp:%d' % client.port)
+        self.adb.forward('--remove tcp:%d' % client.host_port)
 
     def _is_timestamp_in_range(self, target, begin_time, end_time):
         low = mobly_logger.logline_timestamp_comparator(begin_time,
diff --git a/mobly/controllers/android_device_lib/jsonrpc_client_base.py b/mobly/controllers/android_device_lib/jsonrpc_client_base.py
index 3f7d6f5..dc6b75d 100644
--- a/mobly/controllers/android_device_lib/jsonrpc_client_base.py
+++ b/mobly/controllers/android_device_lib/jsonrpc_client_base.py
@@ -41,7 +41,6 @@ from builtins import str
 import json
 import logging
 import socket
-import sys
 import threading
 import time
 
@@ -94,15 +93,25 @@ class JsonRpcClientBase(object):
     of communication.
 
     Attributes:
-        uid: (int) The uid of this session.
+        host_port: (int) The host port of this RPC client.
+        device_port: (int) The device port of this RPC client.
         app_name: (str) The user-visible name of the app being communicated
-          with. Must be set by the superclass.
+                  with.
+        uid: (int) The uid of this session.
     """
-    def __init__(self, adb_proxy):
+
+    def __init__(self, host_port, device_port, app_name, adb_proxy):
         """
         Args:
-            adb_proxy: adb.AdbProxy, The adb proxy to use to start the app
+            host_port: (int) The host port of this RPC client.
+            device_port: (int) The device port of this RPC client.
+            app_name: (str) The user-visible name of the app being communicated
+                      with.
+            adb_proxy: (adb.AdbProxy) The adb proxy to use to start the app.
         """
+        self.host_port = host_port
+        self.device_port = device_port
+        self.app_name = app_name
         self.uid = None
         self._adb = adb_proxy
         self._client = None  # prevent close errors on connect failure
@@ -139,16 +148,6 @@ class JsonRpcClientBase(object):
         """
         raise NotImplementedError()
 
-    def _is_app_running(self):
-        """Checks if the app is currently running on an android device.
-
-        Must be implemented by subclasses.
-
-        Returns:
-            True if the app is running, False otherwise.
-        """
-        raise NotImplementedError()
-
     # Rest of the client methods.
 
     def check_app_installed(self):
@@ -161,7 +160,7 @@ class JsonRpcClientBase(object):
 
         Args:
             wait_time: float, The time to wait for the app to come up before
-                raising an error.
+                       raising an error.
 
         Raises:
             AppStartError: When the app was not able to be started.
@@ -175,8 +174,7 @@ class JsonRpcClientBase(object):
         raise AppStartError(
             '%s failed to start on %s.' % (self.app_name, self._adb.serial))
 
-    def connect(self, port, addr='localhost', uid=UNKNOWN_UID,
-                cmd=JsonRpcCommand.INIT):
+    def connect(self, uid=UNKNOWN_UID, cmd=JsonRpcCommand.INIT):
         """Opens a connection to a JSON RPC server.
 
         Opens a connection to a remote client. The connection attempt will time
@@ -185,12 +183,8 @@ class JsonRpcClientBase(object):
         as well.
 
         Args:
-            port: int, The port this client should connect to.
-            addr: str, The address this client should connect to.
             uid: int, The uid of the session to join, or UNKNOWN_UID to start a
                  new session.
-            connection_timeout: int, The time to wait for the connection to come
-                 up.
             cmd: JsonRpcCommand, The command to use for creating the connection.
 
         Raises:
@@ -200,13 +194,12 @@ class JsonRpcClientBase(object):
         """
         self._counter = self._id_counter()
         try:
-            self._conn = socket.create_connection((addr, port),
+            self._conn = socket.create_connection(('127.0.0.1', self.host_port),
                                                   _SOCKET_TIMEOUT)
             self._conn.settimeout(_SOCKET_TIMEOUT)
         except (socket.timeout, socket.error, IOError):
             logging.exception("Failed to create socket connection!")
             raise
-        self.port = port
         self._client = self._conn.makefile(mode="brw")
 
         resp = self._cmd(cmd, uid)
@@ -271,6 +264,20 @@ class JsonRpcClientBase(object):
             raise ProtocolError(ProtocolError.MISMATCHED_API_ID)
         return result['result']
 
+    def _is_app_running(self):
+        """Checks if the app is currently running on an android device.
+
+        May be overridden by subclasses with custom sanity checks.
+        """
+        running = False
+        try:
+          self.connect()
+          running = True
+        finally:
+          self.close()
+          # This 'return' squashes exceptions from connect()
+          return running
+
     def __getattr__(self, name):
         """Wrapper for python magic to turn method calls into RPC calls."""
         def rpc_call(*args):
diff --git a/mobly/controllers/android_device_lib/sl4a_client.py b/mobly/controllers/android_device_lib/sl4a_client.py
index 1842eb5..71963d0 100644
--- a/mobly/controllers/android_device_lib/sl4a_client.py
+++ b/mobly/controllers/android_device_lib/sl4a_client.py
@@ -27,13 +27,25 @@ _LAUNCH_CMD = (
 
 
 class Sl4aClient(jsonrpc_client_base.JsonRpcClientBase):
-    def __init__(self, adb_proxy):
-        super(Sl4aClient, self).__init__(adb_proxy)
-        self.app_name = 'SL4A'
+    """A client for interacting with SL4A using Mobly Snippet Lib.
+
+    See superclass documentation for a list of public attributes.
+    """
+
+    def __init__(self, host_port, adb_proxy):
+        """Initializes an Sl4aClient.
+
+        Args:
+            host_port: (int) The host port of this RPC client.
+            adb_proxy: (adb.AdbProxy) The adb proxy to use to start the app.
+        """
+        super(Sl4aClient, self).__init__(
+            host_port=host_port, device_port=DEVICE_SIDE_PORT, app_name='SL4A',
+            adb_proxy=adb_proxy)
 
     def _do_start_app(self):
         """Overrides superclass."""
-        self._adb.shell(_LAUNCH_CMD.format(DEVICE_SIDE_PORT))
+        self._adb.shell(_LAUNCH_CMD.format(self.device_port))
 
     def stop_app(self):
         """Overrides superclass."""
@@ -49,16 +61,3 @@ class Sl4aClient(jsonrpc_client_base.JsonRpcClientBase):
             if (e.ret_code == 1) and (not e.stdout) and (not e.stderr):
                 return False
             raise
-
-    def _is_app_running(self):
-        """Overrides superclass."""
-        # Grep for process with a preceding S which means it is truly started.
-        try:
-            out = self._adb.shell(
-                'ps | grep "S com.googlecode.android_scripting"').decode(
-                    'utf-8').strip()
-            return bool(out)
-        except adb.AdbError as e:
-            if (e.ret_code == 1) and (not e.stdout) and (not e.stderr):
-                return False
-            raise
diff --git a/mobly/controllers/android_device_lib/snippet_client.py b/mobly/controllers/android_device_lib/snippet_client.py
index 370d852..ebcae0f 100644
--- a/mobly/controllers/android_device_lib/snippet_client.py
+++ b/mobly/controllers/android_device_lib/snippet_client.py
@@ -15,7 +15,6 @@
 # limitations under the License.
 """JSON RPC interface to Mobly Snippet Lib."""
 import logging
-import socket
 
 from mobly.controllers.android_device_lib import adb
 from mobly.controllers.android_device_lib import jsonrpc_client_base
@@ -32,42 +31,43 @@ class Error(Exception):
 
 
 class SnippetClient(jsonrpc_client_base.JsonRpcClientBase):
-    def __init__(self, package, port, adb_proxy):
-        """Initialzies a SnippetClient.
+    """A client for interacting with snippet APKs using Mobly Snippet Lib.
+
+    See superclass documentation for a list of public attributes.
+    """
+
+    def __init__(self, package, host_port, adb_proxy):
+        """Initializes a SnippetClient.
   
         Args:
-          package: (str) The package name of the apk where the snippets are
-            defined.
-          port: (int) The port at which to start the snippet client. Note that
-            the same port will currently be used for both the device and host
-            side of the connection.
-            TODO(adorokhine): allocate a distinct free port for both sides of
-            the connection; it is not safe in general to reuse the same one.
-          adb_proxy: (adb.AdbProxy) The adb proxy to use to start the app.
+            package: (str) The package name of the apk where the snippets are
+                     defined.
+            host_port: (int) The port at which to start the snippet client. Note
+                       that the same port will currently be used for both the
+                       device and host side of the connection.
+            adb_proxy: (adb.AdbProxy) The adb proxy to use to start the app.
         """
-        super(SnippetClient, self).__init__(adb_proxy)
-        self.app_name = package
-        self._package = package
-        self._port = port
-
-    @property
-    def package(self):
-        return self._package
+        # TODO(adorokhine): Don't assume that a free host-side port is free on
+        # the device as well. Both sides should allocate a unique port.
+        super(SnippetClient, self).__init__(
+            host_port=host_port, device_port=host_port, app_name=package,
+            adb_proxy=adb_proxy)
+        self.package = package
 
     def _do_start_app(self):
         """Overrides superclass."""
-        cmd = _LAUNCH_CMD.format(self._port, self._package)
+        cmd = _LAUNCH_CMD.format(self.device_port, self.package)
         # Use info here so people know exactly what's happening here, which is
         # helpful since they need to create their own instrumentations and
         # manifest.
         logging.info('Launching snippet apk with: %s', cmd)
-        self._adb.shell(_LAUNCH_CMD.format(self._port, self._package))
+        self._adb.shell(cmd)
 
     def stop_app(self):
         """Overrides superclass."""
-        cmd = _STOP_CMD.format(self._package)
+        cmd = _STOP_CMD.format(self.package)
         logging.info('Stopping snippet apk with: %s', cmd)
-        out = self._adb.shell(_STOP_CMD.format(self._package)).decode('utf-8')
+        out = self._adb.shell(_STOP_CMD.format(self.package)).decode('utf-8')
         if 'OK (0 tests)' not in out:
             raise Error('Failed to stop existing apk. Unexpected output: %s' %
                         out)
@@ -77,22 +77,9 @@ class SnippetClient(jsonrpc_client_base.JsonRpcClientBase):
         try:
             out = self._adb.shell(
                 'pm list instrumentation | grep ^instrumentation:%s/' %
-                self._package).decode('utf-8')
+                self.package).decode('utf-8')
             return bool(out)
         except adb.AdbError as e:
             if (e.ret_code == 1) and (not e.stdout) and (not e.stderr):
                 return False
             raise
-
-    def _is_app_running(self):
-        """Overrides superclass."""
-        # While instrumentation is running, 'ps' only shows the package of the
-        # main apk. However, the main apk might be running for other reasons.
-        # Instead of grepping the process tree, this is implemented by seeing if
-        # our destination port is alive.
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        try:
-            result = sock.connect_ex(('127.0.0.1', self._port))
-            return result == 0
-        finally:
-            sock.close()
diff --git a/tools/snippet_shell.py b/tools/snippet_shell.py
index 0283f30..0991747 100755
--- a/tools/snippet_shell.py
+++ b/tools/snippet_shell.py
@@ -28,7 +28,6 @@ from __future__ import print_function
 
 import argparse
 import logging
-import sys
 
 from mobly.controllers.android_device_lib import jsonrpc_shell_base
 


Test patch:
diff --git a/tests/mobly/controllers/android_device_lib/jsonrpc_client_base_test.py b/tests/mobly/controllers/android_device_lib/jsonrpc_client_base_test.py
index ce01bad..19dc109 100755
--- a/tests/mobly/controllers/android_device_lib/jsonrpc_client_base_test.py
+++ b/tests/mobly/controllers/android_device_lib/jsonrpc_client_base_test.py
@@ -46,7 +46,9 @@ class MockSocketFile(object):
 
 class FakeRpcClient(jsonrpc_client_base.JsonRpcClientBase):
     def __init__(self):
-      super(FakeRpcClient, self).__init__(adb_proxy=None)
+        super(FakeRpcClient, self).__init__(
+            host_port=80, device_port=90, app_name='FakeRpcClient',
+            adb_proxy=None)
 
 
 class JsonRpcClientBaseTest(unittest.TestCase):
@@ -78,7 +80,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         mock_create_connection.side_effect = IOError()
         with self.assertRaises(IOError):
             client = FakeRpcClient()
-            client.connect(port=80)
+            client.connect()
 
     @mock.patch('socket.create_connection')
     def test_connect_timeout(self, mock_create_connection):
@@ -90,7 +92,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         mock_create_connection.side_effect = socket.timeout
         with self.assertRaises(socket.timeout):
             client = FakeRpcClient()
-            client.connect(port=80)
+            client.connect()
 
     @mock.patch('socket.create_connection')
     def test_handshake_error(self, mock_create_connection):
@@ -104,7 +106,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         mock_create_connection.return_value = fake_conn
         with self.assertRaises(jsonrpc_client_base.ProtocolError):
             client = FakeRpcClient()
-            client.connect(port=80)
+            client.connect()
 
     @mock.patch('socket.create_connection')
     def test_connect_handshake(self, mock_create_connection):
@@ -118,7 +120,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         mock_create_connection.return_value = fake_conn
 
         client = FakeRpcClient()
-        client.connect(port=80)
+        client.connect()
 
         self.assertEqual(client.uid, 1)
 
@@ -135,7 +137,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         mock_create_connection.return_value = fake_conn
 
         client = FakeRpcClient()
-        client.connect(port=80)
+        client.connect()
 
         self.assertEqual(client.uid, jsonrpc_client_base.UNKNOWN_UID)
 
@@ -149,7 +151,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         fake_file = self.setup_mock_socket_file(mock_create_connection)
 
         client = FakeRpcClient()
-        client.connect(port=80)
+        client.connect()
 
         fake_file.resp = None
 
@@ -169,7 +171,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         fake_file = self.setup_mock_socket_file(mock_create_connection)
 
         client = FakeRpcClient()
-        client.connect(port=80)
+        client.connect()
 
         fake_file.resp = MOCK_RESP_WITH_ERROR
 
@@ -186,7 +188,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         fake_file = self.setup_mock_socket_file(mock_create_connection)
 
         client = FakeRpcClient()
-        client.connect(port=80)
+        client.connect()
 
         fake_file.resp = (MOCK_RESP_TEMPLATE % 52).encode('utf8')
 
@@ -205,7 +207,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         fake_file = self.setup_mock_socket_file(mock_create_connection)
 
         client = FakeRpcClient()
-        client.connect(port=80)
+        client.connect()
 
         fake_file.resp = None
 
@@ -224,7 +226,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         fake_file = self.setup_mock_socket_file(mock_create_connection)
 
         client = FakeRpcClient()
-        client.connect(port=80)
+        client.connect()
 
         result = client.some_rpc(1, 2, 3)
         self.assertEquals(result, 123)
@@ -243,7 +245,7 @@ class JsonRpcClientBaseTest(unittest.TestCase):
         fake_file = self.setup_mock_socket_file(mock_create_connection)
 
         client = FakeRpcClient()
-        client.connect(port=80)
+        client.connect()
 
         for i in range(0, 10):
             fake_file.resp = (MOCK_RESP_TEMPLATE % i).encode('utf-8')