tor-browser

hooks_throttling.py (6109B)

---

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
"""
Drives the throttling feature when the test calls our
controlled server.
"""

import http.client
import json
import os
import sys
import time
from urllib.parse import urlparse

from mozperftest.test.browsertime import add_option
from mozperftest.utils import get_tc_secret

ENDPOINTS = {
   "linux": "h3.dev.mozaws.net",
   "darwin": "h3.mac.dev.mozaws.net",
   "win32": "h3.win.dev.mozaws.net",
}
CTRL_SERVER = ENDPOINTS[sys.platform]
TASK_CLUSTER = "TASK_ID" in os.environ.keys()
_SECRET = {
   "throttler_host": f"https://{CTRL_SERVER}/_throttler",
   "throttler_key": os.environ.get("WEBNETEM_KEY", ""),
}
if TASK_CLUSTER:
   _SECRET.update(get_tc_secret())

if _SECRET["throttler_key"] == "":
   if TASK_CLUSTER:
       raise Exception("throttler_key not found in secret")
   raise Exception("WEBNETEM_KEY not set")

_TIMEOUT = 30
WAIT_TIME = 60 * 10
IDLE_TIME = 10
BREATHE_TIME = 20


class Throttler:
   def __init__(self, env, host, key):
       self.env = env
       self.host = host
       self.key = key
       self.verbose = env.get_arg("verbose", False)
       self.logger = self.verbose and self.env.info or self.env.debug

   def log(self, msg):
       self.logger("[throttler] " + msg)

   def _request(self, action, data=None):
       kw = {}
       headers = {b"X-WEBNETEM-KEY": self.key}
       verb = data is None and "GET" or "POST"
       if data is not None:
           data = json.dumps(data)
           headers[b"Content-type"] = b"application/json"

       parsed = urlparse(self.host)
       server = parsed.netloc
       path = parsed.path
       if action != "status":
           path += "/" + action

       self.log(f"Calling {verb} {path}")
       conn = http.client.HTTPSConnection(server, timeout=_TIMEOUT)
       conn.request(verb, path, body=data, headers=headers, **kw)
       resp = conn.getresponse()
       res = resp.read()
       if resp.status >= 400:
           raise Exception(res)
       res = json.loads(res)
       return res

   def start(self, data=None):
       self.log("Starting")
       now = time.time()
       acquired = False

       while time.time() - now < WAIT_TIME:
           status = self._request("status")
           if status.get("test_running"):
               # a test is running
               self.log("A test is already controlling the server")
               self.log(f"Waiting {IDLE_TIME} seconds")
           else:
               try:
                   self._request("start_test")
                   acquired = True
                   break
               except Exception:
                   # we got beat in the race
                   self.log("Someone else beat us")
           time.sleep(IDLE_TIME)

       if not acquired:
           raise Exception("Could not acquire the test server")

       if data is not None:
           self._request("shape", data)

   def stop(self):
       self.log("Stopping")
       try:
           self._request("reset")
       finally:
           self._request("stop_test")


def get_throttler(env):
   host = _SECRET["throttler_host"]
   key = _SECRET["throttler_key"].encode()
   return Throttler(env, host, key)


_PROTOCOL = "h2", "h3"
_PAGE = "gallery", "news", "shopping", "photoblog"

# set the network condition here.
# each item has a name and some netem options:
#
# loss_ratio: specify percentage of packets that will be lost
# loss_corr: specify a correlation factor for the random packet loss
# dup_ratio: specify percentage of packets that will be duplicated
# delay: specify an overall delay for each packet
# jitter: specify amount of jitter in milliseconds
# delay_jitter_corr: specify a correlation factor for the random jitter
# reorder_ratio: specify percentage of packets that will be reordered
# reorder_corr: specify a correlation factor for the random reordering
#
_THROTTLING = (
   {"name": "full"},  # no throttling.
   {"name": "one", "delay": "20"},
   {"name": "two", "delay": "50"},
   {"name": "three", "delay": "100"},
   {"name": "four", "delay": "200"},
   {"name": "five", "delay": "300"},
)


def get_test():
   """Iterate on test conditions.

   For each cycle, we return a combination of: protocol, page, throttling
   settings. Each combination has a name, and that name will be used along with
   the protocol as a prefix for each metrics.
   """
   for proto in _PROTOCOL:
       for page in _PAGE:
           url = f"https://{CTRL_SERVER}/{page}.html"
           for throttler_settings in _THROTTLING:
               yield proto, page, url, throttler_settings


combo = get_test()


def before_cycle(metadata, env, cycle, script):
   global combo
   if "throttlable" not in script["tags"]:
       return
   throttler = get_throttler(env)
   try:
       proto, page, url, throttler_settings = next(combo)
   except StopIteration:
       combo = get_test()
       proto, page, url, throttler_settings = next(combo)

   # setting the url for the browsertime script
   add_option(env, "browsertime.url", url, overwrite=True)

   # enabling http if needed
   if proto == "h3":
       add_option(env, "firefox.preference", "network.http.http3.enable:true")

   # prefix used to differenciate metrics
   name = throttler_settings["name"]
   script["name"] = f"{name}_{proto}_{page}"

   # throttling the controlled server if needed
   if throttler_settings != {"name": "full"}:
       env.info("Calling the controlled server")
       throttler.start(throttler_settings)
   else:
       env.info("No throttling for this call")
       throttler.start()


def after_cycle(metadata, env, cycle, script):
   if "throttlable" not in script["tags"]:
       return
   throttler = get_throttler(env)
   try:
       throttler.stop()
   except Exception:
       pass

   # give a chance for a competitive job to take over
   time.sleep(BREATHE_TIME)