benchmark: serve files via twisted.web.static

splash/benchmark/benchmark.py

@@ -11,7 +11,6 @@
 import logging
 import os
 import random
-import shutil
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
 from glob import glob
 from multiprocessing.pool import ThreadPool
@@ -68,7 +67,7 @@ def make_render_png_lua_req(splash, params):
 PORT = 8806
 #: Combinations of width & height to test.
 WIDTH_HEIGHT = [(None, None), (500, None), (None, 500), (500, 500)]
-#: Splash log filename.
+#: Splash log filename (set to None to put it to stdout).
 SPLASH_LOG = 'splash.log'
 #: This script is used to collect maxrss & cpu time from splash process.
 GET_PERF_STATS_SCRIPT = """
@@ -172,7 +171,7 @@ def main():
                         '--disable-xvfb',
                         '--max-timeout=600'])
 
-    with splash, serve_files(PORT, args.sites_dir):
+    with splash, serve_files(port=PORT, directory=args.sites_dir):
         start_time = time()
         results = parallel_map(invoke_request, generate_requests(splash, args),
                                args.thread_count)

splash/benchmark/file_server.py

@@ -1,44 +1,56 @@
 #!/usr/bin/env python
 
-"""
-Simple static file server.
-"""
+"""Simple static file server."""
 
 import argparse
 import os
-import SimpleHTTPServer
-import SocketServer
 import subprocess
+import time
 from contextlib import contextmanager
 
+from twisted.internet import reactor
+from twisted.web.server import Site
+from twisted.web.static import File
 
-parser = argparse.ArgumentParser(description=__doc__)
-parser.add_argument('port', type=int, help='Port number to listen at')
-parser.add_argument('directory', type=str, help='Directory to serve')
+import requests
 
-
-class ReusingTCPServer(SocketServer.TCPServer):
-    allow_reuse_address = True
-
-
-class RequestHandler(SimpleHTTPServer.SimpleHTTPRequestHandler):
-    def address_string(self):
-        return "fileserver"
+parser = argparse.ArgumentParser("")
+parser.add_argument('--port', type=int)
+parser.add_argument('--directory', help='Directory to be served')
 
 
 @contextmanager
-def serve_files(port, directory):
-    """Serve files from current directory statically in a subprocess."""
-    site_server = subprocess.Popen(['python', '-m', __name__,
-                                    str(port), directory])
+def serve_files(port, directory, logfile=None):
+    """Serve files from specified directory statically in a subprocess."""
+    command = ['twistd',
+               '-n',    # don't daemonize
+               'web',   # start web component
+               '--port', str(int(port)),
+               '--path', os.path.abspath(directory), ]
+    if logfile is not None:
+        command += ['--logfile', logfile]
+    site_server = subprocess.Popen(command)
     try:
+        # It might take some time to bring up the server, wait for up to 10s.
+        for i in xrange(100):
+            try:
+                requests.get('http://localhost:%d' % port)
+            except requests.ConnectionError:
+                time.sleep(0.1)
+            else:
+                break
         yield
     finally:
         site_server.terminate()
 
 
-if __name__ == '__main__':
+def main():
     args = parser.parse_args()
-    os.chdir(args.directory)
-    server = ReusingTCPServer(("", args.port), RequestHandler)
-    server.serve_forever()
+    resource = File(os.path.abspath(args.directory))
+    site = Site(resource)
+    reactor.listenTCP(args.port, site)
+    reactor.run()
+
+
+if __name__ == '__main__':
+    main()