1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5from __future__ import absolute_import
6
7import os
8import shutil
9import socket
10
11from logger.logger import RaptorLogger
12from wptserve import server, handlers
13
14LOG = RaptorLogger(component="raptor-benchmark")
15here = os.path.abspath(os.path.dirname(__file__))
16
17
18class Benchmark(object):
19    """utility class for running benchmarks in raptor"""
20
21    def __init__(self, config, test):
22        self.config = config
23        self.test = test
24
25        # bench_dir is where we will download all mitmproxy required files
26        # when running locally it comes from obj_path via mozharness/mach
27        if self.config.get("obj_path", None) is not None:
28            self.bench_dir = self.config.get("obj_path")
29        else:
30            # in production it is ../tasks/task_N/build/tests/raptor/raptor/...
31            # 'here' is that path, we can start with that
32            self.bench_dir = here
33
34        # now add path for benchmark source; locally we put it in a raptor benchmarks
35        # folder; in production the files are automatically copied to a different dir
36        if self.config.get("run_local", False):
37            self.bench_dir = os.path.join(
38                self.bench_dir, "testing", "raptor", "benchmarks"
39            )
40        else:
41            self.bench_dir = os.path.join(
42                self.bench_dir, "tests", "webkit", "PerformanceTests"
43            )
44
45            # Some benchmarks may have been downloaded from a fetch task, make
46            # sure they get copied over.
47            fetches_dir = os.environ.get("MOZ_FETCHES_DIR")
48            if (
49                test.get("fetch_task", False)
50                and fetches_dir
51                and os.path.isdir(fetches_dir)
52            ):
53                for name in os.listdir(fetches_dir):
54                    if test.get("fetch_task").lower() in name.lower():
55                        path = os.path.join(fetches_dir, name)
56                        if os.path.isdir(path):
57                            shutil.copytree(path, os.path.join(self.bench_dir, name))
58
59        LOG.info("bench_dir contains:")
60        LOG.info(os.listdir(self.bench_dir))
61
62        # now have the benchmark source ready, go ahead and serve it up!
63        self.start_http_server()
64
65    def start_http_server(self):
66        self.write_server_headers()
67
68        # pick a free port
69        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
70        sock.bind(("", 0))
71        self.host = self.config["host"]
72        self.port = sock.getsockname()[1]
73        sock.close()
74        _webserver = "%s:%d" % (self.host, self.port)
75
76        self.httpd = self.setup_webserver(_webserver)
77        self.httpd.start()
78
79    def write_server_headers(self):
80        # to add specific headers for serving files via wptserve, write out a headers dir file
81        # see http://wptserve.readthedocs.io/en/latest/handlers.html#file-handlers
82        LOG.info("writing wptserve headers file")
83        headers_file = os.path.join(self.bench_dir, "__dir__.headers")
84        file = open(headers_file, "w")
85        file.write("Access-Control-Allow-Origin: *")
86        file.close()
87        LOG.info("wrote wpt headers file: %s" % headers_file)
88
89    def setup_webserver(self, webserver):
90        LOG.info("starting webserver on %r" % webserver)
91        LOG.info("serving benchmarks from here: %s" % self.bench_dir)
92        self.host, self.port = webserver.split(":")
93
94        return server.WebTestHttpd(
95            host=self.host,
96            port=int(self.port),
97            doc_root=self.bench_dir,
98            routes=[("GET", "*", handlers.file_handler)],
99        )
100
101    def stop_serve(self):
102        LOG.info("TODO: stop serving benchmark source")
103        pass
104