From ef8ed3a11c62fb503c40a16a0cdb571de7a48c23 Mon Sep 17 00:00:00 2001 From: Mole Shang <135e2@135e2.dev> Date: Fri, 21 Nov 2025 22:01:10 +0800 Subject: [PATCH 1/3] feat: implement on-disk download streaming --- src/proxpi/_cache.py | 372 +++++++++++++++++++++++++++++++++++-------- src/proxpi/server.py | 29 ++++ 2 files changed, 335 insertions(+), 66 deletions(-) diff --git a/src/proxpi/_cache.py b/src/proxpi/_cache.py index 421f7c5..81ccfd7 100644 --- a/src/proxpi/_cache.py +++ b/src/proxpi/_cache.py @@ -12,10 +12,14 @@ import functools import posixpath import threading +import queue import dataclasses import urllib.parse + import requests +from requests.exceptions import RequestException +from urllib3.util import SKIP_HEADER import lxml.etree INDEX_URL = os.environ.get("PROXPI_INDEX_URL", "https://pypi.org/simple/") @@ -40,19 +44,20 @@ CACHE_SIZE = int(os.environ.get("PROXPI_CACHE_SIZE", 5368709120)) CACHE_DIR = os.environ.get("PROXPI_CACHE_DIR") -DOWNLOAD_TIMEOUT = float(os.environ.get("PROXPI_DOWNLOAD_TIMEOUT", 0.9)) CONNECT_TIMEOUT = ( float(os.environ["PROXPI_CONNECT_TIMEOUT"]) if os.environ.get("PROXPI_CONNECT_TIMEOUT") - else None + else 3.1 ) READ_TIMEOUT = ( float(os.environ["PROXPI_READ_TIMEOUT"]) if os.environ.get("PROXPI_READ_TIMEOUT") - else None + else 20.0 ) +CHUNK_SIZE: t.Final[int] = 16 * 1024 + logger = logging.getLogger(__name__) _name_normalise_re = re.compile("[-_.]+") _hostname_normalise_pattern = re.compile(r"[^a-z0-9]+") @@ -289,6 +294,210 @@ class NotFound(ValueError): pass +@dataclasses.dataclass +class Subscriber: + """Individual subscriber to a download request.""" + + id: int + download_path: str + read_position: int = 0 + # we shall use other methods to let pub to replace first, then sub read + notify_queue: queue.Queue[int | RequestException] = queue.Queue() + + @staticmethod + def _read_to(f, start, end) -> int: + """Yield chunks from start to end. + + If end=-1, use file size as the end + + Returns: + position after read + """ + + if end == -1: + end = f.seek(0, os.SEEK_END) + + while True: + bytes_to_read = min( + end - start, + CHUNK_SIZE, + ) + if bytes_to_read <= 0: # No new data available atm + break + f.seek(start) + chunk = f.read(bytes_to_read) + assert chunk, "should have data chunk if bytes_to_read!=0" + yield chunk + start += len(chunk) + return start + + def generate(self): + """Stream generator for flask streaming.""" + + with open(self.download_path, "rb", 0) as f: + while True: + # Wait on notification of new data + try: + i = self.notify_queue.get(timeout=READ_TIMEOUT) + if isinstance(i, RequestException): # raise and exit control flow + raise i + + write_position = i + # Read all data available + self.read_position = yield from self._read_to( + f, self.read_position, write_position + ) + if write_position == -1: # receives EOF + break + except queue.Empty: + raise TimeoutError(f"subscriber timeout on {self.download_path}") + + +@dataclasses.dataclass +class DownloadStatus: + """ + Download state tracker. + + Also act as publisher (and do file-write). + """ + + _temp_download_path: str = "" + _upstream_status_code: int = 0 + _upstream_headers: t.Tuple[str, str] = dataclasses.field(default_factory=tuple) + _subscribers: t.Dict[int, Subscriber] = dataclasses.field(default_factory=dict) + _subscriber_counter: int = 0 + _write_position: int = 0 + _write_handle: File = None + _notify_lock: threading.Lock = threading.Lock() + _pub_eof: bool = False + _error: RequestException = None + # cleanup callback + _cleanup: t.Callable = lambda: 0 + _thread: threading.Thread = None + + def _can_cleanup(self): + """Check if pub AND all subs finish.""" + return ( + self._pub_eof + and len(self._subscribers) == 0 + and os.path.exists(self._temp_download_path) + ) + + @staticmethod + def _notify(subscriber: Subscriber, data: int | RequestException): + """ + NOTE: ALWAYS acquire _notify_lock before calling + """ + try: + subscriber.notify_queue.put_nowait(data) + except queue.Full: + logger.debug(f"notify fail on {subscriber}") + pass + + def _notify_all(self, data: int | RequestException): + with self._notify_lock: + for subscriber in self._subscribers.values(): + self._notify(subscriber, data) + + @property + def status_code(self) -> int: + """Upstream status code""" + return self._upstream_status_code + + @property + def headers(self) -> t.List[t.Tuple[str, str]]: + """Reverse proxy stripped headers""" + # https://www.rfc-editor.org/rfc/rfc2616#section-13.5.1 + excluded_headers = [ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", + ] + return [ + (name, value) + for (name, value) in self._upstream_headers + if name.lower() not in excluded_headers + ] + + @property + def ready(self) -> bool: + """Whether sub is ready to receive""" + return ( + os.path.exists(self._temp_download_path) + and os.path.getsize(self._temp_download_path) > 0 + ) + + def is_alive(self) -> bool: + return ( + (self._thread and self._thread.is_alive()) # Fetch thread is still alive + or len(self._subscribers) + != 0 # Or fetch is done, but subs still downloading + ) + + def add_subscriber(self) -> Subscriber: + # Wait until content is ready + while not self.ready: + continue + + self._notify_lock.acquire() + subscriber_id = self._subscriber_counter + self._subscriber_counter += 1 + subscriber = Subscriber(subscriber_id, self._temp_download_path) + self._subscribers[subscriber_id] = subscriber + logger.debug( + f"{self._temp_download_path}: add, now {len(self._subscribers)} subs" + ) + + if self._pub_eof: # send again to wakeup + self._notify(subscriber, self._error) + self._notify_lock.release() + + return subscriber + + def remove_subscriber(self, subscriber: Subscriber): + with self._notify_lock: + self._subscribers.pop(subscriber.id) + logger.debug( + f"{self._temp_download_path}: remove, now {len(self._subscribers)} subs" + ) + if self._can_cleanup(): + self._cleanup() + + def set_cleanup_callback(self, cleanup_func: t.Callable): + self._cleanup = cleanup_func + + def broadcast_chunk(self, chunk: bytes): + """Write and broadcast current position to all subscribers known.""" + if self._write_handle is None: + self._write_handle = open(self._temp_download_path, "ab", 0) + self._write_handle.write(chunk) + self._write_position += len(chunk) + self._notify_all(self._write_position) + + def broadcast_eof(self, error: RequestException | int = -1): + """Broadcast EOF (with exception) to all subscribers known.""" + if self._write_handle: + self._write_handle.close() + self._pub_eof = True + self._error = error + self._notify_all(error) + + +class Downloading(Exception): + """Internal exception to pass DownloadStatus""" + + status: DownloadStatus = None + + def __init__(self, status): + super().__init__("Downloading exception") + self.status = status + + class Thread(threading.Thread): """Exception-storing thread runner.""" @@ -731,7 +940,7 @@ class _FileCache: max_size: int cache_dir: str _cache_dir_provided: t.Union[str, None] - _files: t.Dict[str, t.Union[_CachedFile, Thread]] + _files: t.Dict[str, t.Union[_CachedFile, DownloadStatus]] _evict_lock: threading.Lock _stats: _CacheStats _download_filename_suffix = ".proxpi-partial" @@ -740,7 +949,6 @@ def __init__( self, max_size: int, cache_dir: str = None, - download_timeout: float = 0.9, session: requests.Session = None, ): """Initialise file-cache. @@ -755,11 +963,11 @@ def __init__( self.max_size = max_size self.cache_dir = os.path.abspath(cache_dir or tempfile.mkdtemp()) - self.download_timeout = download_timeout self.session = session or requests.Session() self._cache_dir_provided = cache_dir self._files = {} self._evict_lock = threading.Lock() + self._status_init_lock = threading.Lock() self._stats = _CacheStats(name="Files") self._populate_files_from_existing_cache_dir() @@ -797,62 +1005,75 @@ def _get_key(url: str) -> str: parent = _hostname_normalise_pattern.sub("-", urlsplit.hostname) return posixpath.join(parent, *_split_path(urlsplit.path, posixpath.split)) - def _download_file(self, url: str, path: str): + def _download_file(self, url: str, path: str, status: DownloadStatus): """Download a file. + NOTE: + RequestExceptions will be delegated under + DownloadStatus->Subscriber->Generator + This function normally WON'T raise. + Args: url: URL of file to download path: local path to download to """ - url_masked = _mask_password(url) logger.debug(f"Downloading '{url_masked}' to '{path}'") - response = self.session.get(url, stream=True) - if response.status_code // 100 >= 4: - logger.error( - f"Failed to download '{url_masked}': " - f"status={response.status_code}, body={response.text}" - ) - return + + # Set temp download path in the status object + status._temp_download_path = path + self._download_filename_suffix parent, _ = os.path.split(path) os.makedirs(parent, exist_ok=True) - download_path = path + self._download_filename_suffix - with open(download_path, mode="wb") as f: - for chunk in response.iter_content(chunk_size=16 * 1024): - f.write(chunk) - os.replace(download_path, path) - key = self._get_key(url) - self._files[key] = _CachedFile(path, os.stat(path).st_size, 0) - logger.debug(f"Finished downloading '{url_masked}'") + # truncate to zero (or create if not exist) + open(status._temp_download_path, "wb").close() - def _wait_for_existing_download(self, url: str) -> bool: - """Wait for existing download, if any. + try: + response = self.session.get(url, stream=True) + status._upstream_status_code = response.status_code + status._upstream_headers = response.raw.headers.items() + if status.status_code // 100 >= 4: + raise RequestException( + f"Failed to fetch '{url_masked}': " + f"status={response.status_code}, body={response.text}" + ) + for chunk in response.iter_content(chunk_size=CHUNK_SIZE): + # Broadcast chunk to all subscribers + status.broadcast_chunk(chunk) + + except RequestException as e: + logger.error(f"{e.__class__.__name__} when fetching: {e}") + + def cleanup(): + os.unlink(status._temp_download_path) + # Remove entry + with self._status_init_lock: + key = self._get_key(url) + if key in self._files: + del self._files[key] + + status.set_cleanup_callback(cleanup) + status.broadcast_eof(e) + return - Returns: - whether the wait is given up (ie if time-out was reached or - exception was encountered) - """ + def cleanup(): + os.replace(status._temp_download_path, path) + with self._status_init_lock: + key = self._get_key(url) + # Mark the file as cached + self._files[key] = _CachedFile(path, os.stat(path).st_size, 0) + logger.debug(f"All subs finished downloading '{url_masked}'") - file = self._files.get(url) - if isinstance(file, Thread): - url_masked = _mask_password(url) - logger.debug(f"Waiting for existing download of: {url_masked}") - try: - file.join(self.download_timeout) - except Exception as e: - if file.exc and file == self._files[url]: - self._files.pop(url, None) - logger.error(f"Failed to download '{url_masked}'", exc_info=e) - return True - if isinstance(self._files[url], Thread): - return True # default to original URL (due to timeout or HTTP error) - return False + status.set_cleanup_callback(cleanup) + status.broadcast_eof() + logger.debug(f"Finished fetching '{url_masked}'") def _get_cached(self, url: str) -> t.Union[str, None]: """Get file from cache.""" if url in self._files: file = self._files[url] - assert isinstance(file, _CachedFile) + if file and not isinstance(file, _CachedFile): + # still downloading + return None file.n_hits += 1 self._stats.add_hit(key=url) return file.path @@ -864,9 +1085,13 @@ def _start_downloading(self, url: str): key = self._get_key(url) path = os.path.join(self.cache_dir, *_split_path(key, posixpath.split)) - thread = Thread(target=self._download_file, args=(url, path)) - self._files[key] = thread - thread.start() + status = DownloadStatus() + status._thread = Thread(target=self._download_file, args=(url, path, status)) + # Store DownloadStatus to track if download is still running + self._files[key] = status + status._thread.start() + self._status_init_lock.release() + raise Downloading(status) def _evict_lfu(self, url: str): """Evict least-frequently-used files until under max cache size.""" @@ -895,16 +1120,32 @@ def get(self, url: str) -> str: if self.max_size == 0: return url key = self._get_key(url) - path = url - given_up = self._wait_for_existing_download(key) - if not given_up: - path = self._get_cached(key) - if not path: - self._start_downloading(url) - with self._evict_lock: - self._evict_lfu(url) - path = self.get(url) - return path + + # First, check if file is already cached + path = self._get_cached(key) + if path: + return path + + self._status_init_lock.acquire() + # Now check if still downloading + file = self._files.get(key) + + # There's an ongoing download - get the status to allow subscription + if isinstance(file, DownloadStatus): + assert file.is_alive(), "should have an alive Downloading thread" + self._status_init_lock.release() + # Raise Downloading exception to signal server + # to use the ongoing download stream + raise Downloading(file) + return url # unreachable + # No downloading thread found, + # evict and start a new download + else: + with self._evict_lock: + self._evict_lfu(url) + # NOTE: lock would be released inside _start_downloading() + self._start_downloading(url) + return url # unreachable @dataclasses.dataclass @@ -932,17 +1173,14 @@ def from_config(cls): if proxpi_version: session.headers["User-Agent"] = f"proxpi/{proxpi_version}" - if CONNECT_TIMEOUT and READ_TIMEOUT: - session.default_timeout = (CONNECT_TIMEOUT, READ_TIMEOUT) - elif CONNECT_TIMEOUT: - session.default_timeout = (CONNECT_TIMEOUT, 20.0) - elif READ_TIMEOUT: - session.default_timeout = (3.1, READ_TIMEOUT) + session.default_timeout = (CONNECT_TIMEOUT, READ_TIMEOUT) + + # Accept no compression here, + # as our proxy needs a correct Content-Length + session.headers["Accept-Encoding"] = SKIP_HEADER root_cache = cls._index_cache_cls(INDEX_URL, INDEX_TTL, session) - file_cache = cls._file_cache_cls( - CACHE_SIZE, CACHE_DIR, DOWNLOAD_TIMEOUT, session - ) + file_cache = cls._file_cache_cls(CACHE_SIZE, CACHE_DIR, session) if len(EXTRA_INDEX_URLS) != len(EXTRA_INDEX_TTLS): raise RuntimeError( f"Number of extra index URLs doesn't equal number of extra index " @@ -1028,6 +1266,8 @@ def get_file(self, package_name: str, file_name: str) -> str: Raises: NotFound: if project doesn't exist in any index or file doesn't exist in project + Downloading: if there's already a thread fetching so that + we can reuse the DownloadStatus """ try: diff --git a/src/proxpi/server.py b/src/proxpi/server.py index 3fc0ab8..ee4f592 100644 --- a/src/proxpi/server.py +++ b/src/proxpi/server.py @@ -11,6 +11,8 @@ import jinja2 import werkzeug.exceptions +from requests.exceptions import RequestException + from . import _cache try: @@ -211,6 +213,33 @@ def get_file(package_name: str, file_name: str): except _cache.NotFound: flask.abort(404) raise + except _cache.Downloading as e: + status = e.status + subscriber = status.add_subscriber() + + def stream(): + try: + yield from subscriber.generate() + except GeneratorExit: + logger.info("User cancelled request") + pass + except ( + RequestException or TimeoutError + ) as e: # can't edit status code now, give a hint in stream + logger.error(f"{e.__class__.__name__} on stream: {e}") + yield f"***[proxpi] server fetch request failed: {e}***".encode() + return + finally: + status.remove_subscriber(subscriber) + + response = flask.Response( + stream(), + status.status_code, + status.headers, + direct_passthrough=True, + ) + return response + scheme = urllib.parse.urlparse(path).scheme if scheme and scheme != "file": return flask.redirect(path) From 6461d0b7c2c62b4e6a80619c9f62a97fe6f97aea Mon Sep 17 00:00:00 2001 From: Mole Shang <135e2@135e2.dev> Date: Mon, 24 Nov 2025 16:12:59 +0800 Subject: [PATCH 2/3] feat: support setting conn pool size Since we now use streaming, set larger pool size to avoid exhausting the requests connection pool. --- src/proxpi/_cache.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/proxpi/_cache.py b/src/proxpi/_cache.py index 81ccfd7..bffb59b 100644 --- a/src/proxpi/_cache.py +++ b/src/proxpi/_cache.py @@ -55,6 +55,7 @@ if os.environ.get("PROXPI_READ_TIMEOUT") else 20.0 ) +POOL_SIZE = int(os.environ.get("PROXPI_POOL_SIZE", 100)) CHUNK_SIZE: t.Final[int] = 16 * 1024 @@ -1169,6 +1170,11 @@ def from_config(cls): """Create cache from configuration.""" session = Session() session.verify = not DISABLE_INDEX_SSL_VERIFICATION + adapter = requests.adapters.HTTPAdapter( + pool_connections=POOL_SIZE, pool_maxsize=POOL_SIZE + ) + session.mount("http://", adapter) + session.mount("https://", adapter) proxpi_version = get_proxpi_version() if proxpi_version: session.headers["User-Agent"] = f"proxpi/{proxpi_version}" From 1df7b0ccd469c7e98fc8428c564707353ded6185 Mon Sep 17 00:00:00 2001 From: Mole Shang <135e2@135e2.dev> Date: Fri, 30 Jan 2026 16:11:13 +0800 Subject: [PATCH 3/3] test: add streaming ... And since we now directly serve from server, there would be no redirects on cache fail. --- tests/test_integration.py | 87 ++++++++++++++++++++++++++++++++++++--- tests/test_pypi.py | 36 +++++++++++----- 2 files changed, 107 insertions(+), 16 deletions(-) diff --git a/tests/test_integration.py b/tests/test_integration.py index 0a45181..fa0b058 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -1,18 +1,22 @@ """Test ``proxpi`` server.""" import enum +import time import hashlib import logging import pathlib import warnings +import threading import posixpath import contextlib +import concurrent.futures from urllib import parse as urllib_parse from unittest import mock import pytest import requests import proxpi.server +import proxpi._cache import packaging.specifiers import starlette.applications @@ -475,11 +479,7 @@ def test_download_file_failed(mock_root_index, server, readonly_package_dir): f"{server}/index/numpy/numpy-1.23.1.tar.gz", allow_redirects=False, ) - assert response.status_code // 100 == 3 - url_parsed = urllib_parse.urlsplit(response.headers["location"]) - mock_root_index_parsed = urllib_parse.urlsplit(mock_root_index) - assert url_parsed.netloc == mock_root_index_parsed.netloc - assert posixpath.split(url_parsed.path)[1] == "numpy-1.23.1.tar.gz" + assert response.status_code == 200 @pytest.mark.parametrize("file_mime_type", ["application/octet-stream", None]) @@ -502,3 +502,80 @@ def test_download_file_representation(server, tmp_path, file_mime_type): assert response.headers["Content-Type"] == "application/x-tar+gzip" assert not response.headers.get("Content-Encoding") response.close() + + +def test_streaming_success(server, tmp_path): + """Test multiple clients streaming the same file concurrently.""" + chunk_size = 1024 + file_content = bytes.fromhex("DEADBEEF") * chunk_size + + download_file = tmp_path / "downloading_file" + # Create the file so it exists for subscribers + download_file.touch() + + status = proxpi.server._cache.DownloadStatus() + status._temp_download_path = str(download_file) + status._upstream_status_code = 200 + status._upstream_headers = [("Content-Type", "application/octet-stream")] + + def simulate_download(): + time.sleep(1.0) + for i in range(0, len(file_content), chunk_size): + chunk = file_content[i : i + chunk_size] + status.broadcast_chunk(chunk) + time.sleep(0.01) + status.broadcast_eof() + + def client_request(): + return requests.get(f"{server}/index/p/f", stream=True) + + with mock.patch.object( + proxpi.server.cache, + "get_file", + side_effect=proxpi.server._cache.Downloading(status), + ): + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + # Start download simulation + downloader = executor.submit(simulate_download) + + # Start clients + futures = [executor.submit(client_request) for _ in range(3)] + + downloader.result() + responses = [f.result() for f in futures] + + for resp in responses: + assert resp.status_code == 200 + assert resp.content == file_content + + +def test_streaming_failure(server, tmp_path): + """Test streaming failure propagation.""" + download_file = tmp_path / "fail_file" + download_file.touch() + + status = proxpi.server._cache.DownloadStatus() + status._temp_download_path = str(download_file) + status._upstream_status_code = 200 + status._upstream_headers = [("Content-Type", "application/octet-stream")] + + def simulate_fail(): + time.sleep(1.0) + status.broadcast_chunk(b"start") + time.sleep(0.1) + status.broadcast_eof(requests.exceptions.ConnectionError("Oh no")) + + with mock.patch.object( + proxpi.server.cache, + "get_file", + side_effect=proxpi.server._cache.Downloading(status), + ): + # Start simulation in background + t = threading.Thread(target=simulate_fail) + t.start() + + resp = requests.get(f"{server}/index/p/f", stream=True) + content = resp.content + t.join() + + assert b"Oh no" in content diff --git a/tests/test_pypi.py b/tests/test_pypi.py index aa17719..31bab91 100644 --- a/tests/test_pypi.py +++ b/tests/test_pypi.py @@ -3,6 +3,7 @@ import sys import logging import subprocess +import concurrent.futures import pytest @@ -19,28 +20,41 @@ def server(): yield from _utils.make_server(proxpi.server.app) -def test_pip_download(server, tmp_path): - """Test package installation.""" +def run_pip(server, dest, pkgs): args = [ sys.executable, "-m", "pip", "--no-cache-dir", "download", - "--index-url", f"{server}/index/", + "--index-url", + f"{server}/index/", ] - - p = subprocess.run( - [*args, "--dest", str(tmp_path / "dest1"), "Jinja2", "marshmallow"], - ) + p = subprocess.run([*args, "--dest", str(dest), *pkgs], check=True) assert p.returncode == 0 - contents = list((tmp_path / "dest1").iterdir()) + return list(dest.iterdir()) + + +def test_pip_download(server, tmp_path): + """Test package installation.""" + contents = run_pip(server, tmp_path / "dest1", ["Jinja2", "marshmallow"]) print(contents) assert any("jinja2" in p.name.lower() for p in contents) assert any("marshmallow" in p.name.lower() for p in contents) - subprocess.run([*args, "--dest", str(tmp_path / "dest2"), "Jinja2"]) - assert p.returncode == 0 - contents = list((tmp_path / "dest2").iterdir()) + contents = run_pip(server, tmp_path / "dest2", ["Jinja2"]) print(contents) assert any("jinja2" in p.name.lower() for p in contents) + + +def test_concurrent_pip_download(server, tmp_path): + """Test concurrent package installation.""" + dest = tmp_path / "concurrent_dest" + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(run_pip, server, dest / str(i), ["flask"]) for i in range(4) + ] + results = [f.result() for f in futures] + + for contents in results: + assert any("flask" in p.name.lower() for p in contents)