Bug 1537763 - [mozproxy] Add a record mode; r=tarek

Differential Revision: https://phabricator.services.mozilla.com/D84399

python/mozperftest/mozperftest/system/proxy.py

@@ -5,11 +5,12 @@ import json
 import os
 import re
 import signal
+import tempfile
 import threading
 
 from mozlog import get_proxy_logger
 from mozperftest.layers import Layer
-from mozperftest.utils import install_package
+from mozperftest.utils import download_file, install_package
 from mozprocess import ProcessHandler
 
 
@@ -63,9 +64,23 @@ class ProxyRunner(Layer):
     name = "proxy"
     activated = False
 
+    arguments = {
+        "record": {
+            "type": str,
+            "default": None,
+            "help": "Generate a recording of the network requests during this test.",
+        },
+        "replay": {
+            "type": str,
+            "default": None,
+            "help": "A generated recording to play back during this test.",
+        },
+    }
+
     def __init__(self, env, mach_cmd):
         super(ProxyRunner, self).__init__(env, mach_cmd)
         self.proxy = None
+        self.tmpdir = None
 
     def setup(self):
         # Install mozproxy and its vendored deps.
@@ -77,18 +92,33 @@ class ProxyRunner(Layer):
     def run(self, metadata):
         self.metadata = metadata
 
+        replay_file = self.get_arg("replay")
+        if replay_file is not None and replay_file.startswith("http"):
+            self.tmpdir = tempfile.TemporaryDirectory()
+            target = os.path.join(self.tmpdir.name, "recording.dump")
+            self.info("Downloading %s" % replay_file)
+            download_file(replay_file, target)
+            replay_file = target
+
         self.info("Setting up the proxy")
         # replace with artifacts
+        command = [
+            "mozproxy",
+            "--local",
+            "--binary=" + self.mach_cmd.get_binary_path(),
+            "--topsrcdir=" + self.mach_cmd.topsrcdir,
+            "--objdir=" + self.mach_cmd.topobjdir,
+        ]
+        if self.get_arg("record"):
+            command.extend(["--record", self.get_arg("record")])
+        elif replay_file:
+            command.append(replay_file)
+        else:
+            command.append(os.path.join(HERE, "example.dump"))
+
         self.output_handler = OutputHandler()
         self.proxy = ProcessHandler(
-            [
-                "mozproxy",
-                "--local",
-                "--binary=" + self.mach_cmd.get_binary_path(),
-                "--topsrcdir=" + self.mach_cmd.topsrcdir,
-                "--objdir=" + self.mach_cmd.topobjdir,
-                os.path.join(HERE, "example.dump"),
-            ],
+            command,
             processOutputLine=self.output_handler,
             onFinish=self.output_handler.finished,
         )
@@ -115,8 +145,21 @@ class ProxyRunner(Layer):
         return metadata
 
     def teardown(self):
+        err = None
         if self.proxy is not None:
-            kill_signal = getattr(signal, "CTRL_BREAK_EVENT", signal.SIGINT)
-            os.kill(self.proxy.pid, kill_signal)
-            self.proxy.wait()
+            returncode = self.proxy.wait(0)
+            if returncode is not None:
+                err = ValueError(
+                    "mozproxy terminated early with return code %d" % returncode
+                )
+            else:
+                kill_signal = getattr(signal, "CTRL_BREAK_EVENT", signal.SIGINT)
+                os.kill(self.proxy.pid, kill_signal)
+                self.proxy.wait()
             self.proxy = None
+        if self.tmpdir is not None:
+            self.tmpdir.cleanup()
+            self.tmpdir = None
+
+        if err:
+            raise err

python/mozperftest/mozperftest/tests/test_proxy.py

@@ -2,6 +2,10 @@
 import mozunit
 import os
 import pytest
+import shutil
+import tempfile
+import time
+from unittest import mock
 
 from mozbuild.base import MozbuildObject
 from mozperftest.tests.support import get_running_env
@@ -9,6 +13,7 @@ from mozperftest.environment import SYSTEM
 from mozperftest.utils import install_package, silence
 
 here = os.path.abspath(os.path.dirname(__file__))
+example_dump = os.path.join(here, "..", "system", "example.dump")
 
 
 @pytest.fixture(scope="module")
@@ -23,13 +28,48 @@ def install_mozproxy():
     return build
 
 
-def test_proxy(install_mozproxy):
+def mock_download_file(url, dest):
+    shutil.copyfile(example_dump, dest)
+
+
+def test_replay(install_mozproxy):
     mach_cmd, metadata, env = get_running_env(proxy=True)
     system = env.layers[SYSTEM]
+    env.set_arg("proxy-replay", example_dump)
 
     # XXX this will run for real, we need to mock HTTP calls
     with system as proxy, silence():
         proxy(metadata)
+        # Give mitmproxy a bit of time to start up so we can verify that it's
+        # actually running before we tear things down.
+        time.sleep(5)
+
+
+@mock.patch("mozperftest.system.proxy.download_file", mock_download_file)
+def test_replay_url(install_mozproxy):
+    mach_cmd, metadata, env = get_running_env(proxy=True)
+    system = env.layers[SYSTEM]
+    env.set_arg("proxy-replay", "http://example.dump")
+
+    # XXX this will run for real, we need to mock HTTP calls
+    with system as proxy, silence():
+        proxy(metadata)
+        # Give mitmproxy a bit of time to start up so we can verify that it's
+        # actually running before we tear things down.
+        time.sleep(5)
+
+
+def test_record(install_mozproxy):
+    mach_cmd, metadata, env = get_running_env(proxy=True)
+    system = env.layers[SYSTEM]
+    with tempfile.TemporaryDirectory() as tmpdir:
+        recording = os.path.join(tmpdir, "recording.dump")
+        env.set_arg("proxy-record", recording)
+
+        # XXX this will run for real, we need to mock HTTP calls
+        with system as proxy, silence():
+            proxy(metadata)
+        assert os.path.exists(recording)
 
 
 if __name__ == "__main__":

