diff --git a/requirements.txt b/requirements.txt index 3573f8b..6a943f0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,5 @@ plumbum==1.9.0 requests>=2.31.0 -filesplit==3.0.2 pulpcore-client==3.68.0 scikit_build==0.18.1 cerberus==1.3.5 diff --git a/sign_node/config.py b/sign_node/config.py index b1a4129..b091f86 100644 --- a/sign_node/config.py +++ b/sign_node/config.py @@ -18,6 +18,7 @@ DEFAULT_PULP_USER = "pulp" DEFAULT_PULP_PASSWORD = "test_pwd" DEFAULT_PULP_CHUNK_SIZE = 8388608 # 8 MiB +DEFAULT_UPLOAD_WORKERS = 4 # Max file size to allow parallel upload for DEFAULT_PARALLEL_FILE_UPLOAD_SIZE = 524288000 # 500 MB DEFAULT_PGP_PASSWORD = "test_pwd" @@ -64,6 +65,7 @@ def __init__(self, config_file=None, **cmd_args): "pulp_user": DEFAULT_PULP_USER, "pulp_password": DEFAULT_PULP_PASSWORD, "pulp_chunk_size": DEFAULT_PULP_CHUNK_SIZE, + "upload_workers": DEFAULT_UPLOAD_WORKERS, "parallel_upload_file_size": DEFAULT_PARALLEL_FILE_UPLOAD_SIZE, "dev_pgp_key_password": DEFAULT_PGP_PASSWORD, 'sentry_dsn': DEFAULT_SENTRY_DSN, @@ -90,6 +92,7 @@ def __init__(self, config_file=None, **cmd_args): "pulp_user": {"type": "string", "nullable": False}, "pulp_password": {"type": "string", "nullable": False}, "pulp_chunk_size": {"type": "integer", "nullable": False}, + "upload_workers": {"type": "integer", "nullable": False}, "parallel_upload_file_size": {"type": "integer", "nullable": False}, "jwt_token": {"type": "string", "required": True}, "dev_pgp_key_password": {"type": "string", "nullable": False}, diff --git a/sign_node/signer.py b/sign_node/signer.py index 8499258..4e9faab 100644 --- a/sign_node/signer.py +++ b/sign_node/signer.py @@ -63,6 +63,7 @@ def __init__(self, config, password_db, gpg): self.__config.pulp_user, self.__config.pulp_password, self.__config.pulp_chunk_size, + upload_workers=self.__config.upload_workers, ) self.__working_dir_path = Path(self.__config.working_dir) self.__download_credentials = { diff --git a/sign_node/uploaders/pulp.py b/sign_node/uploaders/pulp.py index 58f8f4a..4c4b377 100644 --- a/sign_node/uploaders/pulp.py +++ b/sign_node/uploaders/pulp.py @@ -1,13 +1,11 @@ -import csv import logging +import math import os import tempfile import time -import shutil import typing +from concurrent.futures import ThreadPoolExecutor, as_completed from typing import List - -from fsplit.filesplit import Filesplit from pulpcore.client.pulpcore.configuration import Configuration from pulpcore.client.pulpcore.api_client import ApiClient from pulpcore.client.pulpcore.api.tasks_api import TasksApi @@ -31,7 +29,8 @@ class PulpBaseUploader(BaseUploader): Handles uploads to Pulp server. """ - def __init__(self, host: str, username: str, password: str, chunk_size: int): + def __init__(self, host: str, username: str, password: str, chunk_size: int, + upload_workers: int = 4): """ Initiate uploader. @@ -45,13 +44,15 @@ def __init__(self, host: str, username: str, password: str, chunk_size: int): User password. chunk_size : int Size of chunk to split files during the upload. + upload_workers : int + Maximum number of concurrent workers for chunked uploads. """ api_client = self._prepare_api_client(host, username, password) self._uploads_client = UploadsApi(api_client=api_client) self._tasks_client = TasksApi(api_client=api_client) self._artifacts_client = ArtifactsApi(api_client=api_client) - self._file_splitter = Filesplit() self._chunk_size = chunk_size + self._upload_workers = upload_workers self._logger = logging.getLogger(__file__) @staticmethod @@ -137,38 +138,88 @@ def _commit_upload(self, reference: str, file_sha256: str) -> str: task_result = self._wait_for_task_completion(response.task) return task_result.created_resources[0] + def _create_artifact_direct(self, file_path: str, file_sha256: str) -> str: + response = self._artifacts_client.create(file_path, sha256=file_sha256) + return response.pulp_href + def _put_large_file(self, file_path: str, reference: str): - temp_dir = tempfile.mkdtemp(prefix="pulp_uploader_") - try: - lower_bytes_limit = 0 - total_size = os.path.getsize(file_path) - self._file_splitter.split(file_path, self._chunk_size, output_dir=temp_dir) - manifest_path = os.path.join(temp_dir, 'fs_manifest.csv') - for meta in csv.DictReader(open(manifest_path, 'r')): - split_file_path = os.path.join(temp_dir, meta['filename']) - upper_bytes_limit = lower_bytes_limit + int(meta['filesize']) - 1 + total_size = os.path.getsize(file_path) + + # Build list of (offset, length) chunk descriptors + chunks = [] + offset = 0 + while offset < total_size: + length = min(self._chunk_size, total_size - offset) + chunks.append((offset, length)) + offset += length + + def _upload_chunk(chunk_offset: int, chunk_length: int): + """Read a byte range from source and upload via a temp file.""" + tmp_path = None + try: + with open(file_path, 'rb') as src: + src.seek(chunk_offset) + data = src.read(chunk_length) + tmp = tempfile.NamedTemporaryFile(delete=False, + prefix='pulp_chunk_') + tmp_path = tmp.name + tmp.write(data) + tmp.close() + end_byte = chunk_offset + chunk_length - 1 + content_range = ( + f'bytes {chunk_offset}-{end_byte}/{total_size}' + ) self._uploads_client.update( - f'bytes {lower_bytes_limit}-{upper_bytes_limit}/' - f'{total_size}', - reference, split_file_path) - lower_bytes_limit += int(meta['filesize']) - finally: - if temp_dir and os.path.exists(temp_dir): - shutil.rmtree(temp_dir) + content_range, reference, tmp_path, + ) + finally: + if tmp_path and os.path.exists(tmp_path): + os.unlink(tmp_path) + + with ThreadPoolExecutor(max_workers=self._upload_workers) as executor: + futures = { + executor.submit(_upload_chunk, off, length): (off, length) + for off, length in chunks + } + for future in as_completed(futures): + future.result() # propagates any exception def _send_file(self, file_path: str) -> typing.Tuple[str, str]: file_sha256 = hash_file(file_path, hash_type="sha256") reference = self.check_if_artifact_exists(file_sha256) if reference: return file_sha256, reference - reference, file_size = self._create_upload(file_path) + file_size = os.path.getsize(file_path) + file_name = os.path.basename(file_path) + start_time = time.time() + if file_size < self._chunk_size: + artifact_href = self._create_artifact_direct(file_path, file_sha256) + elapsed = time.time() - start_time + self._logger.info( + 'Upload complete: %s (%d bytes) via direct artifact in %.2fs', + file_name, file_size, elapsed, + ) + return file_sha256, artifact_href + reference, _ = self._create_upload(file_path) if file_size > self._chunk_size: self._put_large_file(file_path, reference) + num_chunks = math.ceil(file_size / self._chunk_size) + artifact_href = self._commit_upload(reference, file_sha256) + elapsed = time.time() - start_time + self._logger.info( + 'Upload complete: %s (%d bytes) via %d chunks in %.2fs', + file_name, file_size, num_chunks, elapsed, + ) else: self._uploads_client.update( f"bytes 0-{file_size - 1}/{file_size}", reference, file_path ) - artifact_href = self._commit_upload(reference, file_sha256) + artifact_href = self._commit_upload(reference, file_sha256) + elapsed = time.time() - start_time + self._logger.info( + 'Upload complete: %s (%d bytes) via single chunk in %.2fs', + file_name, file_size, elapsed, + ) return file_sha256, artifact_href def check_if_artifact_exists(self, sha256: str) -> str: diff --git a/tests/sign_node/uploaders/test_pulp.py b/tests/sign_node/uploaders/test_pulp.py index c03b3fc..5b799ee 100644 --- a/tests/sign_node/uploaders/test_pulp.py +++ b/tests/sign_node/uploaders/test_pulp.py @@ -1,5 +1,7 @@ import operator import os +import re +import threading from unittest.mock import Mock, patch from pyfakefs.fake_filesystem_unittest import TestCase @@ -112,3 +114,357 @@ def read(self, task_href): file_sha256, artifact_href = uploader._send_file(f_path) assert file_sha256 == f_hash assert artifact_href == f_href + + def test_send_file_direct_artifact(self): + """Verify that small files (< chunk_size) use the direct artifact + creation path and never touch UploadsApi.""" + f_path = '/build_dir/small.rpm' + self.fs.create_file(f_path, contents='small data') + f_hash = hash_file(f_path, hash_type='sha256') + artifact_href = '/pulp/api/v3/artifacts/direct-small/' + chunk_size = 100 # larger than the 10-byte file + + class MockArtifactsApi: + def __init__(self, *_, **__): + pass + + def create(_, file_path, sha256=None): + assert file_path == f_path + assert sha256 == f_hash + response = Mock() + response.pulp_href = artifact_href + return response + + def list(_, sha256=None): + response = Mock() + response.results = [] + return response + + mock_uploads = Mock() + mock_tasks = Mock() + + with ( + patch('sign_node.uploaders.pulp.ArtifactsApi', + new=MockArtifactsApi), + patch('sign_node.uploaders.pulp.UploadsApi', + return_value=mock_uploads), + patch('sign_node.uploaders.pulp.TasksApi', + return_value=mock_tasks), + ): + uploader = PulpRpmUploader('localhost', 'user', 'password', + chunk_size) + result_hash, result_href = uploader._send_file(f_path) + + assert result_hash == f_hash + assert result_href == artifact_href + # UploadsApi methods must never be called for small files + assert not mock_uploads.create.called + assert not mock_uploads.update.called + assert not mock_uploads.commit.called + + def test_send_large_file(self): + """Verify parallel chunked upload produces correct Content-Range + headers and the right number of update() calls.""" + f_path = '/build_dir/large.rpm' + self.fs.create_file(f_path, contents='x' * 100) + f_hash = hash_file(f_path, hash_type='sha256') + f_size = 100 + chunk_size = 30 # 4 chunks: 30+30+30+10 + upload_href = '/pulp/api/v3/uploads/large-upload/' + artifact_href = '/pulp/api/v3/artifacts/large-artifact/' + + # Thread-safe list to record all update() calls + update_calls = [] + lock = threading.Lock() + + class MockUploadsApi: + def __init__(self, *_, **__): + pass + + def create(_, opts): + assert opts['size'] == f_size + response = Mock() + response.pulp_href = upload_href + return response + + def update(_, content_range, href, file_path): + with lock: + update_calls.append(content_range) + assert href == upload_href + + def commit(_, href, upload_commit): + assert href == upload_href + assert upload_commit['sha256'] == f_hash + response = Mock() + response.task = 'task-large' + return response + + class MockTasksApi: + def __init__(self, *_, **__): + pass + + def read(self, task_href): + assert task_href == 'task-large' + result = Mock() + result.created_resources = [artifact_href] + result.state = 'completed' + return result + + with ( + patch('sign_node.uploaders.pulp.UploadsApi', new=MockUploadsApi), + patch('sign_node.uploaders.pulp.TasksApi', new=MockTasksApi), + patch.object(PulpRpmUploader, 'check_if_artifact_exists', + return_value=None), + ): + uploader = PulpRpmUploader('localhost', 'user', 'password', + chunk_size) + file_sha256, result_href = uploader._send_file(f_path) + + # Correct return values + assert file_sha256 == f_hash + assert result_href == artifact_href + + # Exactly 4 chunks uploaded + assert len(update_calls) == 4 + + # Parse Content-Range headers and verify full coverage + # Expected: bytes 0-29/100, bytes 30-59/100, bytes 60-89/100, + # bytes 90-99/100 + pattern = re.compile(r'bytes (\d+)-(\d+)/(\d+)') + ranges = [] + for cr in update_calls: + m = pattern.match(cr) + assert m, f'Invalid Content-Range header: {cr}' + start, end, total = int(m.group(1)), int(m.group(2)), int(m.group(3)) + assert total == f_size + ranges.append((start, end)) + + # Sort by start offset and verify contiguous, non-overlapping coverage + ranges.sort() + assert ranges == [(0, 29), (30, 59), (60, 89), (90, 99)] + + def test_send_large_file_chunk_error(self): + """Verify that if one chunk upload raises, the exception propagates.""" + f_path = '/build_dir/error.rpm' + self.fs.create_file(f_path, contents='y' * 100) + f_size = 100 + chunk_size = 30 + upload_href = '/pulp/api/v3/uploads/error-upload/' + + call_count = {'n': 0} + count_lock = threading.Lock() + + class MockUploadsApi: + def __init__(self, *_, **__): + pass + + def create(_, opts): + response = Mock() + response.pulp_href = upload_href + return response + + def update(_, content_range, href, file_path): + with count_lock: + call_count['n'] += 1 + if call_count['n'] == 2: + raise RuntimeError('chunk upload failed') + + def commit(_, href, upload_commit): + response = Mock() + response.task = 'task-err' + return response + + class MockTasksApi: + def __init__(self, *_, **__): + pass + + def read(self, task_href): + result = Mock() + result.created_resources = ['/artifact/'] + result.state = 'completed' + return result + + with ( + patch('sign_node.uploaders.pulp.UploadsApi', new=MockUploadsApi), + patch('sign_node.uploaders.pulp.TasksApi', new=MockTasksApi), + patch.object(PulpRpmUploader, 'check_if_artifact_exists', + return_value=None), + ): + uploader = PulpRpmUploader('localhost', 'user', 'password', + chunk_size) + try: + uploader._send_file(f_path) + assert False, 'Expected RuntimeError to propagate' + except RuntimeError as e: + assert 'chunk upload failed' in str(e) + + def test_configurable_upload_workers(self): + """Verify that a custom upload_workers value reaches ThreadPoolExecutor.""" + from concurrent.futures import ThreadPoolExecutor as RealTPE + + f_path = '/build_dir/workers.rpm' + self.fs.create_file(f_path, contents='w' * 100) + f_hash = hash_file(f_path, hash_type='sha256') + f_size = 100 + chunk_size = 30 # triggers large file path (100 > 30) + upload_href = '/pulp/api/v3/uploads/workers-upload/' + artifact_href = '/pulp/api/v3/artifacts/workers-artifact/' + + captured_max_workers = {} + + class MockUploadsApi: + def __init__(self, *_, **__): + pass + + def create(_, opts): + response = Mock() + response.pulp_href = upload_href + return response + + def update(_, content_range, href, file_path): + pass + + def commit(_, href, upload_commit): + response = Mock() + response.task = 'task-workers' + return response + + class MockTasksApi: + def __init__(self, *_, **__): + pass + + def read(self, task_href): + result = Mock() + result.created_resources = [artifact_href] + result.state = 'completed' + return result + + def spy_tpe(*args, **kwargs): + captured_max_workers['value'] = kwargs.get('max_workers') + return RealTPE(*args, **kwargs) + + with ( + patch('sign_node.uploaders.pulp.UploadsApi', new=MockUploadsApi), + patch('sign_node.uploaders.pulp.TasksApi', new=MockTasksApi), + patch.object(PulpRpmUploader, 'check_if_artifact_exists', + return_value=None), + patch('sign_node.uploaders.pulp.ThreadPoolExecutor', + side_effect=spy_tpe), + ): + uploader = PulpRpmUploader('localhost', 'user', 'password', + chunk_size, upload_workers=2) + uploader._send_file(f_path) + + # ThreadPoolExecutor was called with the custom value + assert captured_max_workers.get('value') == 2 + + def test_default_upload_workers(self): + """Verify the default upload_workers is 4 when not explicitly set.""" + with ( + patch('sign_node.uploaders.pulp.UploadsApi'), + patch('sign_node.uploaders.pulp.TasksApi'), + patch('sign_node.uploaders.pulp.ArtifactsApi'), + ): + uploader = PulpRpmUploader('localhost', 'user', 'password', 42) + assert uploader._upload_workers == 4 + + def test_upload_timing_log_direct(self): + """Verify timing log is emitted for direct artifact upload path.""" + f_path = '/build_dir/timed_small.rpm' + self.fs.create_file(f_path, contents='tiny') + f_hash = hash_file(f_path, hash_type='sha256') + artifact_href = '/pulp/api/v3/artifacts/timed-small/' + chunk_size = 100 # larger than file -> direct artifact path + + class MockArtifactsApi: + def __init__(self, *_, **__): + pass + + def create(_, file_path, sha256=None): + response = Mock() + response.pulp_href = artifact_href + return response + + def list(_, sha256=None): + response = Mock() + response.results = [] + return response + + with ( + patch('sign_node.uploaders.pulp.ArtifactsApi', + new=MockArtifactsApi), + patch('sign_node.uploaders.pulp.UploadsApi'), + patch('sign_node.uploaders.pulp.TasksApi'), + self.assertLogs(level='INFO') as cm, + ): + uploader = PulpRpmUploader('localhost', 'user', 'password', + chunk_size) + uploader._send_file(f_path) + + # Find the timing log line + timing_logs = [m for m in cm.output if 'Upload complete' in m] + assert len(timing_logs) >= 1, f'Expected timing log, got: {cm.output}' + log_line = timing_logs[0] + assert 'direct artifact' in log_line + assert 'timed_small.rpm' in log_line + # Verify a time value is present (e.g. "0.00s" or "1.23s") + assert re.search(r'\d+\.\d+s', log_line), \ + f'Expected elapsed time in log: {log_line}' + + def test_upload_timing_log_chunked(self): + """Verify timing log is emitted for chunked upload path.""" + f_path = '/build_dir/timed_large.rpm' + self.fs.create_file(f_path, contents='z' * 100) + f_hash = hash_file(f_path, hash_type='sha256') + f_size = 100 + chunk_size = 30 # triggers large file path + upload_href = '/pulp/api/v3/uploads/timed-large/' + artifact_href = '/pulp/api/v3/artifacts/timed-large/' + + class MockUploadsApi: + def __init__(self, *_, **__): + pass + + def create(_, opts): + response = Mock() + response.pulp_href = upload_href + return response + + def update(_, content_range, href, file_path): + pass + + def commit(_, href, upload_commit): + response = Mock() + response.task = 'task-timed' + return response + + class MockTasksApi: + def __init__(self, *_, **__): + pass + + def read(self, task_href): + result = Mock() + result.created_resources = [artifact_href] + result.state = 'completed' + return result + + with ( + patch('sign_node.uploaders.pulp.UploadsApi', new=MockUploadsApi), + patch('sign_node.uploaders.pulp.TasksApi', new=MockTasksApi), + patch.object(PulpRpmUploader, 'check_if_artifact_exists', + return_value=None), + self.assertLogs(level='INFO') as cm, + ): + uploader = PulpRpmUploader('localhost', 'user', 'password', + chunk_size) + uploader._send_file(f_path) + + # Find the timing log line + timing_logs = [m for m in cm.output if 'Upload complete' in m] + assert len(timing_logs) >= 1, f'Expected timing log, got: {cm.output}' + log_line = timing_logs[0] + assert 'chunks' in log_line + assert 'timed_large.rpm' in log_line + # Verify a time value is present + assert re.search(r'\d+\.\d+s', log_line), \ + f'Expected elapsed time in log: {log_line}'