Removing flake in TemporaryHTTPServer.

Note that I think ReplayServer (wpr_server.py) could still potentially be flaky since it makes use of webpagereplay.py (which hardcodes ports). We may eventually want to subclass it or something to make it get dynamic ports like we do elsewhere. In fact, I ran into the problem today while working on this patch - if a test doesn't exit properly, the server may be left up blocking these hardcoded ports. I've made a new bug for this, since I think it's a bit out of the scope of this issue (crbug.com/157459).

BUG=154229

Review URL: https://chromiumcodereview.appspot.com/11231087

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@168902 0039d316-1c4b-4281-b951-d872f2087c98

5 files changed, 80 insertions, 14 deletions

diff --git a/tools/telemetry/telemetry/cros_browser_backend.py b/tools/telemetry/telemetry/cros_browser_backend.py
index 3ce3cd86..c66c88e 100644
--- a/tools/telemetry/telemetry/cros_browser_backend.py
+++ b/tools/telemetry/telemetry/cros_browser_backend.py
@@ -5,9 +5,9 @@ import logging
 import os
 import socket
 import subprocess
-import time
 
 from telemetry import browser_backend
+from telemetry import util
 
 class CrOSBrowserBackend(browser_backend.BrowserBackend):
   def __init__(self, browser_type, options, is_content_shell, cri):
@@ -116,7 +116,7 @@ class CrOSBrowserBackend(browser_backend.BrowserBackend):
 
   def CreateForwarder(self, *ports):
     assert self._cri
-    return SSHForwarder(self._cri, 'R', *[(port, port) for port in ports])
+    return SSHForwarder(self._cri, 'R', *ports)
 
   def _RestartUI(self):
     if self._cri:
@@ -132,19 +132,30 @@ class SSHForwarder(object):
     self._proc = None
     self._host_port = ports[0][0]
 
+    port_pairs = []
+
+    for port in ports:
+      if port[1] is None:
+        port_pairs.append((port[0], cri.GetRemotePort()))
+      else:
+        port_pairs.append(port)
+
+    if forwarding_flag == 'R':
+      self._device_port = port_pairs[0][0]
+    else:
+      self._device_port = port_pairs[0][1]
+
     self._proc = subprocess.Popen(
       cri.FormSSHCommandLine(
         ['sleep', '999999999'],
         ['-%s%i:localhost:%i' % (forwarding_flag, from_port, to_port)
-        for from_port, to_port in ports]),
+        for from_port, to_port in port_pairs]),
       stdout=subprocess.PIPE,
       stderr=subprocess.PIPE,
       stdin=subprocess.PIPE,
       shell=False)
 
-    # TODO(nduca): How do we wait for the server to come up in a
-    # robust way?
-    time.sleep(1.5)
+    util.WaitFor(lambda: cri.IsHTTPServerRunningOnPort(self._device_port), 60)
 
   @property
   def url(self):
diff --git a/tools/telemetry/telemetry/cros_interface.py b/tools/telemetry/telemetry/cros_interface.py
index 43329fb..91a8574 100644
--- a/tools/telemetry/telemetry/cros_interface.py
+++ b/tools/telemetry/telemetry/cros_interface.py
@@ -12,8 +12,6 @@ import tempfile
 
 from telemetry import util
 
-_next_remote_port = 9224
-
 # TODO(nduca): This whole file is built up around making individual ssh calls
 # for each operation. It really could get away with a single ssh session built
 # around pexpect, I suspect, if we wanted it to be faster. But, this was
@@ -373,7 +371,24 @@ class CrOSInterface(object):
     return stdout
 
   def GetRemotePort(self):