testing/mozbase/mozproxy/mozproxy/backends/mitm/mitm.py

@@ -219,6 +219,15 @@ class Mitmproxy(Playback):
     def stop(self):
         self.stop_mitmproxy_playback()
 
+    def wait(self, timeout=1):
+        """Wait until the mitmproxy process has terminated."""
+        # We wait using this method to allow Windows to respond to the Ctrl+Break
+        # signal so that we can exit cleanly from the command-line driver.
+        while True:
+            returncode = self.mitmproxy_proc.wait(timeout)
+            if returncode is not None:
+                return returncode
+
     def start_mitmproxy_playback(self, mitmdump_path, browser_path):
         """Startup mitmproxy and replay the specified flow file"""
         if self.mitmproxy_proc is not None:
@@ -235,10 +244,12 @@ class Mitmproxy(Playback):
         # add proxy host and port options
         command.extend(["--listen-host", self.host, "--listen-port", str(self.port)])
 
-        if "playback_tool_args" in self.config:
+        if self.config.get("playback_tool_args"):
             LOG.info("Staring Proxy using provided command line!")
             command.extend(self.config["playback_tool_args"])
-        elif "playback_files" in self.config:
+        elif self.config.get("playback_record"):
+            command.extend(["-w", self.config.get("playback_record")])
+        elif self.config.get("playback_files"):
             script = os.path.join(
                 os.path.dirname(os.path.realpath(__file__)),
                 "scripts",

testing/mozbase/mozproxy/mozproxy/backends/mitm/scripts/alternate-server-replay.py

@@ -157,7 +157,7 @@ class AlternateServerPlayback:
                     _PROTO.update(recording_info["http_protocol"])
         except Exception as e:
             ctx.log.error("Could not load recording file! Stopping playback process!")
-            ctx.log.info(e)
+            ctx.log.error(str(e))
             ctx.master.shutdown()
 
     def _hash(self, flow):

testing/mozbase/mozproxy/mozproxy/driver.py

@@ -7,7 +7,6 @@ import argparse
 import os
 import signal
 import sys
-import time
 
 import mozinfo
 import mozlog.commandline
@@ -15,11 +14,16 @@ import mozlog.commandline
 from . import get_playback
 from .utils import LOG, TOOLTOOL_PATHS
 
+EXIT_SUCCESS = 0
+EXIT_EARLY_TERMINATE = 3
+EXIT_EXCEPTION = 4
+
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--local", action="store_true",
                         help="run this locally (i.e. not in production)")
+    parser.add_argument("--record", help="generate a proxy recording")
     parser.add_argument("--tool", default="mitmproxy",
                         help="the playback tool to use (default: %(default)s)")
     parser.add_argument("--host", default="localhost",
@@ -57,6 +61,7 @@ def main():
         playback = get_playback({
             "run_local": args.local,
             "playback_tool": args.tool,
+            "playback_record": args.record,
             "host": args.host,
             "binary": args.binary,
             "obj_path": args.objdir,
@@ -68,12 +73,13 @@ def main():
 
         LOG.info("Proxy running on port %d" % playback.port)
         # Wait for a keyboard interrupt from the caller so we know when to
-        # terminate. We wait using this method to allow Windows to respond to
-        # the Ctrl+Break signal so that we can exit cleanly.
-        while True:
-            time.sleep(1)
+        # terminate.
+        playback.wait()
+        return EXIT_EARLY_TERMINATE
     except KeyboardInterrupt:
         LOG.info("Terminating mozproxy")
         playback.stop()
+        return EXIT_SUCCESS
     except Exception as e:
         LOG.error(str(e), exc_info=True)
+        return EXIT_EXCEPTION

testing/mozbase/mozproxy/tests/test_command_line.py

@@ -1,17 +1,18 @@
 #!/usr/bin/env python
-from __future__ import absolute_import
+from __future__ import absolute_import, print_function
 import json
 import os
 import re
 import signal
 import threading
+import time
 
 import mozunit
 import pytest
 from mozbuild.base import MozbuildObject
 from mozprocess import ProcessHandler
 
-here = os.path.abspath(os.path.dirname(__file__))
+here = os.path.dirname(__file__)
 
 
 # This is copied from <python/mozperftest/mozperftest/utils.py>. It's copied
@@ -48,6 +49,8 @@ class OutputHandler(object):
         if not line.strip():
             return
         line = line.decode("utf-8", errors="replace")
+        # Print the output we received so we have useful logs if a test fails.
+        print(line)
 
         try:
             data = json.loads(line)
@@ -93,7 +96,8 @@ def test_run(install_mozproxy):
          "--local",
          "--binary=firefox",
          "--topsrcdir=" + build.topsrcdir,
-         "--objdir=" + build.topobjdir],
+         "--objdir=" + build.topobjdir,
+         os.path.join(here, "example.dump")],
         processOutputLine=output_handler,
         onFinish=output_handler.finished,
     )
@@ -101,6 +105,9 @@ def test_run(install_mozproxy):
     # The first time we run mozproxy, we need to fetch mitmproxy, which can
     # take a while...
     assert output_handler.port_event.wait(120) is True
+    # Give mitmproxy a bit of time to start up so we can verify that it's
+    # actually running before we kill mozproxy.
+    time.sleep(5)
     _kill_mozproxy(p.pid)
 
     assert p.wait(10) == 0