From 69cb0dd4d712b3227ffdedaac4aa9bca945a238a Mon Sep 17 00:00:00 2001 From: Andrew Gunnerson Date: Tue, 28 Feb 2023 01:50:50 -0500 Subject: [PATCH 1/4] tests: Add native Python parallel download implementation Signed-off-by: Andrew Gunnerson --- tests/downloader.py | 380 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 380 insertions(+) create mode 100644 tests/downloader.py diff --git a/tests/downloader.py b/tests/downloader.py new file mode 100644 index 0000000..eaf31e8 --- /dev/null +++ b/tests/downloader.py @@ -0,0 +1,380 @@ +import argparse +import collections +import copy +import dataclasses +import functools +import json +import os +import queue +import sys +import threading +import time +import typing +import urllib.request + +# This is a Python adaptation of the parallel downloader written for samfusdl. + +MIN_CHUNK_SIZE = 1 * 1024 * 1024 # 1 MiB + +DEFAULT_BUF_SIZE = 16384 +DEFAULT_RETRIES = 3 +DEFAULT_THREADS = 4 +DEFAULT_TIMEOUT = 30 + + +@dataclasses.dataclass +@functools.total_ordering +class Range: + start: int + end: int + + def __lt__(self, other) -> bool: + return (self.start, self.end) < (other.start, other.end) + + def __eq__(self, other) -> bool: + return (self.start, self.end) == (other.start, other.end) + + def __bool__(self) -> bool: + return self.start < self.end + + def size(self) -> int: + return self.end - self.start + + +class DownloadWorker(threading.Thread): + ''' + Thread to download a contiguous range from . + + After each buffer read, the buffer is submitted to . The + download continues only after reading a new range end offset from + . The controller uses this to interrupt this thread sooner + than expected in order to split the work for better resource utilization. + ''' + + def __init__(self, url: str, range: Range, output_queue: queue.Queue, + buf_size: typing.Optional[int] = None, + timeout: typing.Optional[int] = None): + super().__init__() + + self.url = url + self.range = copy.copy(range) + self.input_queue = queue.Queue() + self.output_queue = output_queue + self.timeout = DEFAULT_TIMEOUT if timeout is None else timeout + self.buf_size = DEFAULT_BUF_SIZE if buf_size is None else buf_size + + def run(self): + try: + self._download() + self.output_queue.put((self.ident, None)) + except BaseException as e: + self.output_queue.put((self.ident, e)) + + def _download(self): + req = urllib.request.Request(self.url) + req.add_header('Range', f'bytes={self.range.start}-{self.range.end}') + + with urllib.request.urlopen(req, timeout=self.timeout) as resp: + while self.range: + n = min(self.buf_size, self.range.size()) + data = resp.read(n) + + if not data or len(data) != n: + raise EOFError(f'Expected {n} bytes, but downloaded ' + f'{len(data)} bytes in {self.range}') + + self.output_queue.put((self.ident, data)) + + new_end = self.input_queue.get() + self.range.start += n + self.range.end = new_end + + +class DisplayCallback: + def progress(self, current, total): + raise NotImplementedError() + + def error(self, msg): + raise NotImplementedError() + + def finish(self): + raise NotImplementedError() + + +class DefaultDisplayCallback(DisplayCallback): + def __init__(self, delay_ms=50): + self.current = 0 + self.total = 0 + self.delay_ms = delay_ms + self.last_render = None + + def progress(self, current, total): + self.current = current + self.total = total + + now = time.perf_counter_ns() + + if self.last_render is None or \ + (now - self.last_render) / 1_000_000 > self.delay_ms: + self._clear_line() + self._print(f'{current / 1024 / 1024:.1f} / ' + f'{total / 1024 / 1024:.1f} MiB') + self.last_render = now + + def error(self, msg): + self._clear_line() + self._print(msg, end='\n') + + def finish(self): + self._clear_line() + + def _print(self, *args, **kwargs): + kwargs.setdefault('end', '') + kwargs.setdefault('file', sys.stderr) + kwargs.setdefault('flush', True) + print(*args, **kwargs) + + def _clear_line(self): + self._print('\033[2K\r') + + +def _get_content_length(url: str, timeout: typing.Optional[int] = None) -> int: + timeout = DEFAULT_TIMEOUT if timeout is None else timeout + req = urllib.request.Request(url, method='HEAD') + + with urllib.request.urlopen(req, timeout=timeout) as resp: + return int(resp.headers['content-length']) + + +def _download_ranges(f: typing.BinaryIO, url: str, ranges: list[Range], + display: DisplayCallback, + buf_size: typing.Optional[int] = None, + retries: typing.Optional[int] = None, + threads: typing.Optional[int] = None, + timeout: typing.Optional[int] = None): + retries = DEFAULT_RETRIES if retries is None else retries + threads = DEFAULT_THREADS if threads is None else threads + + workers = {} + # This is our own world view of the ranges that each worker should operate + # on. The worker's own view of its range will lag behind this until the + # next buffer loop cycle. It's fine for some workers' ranges to temporarily + # overlap since all file writes happen on the main thread. The worst case + # is that two workers download an overlapped range and the same file region + # is written twice. + worker_ranges = {} + output_queue = queue.Queue() + error_count = 0 + + file_size = _get_content_length(url, timeout=timeout) + f.truncate(file_size) + + if not ranges: + ranges.append(Range(0, file_size)) + + progress = file_size - sum(r.size() for r in ranges) + + remaining = collections.deque(ranges) + failed = [] + + try: + while True: + # Spawn new workers + while len(workers) < threads: + if not remaining and workers: + # No more ranges to download. Split another worker's range. + ident = max(worker_ranges, + key=lambda i: worker_ranges[i].size()) + old_range = worker_ranges[ident] + size = old_range.size() + + if size >= MIN_CHUNK_SIZE: + new_range = Range(old_range.start + size // 2, + old_range.end) + old_range.end = new_range.start + + remaining.appendleft(new_range) + + if not remaining: + break + + worker_range = remaining.popleft() + worker = DownloadWorker(url, worker_range, output_queue, + buf_size=buf_size, timeout=timeout) + worker.start() + + workers[worker.ident] = worker + worker_ranges[worker.ident] = worker_range + + if not workers: + # Everything is done + break + + ident, result = output_queue.get() + + # Worker completed successfully + if result is None: + workers[ident].join() + del workers[ident] + del worker_ranges[ident] + + # Worker failed + elif isinstance(result, BaseException): + display.error(f'[Worker#{ident}] {result}') + error_count += 1 + + # Retry range if we haven't failed too many times + if error_count < retries: + remaining.appendleft(worker_ranges[ident]) + else: + failed.append(worker_ranges[ident]) + + workers[ident].join() + del workers[ident] + del worker_ranges[ident] + + # Got a buffer, write it to the file + else: + f.seek(worker_ranges[ident].start) + f.write(result) + + progress += len(result) + worker_ranges[ident].start += len(result) + workers[ident].input_queue.put(worker_ranges[ident].end) + + display.progress(progress, file_size) + + display.finish() + + except BaseException: + # Cancel remaining threads + for ident, worker in workers.items(): + worker.input_queue.put(worker_ranges[ident].start) + + for worker in workers.values(): + worker.join() + + display.finish() + + raise + + finally: + # Mutate the original ranges parameter so that the caller knows which + # ranges are still remaining. This includes the not-yet-started, + # failed, and previously-in-progress ranges. + remaining.extend(failed) + remaining.extend(worker_ranges.values()) + + ranges.clear() + ranges.extend(sorted(remaining)) + + if remaining: + raise Exception(f'Download failed with {len(remaining)} chunks left ' + f'({file_size - progress} bytes)') + + +def _open_or_create(path: os.PathLike[str]) -> typing.BinaryIO(): + # Python's open() function has no way to open or create a file for both + # reading and writing without truncating without TOCTOU issues. + return os.fdopen(os.open(path, os.O_RDWR | os.O_CREAT, 0o644), 'r+b') + + +def download_ranges(out: os.PathLike[str], url: str, + initial_ranges: typing.Optional[list[Range]], + display: DisplayCallback, + buf_size: typing.Optional[int] = None, + retries: typing.Optional[int] = None, + threads: typing.Optional[int] = None, + timeout: typing.Optional[int] = None): + ''' + Download to with parallel threads. + + If is specified, only those sections of the file will be + downloaded. The empty regions are left untouched (i.e. filled with zeroes). + A `.state` file is written if the download is interrupted. If the state + file exists when this function is called, is ignored and + the ranges from the state file are used to resume the download. + ''' + + state_file = f'{out}.state' + + try: + with open(state_file, 'r') as f: + ranges_json = json.load(f) + ranges = [Range(r['start'], r['end']) for r in ranges_json] + except FileNotFoundError: + ranges = list(initial_ranges) if initial_ranges else [] + + try: + with _open_or_create(out) as f: + _download_ranges( + f, + url, + ranges, + display, + buf_size=buf_size, + retries=retries, + threads=threads, + timeout=timeout, + ) + finally: + if ranges: + with open(state_file, 'w') as f: + ranges_json = [{'start': r.start, 'end': r.end} + for r in ranges] + json.dump(ranges_json, f, indent=4) + else: + try: + os.unlink(state_file) + except FileNotFoundError: + pass + + +def parse_range(arg: str): + start, delim, end = arg.partition('-') + if not delim: + raise ValueError('Range should be in the form: -') + + start = int(start) + end = int(end) + + return Range(start, end) + + +def parse_args(args=None): + parser = argparse.ArgumentParser() + parser.add_argument('-u', '--url', required=True, + help='URL to download') + parser.add_argument('-o', '--output', required=True, + help='Output path') + parser.add_argument('-r', '--range', action='append', type=parse_range, + help='Half-open range to download ("-")') + parser.add_argument('-t', '--threads', default=4, type=int, + help='Number of parallel threads for downloading') + parser.add_argument('--buf-size', default=DEFAULT_BUF_SIZE, type=int, + help='Buffer size per download thread') + parser.add_argument('--retries', default=DEFAULT_RETRIES, type=int, + help='Maximum retries during download') + parser.add_argument('--timeout', default=DEFAULT_THREADS, type=int, + help='Connection timeout in seconds') + + return parser.parse_args() + + +def main(): + args = parse_args() + + download_ranges( + args.output, + args.url, + args.range, + DefaultDisplayCallback(), + buf_size=args.buf_size, + retries=args.retries, + threads=args.threads, + timeout=args.timeout, + ) + + +if __name__ == '__main__': + main() From db562315e900b885ed29e3639cdcaa4730ae16a8 Mon Sep 17 00:00:00 2001 From: Andrew Gunnerson Date: Tue, 28 Feb 2023 18:13:04 -0500 Subject: [PATCH 2/4] tests: Drop aria2 dependency Signed-off-by: Andrew Gunnerson --- tests/README.md | 4 --- tests/distros/Containerfile.alpine | 2 +- tests/distros/Containerfile.arch | 2 +- tests/distros/Containerfile.arch-wine | 1 - tests/distros/Containerfile.debian | 2 +- tests/distros/Containerfile.fedora | 2 +- tests/distros/Containerfile.opensuse | 2 +- tests/distros/Containerfile.ubuntu | 2 +- tests/distros/Containerfile.ubuntu-lts | 2 +- tests/tests.py | 35 ++++++++++++-------------- 10 files changed, 23 insertions(+), 31 deletions(-) diff --git a/tests/README.md b/tests/README.md index 37e673b..eb0d3db 100644 --- a/tests/README.md +++ b/tests/README.md @@ -14,8 +14,6 @@ For each image listed in the config, the test process will: ## Running the tests -aria2 must be installed before running the tests. It is used for downloading the OTA images. - To test against all of the device OTA images listed in [`tests.conf`](./tests.conf), run: ```bash @@ -24,8 +22,6 @@ python tests/tests.py To only run the tests for a specific device, pass in `-d `. This argument can be specified multiple times (eg. `-d GooglePixel7Pro -d GooglePixel6a`). -To pass in extra arguments to `aria2c`, use `-a=`. This can be used to enable concurrent downloads (eg. `-a=-x4 -a=-s4`). - ## Running the tests in a container The tests can also be run inside a podman container for easy testing on various Linux distros. To do so, run: diff --git a/tests/distros/Containerfile.alpine b/tests/distros/Containerfile.alpine index 4732f12..2733ab5 100644 --- a/tests/distros/Containerfile.alpine +++ b/tests/distros/Containerfile.alpine @@ -1,3 +1,3 @@ FROM docker.io/library/alpine:3.17 -RUN apk add --no-cache aria2 openssl py3-lz4 py3-protobuf +RUN apk add --no-cache openssl py3-lz4 py3-protobuf diff --git a/tests/distros/Containerfile.arch b/tests/distros/Containerfile.arch index 6b6366f..d812932 100644 --- a/tests/distros/Containerfile.arch +++ b/tests/distros/Containerfile.arch @@ -1,4 +1,4 @@ FROM docker.io/archlinux/archlinux:latest -RUN pacman --noconfirm -Syu --needed aria2 openssl python-lz4 python-protobuf \ +RUN pacman --noconfirm -Syu --needed openssl python-lz4 python-protobuf \ && yes | pacman -Scc diff --git a/tests/distros/Containerfile.arch-wine b/tests/distros/Containerfile.arch-wine index a6d512d..8b2f748 100644 --- a/tests/distros/Containerfile.arch-wine +++ b/tests/distros/Containerfile.arch-wine @@ -20,7 +20,6 @@ RUN cat /tmp/pacman.additional.conf >> /etc/pacman.conf \ && rm /tmp/pacman.additional.conf RUN pacman --noconfirm -Sy \ - mingw-w64-x86_64-aria2 \ mingw-w64-x86_64-ca-certificates \ mingw-w64-x86_64-openssl \ mingw-w64-x86_64-python \ diff --git a/tests/distros/Containerfile.debian b/tests/distros/Containerfile.debian index 56cb107..b048dda 100644 --- a/tests/distros/Containerfile.debian +++ b/tests/distros/Containerfile.debian @@ -1,5 +1,5 @@ FROM docker.io/library/debian:11 RUN apt-get -y update \ - && apt-get -y install aria2 openssl python3-lz4 python3-protobuf \ + && apt-get -y install openssl python3-lz4 python3-protobuf \ && find /var/cache/apt/archives /var/lib/apt/lists -mindepth 1 -delete diff --git a/tests/distros/Containerfile.fedora b/tests/distros/Containerfile.fedora index e357a78..7db6ee3 100644 --- a/tests/distros/Containerfile.fedora +++ b/tests/distros/Containerfile.fedora @@ -1,4 +1,4 @@ FROM registry.fedoraproject.org/fedora-toolbox:37 -RUN dnf install -y aria2 openssl python3-lz4 python3-protobuf \ +RUN dnf install -y openssl python3-lz4 python3-protobuf \ && find /var/cache/dnf -mindepth 1 -delete diff --git a/tests/distros/Containerfile.opensuse b/tests/distros/Containerfile.opensuse index 098f957..84d2c89 100644 --- a/tests/distros/Containerfile.opensuse +++ b/tests/distros/Containerfile.opensuse @@ -1,4 +1,4 @@ FROM registry.opensuse.org/opensuse/toolbox:latest -RUN zypper install -y aria2 openssl python3-lz4 python3-protobuf \ +RUN zypper install -y openssl python3-lz4 python3-protobuf \ && zypper clean -a diff --git a/tests/distros/Containerfile.ubuntu b/tests/distros/Containerfile.ubuntu index 7d966a8..7d423e7 100644 --- a/tests/distros/Containerfile.ubuntu +++ b/tests/distros/Containerfile.ubuntu @@ -1,5 +1,5 @@ FROM docker.io/library/ubuntu:22.10 RUN apt-get -y update \ - && apt-get -y install aria2 openssl python3-lz4 python3-protobuf \ + && apt-get -y install openssl python3-lz4 python3-protobuf \ && find /var/cache/apt/archives /var/lib/apt/lists -mindepth 1 -delete diff --git a/tests/distros/Containerfile.ubuntu-lts b/tests/distros/Containerfile.ubuntu-lts index b728af6..10e2b46 100644 --- a/tests/distros/Containerfile.ubuntu-lts +++ b/tests/distros/Containerfile.ubuntu-lts @@ -1,5 +1,5 @@ FROM docker.io/library/ubuntu:22.04 RUN apt-get -y update \ - && apt-get -y install aria2 openssl python3-lz4 python3-protobuf \ + && apt-get -y install openssl python3-lz4 python3-protobuf \ && find /var/cache/apt/archives /var/lib/apt/lists -mindepth 1 -delete diff --git a/tests/tests.py b/tests/tests.py index c28b7e3..df10225 100755 --- a/tests/tests.py +++ b/tests/tests.py @@ -4,7 +4,6 @@ import argparse import configparser import hashlib import os -import subprocess import sys import tempfile import zipfile @@ -14,6 +13,8 @@ from avbroot import main from avbroot import ota from avbroot import util +import downloader + def url_filename(url): return url.split('/')[-1] @@ -40,26 +41,26 @@ def verify_checksum(path, sha256_hex): f'but have {hasher.hexdigest()}') -def download(directory, url, sha256_hex, revalidate=False, extra_args=[]): +def download(directory, url, sha256_hex, revalidate=False): path = os.path.join(directory, url_filename(url)) + validate = True - if os.path.exists(path) and not os.path.exists(path + '.aria2'): - if revalidate: - print('Verifying checksum of', path) - - verify_checksum(path, sha256_hex) + if os.path.exists(path) and not os.path.exists(path + '.state'): + validate = revalidate else: print('Downloading', url, 'to', path) - subprocess.check_call([ - 'aria2c', - *extra_args, - '--auto-file-renaming=false', - f'--checksum=SHA-256={sha256_hex}', - f'--dir={os.path.dirname(path)}', - f'--out={os.path.basename(path)}', + downloader.download_ranges( + path, url, - ]) + None, + downloader.DefaultDisplayCallback(), + ) + + if validate: + print('Verifying checksum of', path) + + verify_checksum(path, sha256_hex) return path @@ -80,7 +81,6 @@ def run_tests(args, config, files_dir): config['magisk']['url'], config['magisk']['sha256'], revalidate=args.revalidate, - extra_args=args.aria2c_arg, ) for device, image in device_configs(args, config): @@ -92,7 +92,6 @@ def run_tests(args, config, files_dir): image['url'], image['sha256'], revalidate=args.revalidate, - extra_args=args.aria2c_arg, ) patched_file = image_file + args.output_file_suffix @@ -159,8 +158,6 @@ def run_tests(args, config, files_dir): def parse_args(): parser = argparse.ArgumentParser() - parser.add_argument('-a', '--aria2c-arg', default=[], action='append', - help='Argument to pass to aria2c') parser.add_argument('-d', '--device', action='append', help='Device image to test against') parser.add_argument('--delete-on-success', action='store_true', From 76eec8fb6e579c32a900703f2e61138637bfbfb7 Mon Sep 17 00:00:00 2001 From: Andrew Gunnerson Date: Tue, 28 Feb 2023 23:46:56 -0500 Subject: [PATCH 3/4] tests/downloader.py: Avoid unnecessary buffer allocations in DownloadWorker Signed-off-by: Andrew Gunnerson --- tests/downloader.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/downloader.py b/tests/downloader.py index eaf31e8..313f77c 100644 --- a/tests/downloader.py +++ b/tests/downloader.py @@ -75,15 +75,18 @@ class DownloadWorker(threading.Thread): req.add_header('Range', f'bytes={self.range.start}-{self.range.end}') with urllib.request.urlopen(req, timeout=self.timeout) as resp: + buf = bytearray(self.buf_size) + buf_view = memoryview(buf) + while self.range: - n = min(self.buf_size, self.range.size()) - data = resp.read(n) + to_read = min(self.buf_size, self.range.size()) + n = resp.readinto(buf_view[:to_read]) - if not data or len(data) != n: + if n != to_read: raise EOFError(f'Expected {n} bytes, but downloaded ' - f'{len(data)} bytes in {self.range}') + f'{to_read} bytes in {self.range}') - self.output_queue.put((self.ident, data)) + self.output_queue.put((self.ident, buf_view[:n])) new_end = self.input_queue.get() self.range.start += n From 003fbf18ba18cea2965208e37c27ff790119ed94 Mon Sep 17 00:00:00 2001 From: Andrew Gunnerson Date: Wed, 1 Mar 2023 00:18:23 -0500 Subject: [PATCH 4/4] tests/downloader.py: Add MiB/s to DefaultDisplayCallback output No fancy algorithm. Just a simple moving average over 5 seconds, updated at 100ms intervals. Signed-off-by: Andrew Gunnerson --- tests/downloader.py | 35 ++++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/downloader.py b/tests/downloader.py index 313f77c..3e6ed86 100644 --- a/tests/downloader.py +++ b/tests/downloader.py @@ -94,10 +94,10 @@ class DownloadWorker(threading.Thread): class DisplayCallback: - def progress(self, current, total): + def progress(self, current: int, total: int): raise NotImplementedError() - def error(self, msg): + def error(self, msg: str): raise NotImplementedError() def finish(self): @@ -105,11 +105,16 @@ class DisplayCallback: class DefaultDisplayCallback(DisplayCallback): - def __init__(self, delay_ms=50): + # Speed is a simple moving average over 5 seconds + AVG_INTERVAL_MS = 100 + AVG_WINDOW_SIZE = 5000 / AVG_INTERVAL_MS + + def __init__(self, interval_ms: int = 50): self.current = 0 self.total = 0 - self.delay_ms = delay_ms + self.interval_ms = interval_ms self.last_render = None + self.avg = collections.deque() def progress(self, current, total): self.current = current @@ -117,11 +122,27 @@ class DefaultDisplayCallback(DisplayCallback): now = time.perf_counter_ns() + if not self.avg or \ + (now - self.avg[-1][0]) / 1_000_000 > self.AVG_INTERVAL_MS: + if len(self.avg) == self.AVG_WINDOW_SIZE: + self.avg.popleft() + + self.avg.append((now, current)) + if self.last_render is None or \ - (now - self.last_render) / 1_000_000 > self.delay_ms: + (now - self.last_render) / 10 ** 6 > self.interval_ms: + current_mib = current / 1024 ** 2 + total_mib = total / 1024 ** 2 + if len(self.avg) == 1: + speed_mib_s = 0 + else: + avg_window_mib = (self.avg[-1][1] - self.avg[0][1]) / 1024 ** 2 + avg_window_secs = (self.avg[-1][0] - self.avg[0][0]) / 10 ** 9 + speed_mib_s = avg_window_mib / avg_window_secs + self._clear_line() - self._print(f'{current / 1024 / 1024:.1f} / ' - f'{total / 1024 / 1024:.1f} MiB') + self._print(f'{current_mib:.1f} / {total_mib:.1f} MiB ' + f'({speed_mib_s:.1f} MiB/s)') self.last_render = now def error(self, msg):