Avoid Python test server port-allocation races

test/py/RunClientServer.py

@@ -23,9 +23,9 @@
 import copy
 import os
 import signal
-import socket
 import subprocess
 import sys
+import tempfile
 import time
 from optparse import OptionParser
 
@@ -103,22 +103,13 @@ def runScriptTest(libdir, genbase, genpydir, script):
         raise Exception("Script subprocess failed, retcode=%d, args: %s" % (ret, ' '.join(script_args)))
 
 
-def pick_unused_port():
-    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    sock.bind(('127.0.0.1', 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    return port
-
-
 def runServiceTest(libdir, genbase, genpydir, server_class, proto, port, use_zlib, use_ssl, verbose):
     env = setup_pypath(libdir, os.path.join(genbase, genpydir))
     # Build command line arguments
     server_args = [sys.executable, relfile('TestServer.py')]
     cli_args = [sys.executable, relfile('TestClient.py')]
     for which in (server_args, cli_args):
         which.append('--protocol=%s' % proto)  # accel, binary, compact or json
-        which.append('--port=%d' % port)  # default to 9090
         if use_zlib:
             which.append('--zlib')
         if use_ssl:
@@ -127,6 +118,10 @@ def runServiceTest(libdir, genbase, genpydir, server_class, proto, port, use_zli
             which.append('-q')
         if verbose == 2:
             which.append('-v')
+    fd, port_file = tempfile.mkstemp(prefix='thrift-test-port-')
+    os.close(fd)
+    server_args.append('--port=%d' % port)
+    server_args.append('--port-file=%s' % port_file)
     # server-specific option to select server class
     server_args.append(server_class)
     # client-specific cmdline options
@@ -145,7 +140,7 @@ def runServiceTest(libdir, genbase, genpydir, server_class, proto, port, use_zli
             popen_kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP
     else:
         popen_kwargs['start_new_session'] = True
-    serverproc = subprocess.Popen(server_args, **popen_kwargs)
+    serverproc = None
 
     def ensureServerAlive():
         if serverproc.poll() is not None:
@@ -154,28 +149,22 @@ def ensureServerAlive():
             raise Exception('Server subprocess %s died, args: %s'
                             % (server_class, ' '.join(server_args)))
 
-    # Wait for the server to start accepting connections on the given port.
-    sleep_time = 0.1  # Seconds
-    max_attempts = 100
-    attempt = 0
-    while True:
-        sock4 = socket.socket()
-        sock6 = socket.socket(socket.AF_INET6)
-        try:
-            if sock4.connect_ex(('127.0.0.1', port)) == 0 \
-                    or sock6.connect_ex(('::1', port)) == 0:
-                break
-            attempt += 1
-            if attempt >= max_attempts:
-                raise Exception("TestServer not ready on port %d after %.2f seconds"
-                                % (port, sleep_time * attempt))
+    try:
+        serverproc = subprocess.Popen(server_args, **popen_kwargs)
+
+        deadline = time.monotonic() + 10
+        while True:
             ensureServerAlive()
-            time.sleep(sleep_time)
-        finally:
-            sock4.close()
-            sock6.close()
+            try:
+                with open(port_file) as input_file:
+                    port = int(input_file.read())
+                break
+            except (OSError, ValueError):
+                if time.monotonic() >= deadline:
+                    raise Exception("TestServer did not report its port after 10 seconds")
+                time.sleep(0.1)
 
-    try:
+        cli_args.append('--port=%d' % port)
         if verbose > 0:
             print('Testing client: %s' % (' '.join(cli_args)))
         ret = subprocess.call(cli_args, env=env)
@@ -185,28 +174,32 @@ def ensureServerAlive():
             print('PY_GEN: %s' % genpydir, file=sys.stderr)
             raise Exception("Client subprocess failed, retcode=%d, args: %s" % (ret, ' '.join(cli_args)))
     finally:
-        # check that server didn't die, but still attempt cleanup
         cleanup_exc = None
+        if serverproc is not None:
+            try:
+                ensureServerAlive()
+            except Exception as exc:
+                cleanup_exc = exc
+            extra_sleep = EXTRA_DELAY.get(server_class, 0)
+            if extra_sleep > 0 and verbose > 0:
+                print('Giving %s (proto=%s,zlib=%s,ssl=%s) an extra %d seconds for child'
+                      'processes to terminate via alarm'
+                      % (server_class, proto, use_zlib, use_ssl, extra_sleep))
+                time.sleep(extra_sleep)
+            sig = signal.SIGKILL if platform.system() != 'Windows' else signal.SIGABRT
+            try:
+                if platform.system() == 'Windows':
+                    os.kill(serverproc.pid, sig)
+                else:
+                    # POSIX: kill the whole process group to reap forked children.
+                    os.killpg(serverproc.pid, sig)
+            except OSError:
+                pass
+            serverproc.wait()
         try:
-            ensureServerAlive()
-        except Exception as exc:
-            cleanup_exc = exc
-        extra_sleep = EXTRA_DELAY.get(server_class, 0)
-        if extra_sleep > 0 and verbose > 0:
-            print('Giving %s (proto=%s,zlib=%s,ssl=%s) an extra %d seconds for child'
-                  'processes to terminate via alarm'
-                  % (server_class, proto, use_zlib, use_ssl, extra_sleep))
-            time.sleep(extra_sleep)
-        sig = signal.SIGKILL if platform.system() != 'Windows' else signal.SIGABRT
-        try:
-            if platform.system() == 'Windows':
-                os.kill(serverproc.pid, sig)
-            else:
-                # POSIX: kill the whole process group to reap forked children.
-                os.killpg(serverproc.pid, sig)
+            os.unlink(port_file)
         except OSError:
             pass
-        serverproc.wait()
         if cleanup_exc:
             raise cleanup_exc
 
@@ -248,8 +241,8 @@ def run(self, conf, test_count):
         if self.verbose > 0:
             print('\nTest run #%d:  (includes %s) Server=%s,  Proto=%s,  zlib=%s,  SSL=%s'
                   % (test_count, genpydir, try_server, try_proto, with_zlib, with_ssl))
-        port = self.port if self.port else pick_unused_port()
-        runServiceTest(self.libdir, self.genbase, genpydir, try_server, try_proto, port, with_zlib, with_ssl, self.verbose)
+        runServiceTest(self.libdir, self.genbase, genpydir, try_server, try_proto,
+                       self.port, with_zlib, with_ssl, self.verbose)
         if self.verbose > 0:
             print('OK: Finished (includes %s)  %s / %s proto / zlib=%s / SSL=%s.   %d combinations tested.'
                   % (genpydir, try_server, try_proto, with_zlib, with_ssl, test_count))
@@ -285,8 +278,9 @@ def run_all_tests(self):
                             if self.verbose > 0:
                                 print('\nTest run #%d:  (includes %s) Server=%s,  Proto=%s,  zlib=%s,  SSL=%s'
                                       % (test_count, genpydir, try_server, try_proto, with_zlib, with_ssl))
-                            port = self.port if self.port else pick_unused_port()
-                            runServiceTest(self.libdir, self.genbase, genpydir, try_server, try_proto, port, with_zlib, with_ssl, self.verbose)
+                            runServiceTest(self.libdir, self.genbase, genpydir,
+                                           try_server, try_proto, self.port,
+                                           with_zlib, with_ssl, self.verbose)
                             if self.verbose > 0:
                                 print('OK: Finished (includes %s)  %s / %s proto / zlib=%s / SSL=%s.   %d combinations tested.'
                                       % (genpydir, try_server, try_proto, with_zlib, with_ssl, test_count))

test/py/TestServer.py

@@ -21,19 +21,59 @@
 import logging
 import os
 import signal
+import socketserver
 import ssl
 import sys
 import time
+from http.server import HTTPServer
 from optparse import OptionParser
 
 from util import local_libpath
 sys.path.insert(0, local_libpath())
 from thrift.protocol import TProtocol, TProtocolDecorator
 from thrift.Thrift import TException
+from thrift.transport import TSocket, TSSLSocket
 
 SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))
 
 
+def report_bound_port(port_file, handle):
+    if port_file:
+        address = handle.getsockname()
+        if not isinstance(address, tuple) or len(address) < 2 \
+                or not isinstance(address[1], int):
+            raise ValueError('port reporting requires a TCP socket')
+        with open(port_file, 'w') as output:
+            output.write(str(address[1]))
+
+
+class TPortReportingServerSocket(TSocket.TServerSocket):
+    def __init__(self, host, port, unix_socket, port_file):
+        super(TPortReportingServerSocket, self).__init__(host, port, unix_socket)
+        self.port_file = port_file
+
+    def listen(self):
+        super(TPortReportingServerSocket, self).listen()
+        report_bound_port(self.port_file, self.handle)
+
+
+class TPortReportingSSLServerSocket(TSSLSocket.TSSLServerSocket):
+    def __init__(self, host, port, port_file, **kwargs):
+        super(TPortReportingSSLServerSocket, self).__init__(host, port, **kwargs)
+        self.port_file = port_file
+
+    def listen(self):
+        super(TPortReportingSSLServerSocket, self).listen()
+        report_bound_port(self.port_file, self.handle)
+
+
+class TLocalHTTPServer(HTTPServer):
+    def server_bind(self):
+        socketserver.TCPServer.server_bind(self)
+        self.server_name = 'localhost'
+        self.server_port = self.server_address[1]
+
+
 class TestHandler(object):
     def __init__(self, options):
         self.options = options
@@ -306,36 +346,53 @@ def main(options):
 
     # Handle THttpServer as a special case
     if server_type == 'THttpServer':
+        # Avoid HTTPServer's external hostname lookup for local ephemeral tests.
+        http_host = '127.0.0.1' if options.port_file else ''
+        http_server_class = TLocalHTTPServer if options.port_file else HTTPServer
         if options.ssl:
             __certfile = os.path.join(os.path.dirname(SCRIPT_DIR), "keys", "server.crt")
             __keyfile = os.path.join(os.path.dirname(SCRIPT_DIR), "keys", "server.key")
-            server = THttpServer.THttpServer(processor, ('', options.port), pfactory, cert_file=__certfile, key_file=__keyfile)
+            server = THttpServer.THttpServer(
+                processor,
+                (http_host, options.port),
+                pfactory,
+                cert_file=__certfile,
+                key_file=__keyfile,
+                server_class=http_server_class,
+            )
         else:
-            server = THttpServer.THttpServer(processor, ('', options.port), pfactory)
+            server = THttpServer.THttpServer(
+                processor,
+                (http_host, options.port),
+                pfactory,
+                server_class=http_server_class,
+            )
+        report_bound_port(options.port_file, server.httpd.socket)
         server.serve()
         sys.exit(0)
 
     # set up server transport and transport factory
 
     host = None
     if options.ssl:
-        from thrift.transport import TSSLSocket
         keys_dir = os.path.join(os.path.dirname(SCRIPT_DIR), 'keys')
         ca_certs = os.path.join(keys_dir, 'client.pem')
         certfile = os.path.join(keys_dir, 'server.crt')
         keyfile = os.path.join(keys_dir, 'server.key')
         ssl_version = getattr(ssl, 'PROTOCOL_TLS_SERVER', ssl.PROTOCOL_TLSv1)
-        transport = TSSLSocket.TSSLServerSocket(
+        transport = TPortReportingSSLServerSocket(
             host,
             options.port,
+            options.port_file,
             certfile=certfile,
             keyfile=keyfile,
             ca_certs=ca_certs,
             cert_reqs=ssl.CERT_REQUIRED,
             ssl_version=ssl_version,
         )
     else:
-        transport = TSocket.TServerSocket(host, options.port, options.domain_socket)
+        transport = TPortReportingServerSocket(
+            host, options.port, options.domain_socket, options.port_file)
     tfactory = TTransport.TBufferedTransportFactory()
     if options.trans == 'buffered':
         tfactory = TTransport.TBufferedTransportFactory()
@@ -403,6 +460,8 @@ def exit_gracefully(signum, frame):
                       help='include this directory to sys.path for locating generated code')
     parser.add_option("--port", type="int", dest="port",
                       help="port number for server to listen on")
+    parser.add_option("--port-file", dest="port_file",
+                      help="write the bound TCP port to this file")
     parser.add_option("--zlib", action="store_true", dest="zlib",
                       help="use zlib wrapper for compressed transport")
     parser.add_option("--ssl", action="store_true", dest="ssl",
@@ -433,7 +492,6 @@ def exit_gracefully(signum, frame):
     from thrift.TMultiplexedProcessor import TMultiplexedProcessor
     from thrift.transport import THeaderTransport
     from thrift.transport import TTransport
-    from thrift.transport import TSocket
     from thrift.transport import TZlibTransport
     from thrift.protocol import TBinaryProtocol
     from thrift.protocol import TCompactProtocol