Merge pull request #72 from chenxiaolong/parallel_download

tests: Add native Python parallel download implementation

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 <device>`. This argument can be specified multiple times (eg. `-d GooglePixel7Pro -d GooglePixel6a`).
 
-To pass in extra arguments to `aria2c`, use `-a=<arg>`. 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:

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

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

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 \

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

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

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

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

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

tests/downloader.py

@@ -0,0 +1,404 @@
+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 <url>.
+
+    After each buffer read, the buffer is submitted to <output_queue>. The
+    download continues only after reading a new range end offset from
+    <input_queue>. 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:
+            buf = bytearray(self.buf_size)
+            buf_view = memoryview(buf)
+
+            while self.range:
+                to_read = min(self.buf_size, self.range.size())
+                n = resp.readinto(buf_view[:to_read])
+
+                if n != to_read:
+                    raise EOFError(f'Expected {n} bytes, but downloaded '
+                                   f'{to_read} bytes in {self.range}')
+
+                self.output_queue.put((self.ident, buf_view[:n]))
+
+                new_end = self.input_queue.get()
+                self.range.start += n
+                self.range.end = new_end
+
+
+class DisplayCallback:
+    def progress(self, current: int, total: int):
+        raise NotImplementedError()
+
+    def error(self, msg: str):
+        raise NotImplementedError()
+
+    def finish(self):
+        raise NotImplementedError()
+
+
+class DefaultDisplayCallback(DisplayCallback):
+    # 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.interval_ms = interval_ms
+        self.last_render = None
+        self.avg = collections.deque()
+
+    def progress(self, current, total):
+        self.current = current
+        self.total = total
+
+        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) / 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_mib:.1f} / {total_mib:.1f} MiB '
+                        f'({speed_mib_s:.1f} MiB/s)')
+            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 <url> to <out> with <threads> parallel threads.
+
+    If <initial_ranges> 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, <initial_ranges> 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>-<end>')
+
+    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 ("<start>-<end>")')
+    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()

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',