-    global _next_remote_port
-    port = _next_remote_port
-    _next_remote_port += 1
-    return port
+    netstat = self.GetAllCmdOutput(['netstat', '-ant'])
+    netstat = netstat[0].split('\n')
+    ports_in_use = []
+
+    for line in netstat[2:]:
+      if not line:
+        continue
+      address_in_use = line.split()[3]
+      port_in_use = address_in_use.split(':')[-1]
+      ports_in_use.append(int(port_in_use))
+
+    return sorted(ports_in_use)[-1] + 1
+
+  def IsHTTPServerRunningOnPort(self, port):
+    wget_output = self.GetAllCmdOutput(
+        ['wget', 'localhost:%i' % (port), '-T1', '-t1'])
+
+    if 'Connection refused' in wget_output[1]:
+      return False
+
+    return True
diff --git a/tools/telemetry/telemetry/cros_interface_unittest.py b/tools/telemetry/telemetry/cros_interface_unittest.py
index a445a5c..4609b65 100644
--- a/tools/telemetry/telemetry/cros_interface_unittest.py
+++ b/tools/telemetry/telemetry/cros_interface_unittest.py
@@ -6,7 +6,9 @@
 # actually talking to the device. This would improve our coverage quite
 # a bit.
 import unittest
+import socket
 
+from telemetry import cros_browser_backend
 from telemetry import cros_interface
 from telemetry import options_for_unittests
 from telemetry import run_tests
@@ -109,3 +111,38 @@ class CrOSInterfaceTest(unittest.TestCase):
 
     self.assertTrue(cri.IsServiceRunning('openssh-server'))
 
+  @run_tests.RequiresBrowserOfType('cros-chrome')
+  def testGetRemotePortAndIsHTTPServerRunningOnPort(self):
+    remote = options_for_unittests.Get().cros_remote
+    cri = cros_interface.CrOSInterface(
+      remote,
+      options_for_unittests.Get().cros_ssh_identity)
+
+    # Create local server.
+    sock = socket.socket()
+    sock.bind(('', 0))
+    port = sock.getsockname()[1]
+    sock.listen(0)
+
+    # Get remote port and ensure that it was unused.
+    remote_port = cri.GetRemotePort()
+    self.assertFalse(cri.IsHTTPServerRunningOnPort(remote_port))
+
+    # Forward local server's port to remote device's remote_port.
+    forwarder = cros_browser_backend.SSHForwarder(
+        cri, 'R', (remote_port, port))
+
+    # At this point, remote device should be able to connect to local server.
+    self.assertTrue(cri.IsHTTPServerRunningOnPort(remote_port))
+
+    # Next remote port shouldn't be the same as remote_port, since remote_port
+    # is now in use.
+    self.assertTrue(cri.GetRemotePort() != remote_port)
+
+    # Close forwarder and local server ports.
+    forwarder.Close()
+    sock.close()
+
+    # Device should no longer be able to connect to remote_port since it is no
+    # longer in use.
+    self.assertFalse(cri.IsHTTPServerRunningOnPort(remote_port))
diff --git a/tools/telemetry/telemetry/temporary_http_server.py b/tools/telemetry/telemetry/temporary_http_server.py
index abd7300..f6ab8c9 100644
--- a/tools/telemetry/telemetry/temporary_http_server.py
+++ b/tools/telemetry/telemetry/temporary_http_server.py
@@ -31,7 +31,7 @@ class TemporaryHTTPServer(object):
         cwd=self._path,
         stdout=self._devnull, stderr=self._devnull)
 
-    self._forwarder = browser_backend.CreateForwarder(self._host_port)
+    self._forwarder = browser_backend.CreateForwarder((self._host_port, None))
 
     def IsServerUp():
       return not socket.socket().connect_ex(('localhost', self._host_port))
diff --git a/tools/telemetry/telemetry/wpr_server.py b/tools/telemetry/telemetry/wpr_server.py
index 9789d95..ec8dd05 100644
--- a/tools/telemetry/telemetry/wpr_server.py
+++ b/tools/telemetry/telemetry/wpr_server.py
@@ -21,8 +21,11 @@ class ReplayServer(object):
     self._web_page_replay = None
     self._is_record_mode = is_record_mode
 
+    # Note: This can cause flake if server doesn't shut down properly and keeps
+    # ports tied up. See crbug.com/157459.
     self._forwarder = browser_backend.CreateForwarder(
-        webpagereplay.HTTP_PORT, webpagereplay.HTTPS_PORT)
+        (webpagereplay.HTTP_PORT, webpagereplay.HTTP_PORT),
+        (webpagereplay.HTTPS_PORT, webpagereplay.HTTPS_PORT))
 
     options = []
     if self._is_record_mode: