diff --git a/kittens/transfer/librsync.py b/kittens/transfer/librsync.py index 60226188b..57fddca20 100644 --- a/kittens/transfer/librsync.py +++ b/kittens/transfer/librsync.py @@ -16,6 +16,8 @@ if TYPE_CHECKING: class StreamingJob: + expected_input_size = IO_BUFFER_SIZE + def __init__(self, job: 'JobCapsule', output_buf_size: int = IO_BUFFER_SIZE): self.job = job self.finished = False @@ -49,9 +51,9 @@ class StreamingJob: pass -def drive_job_on_file(f: IO[bytes], job: 'JobCapsule') -> Iterator[memoryview]: - sj = StreamingJob(job) - input_buf = bytearray(IO_BUFFER_SIZE) +def drive_job_on_file(f: IO[bytes], job: 'JobCapsule', input_buf_size: int = IO_BUFFER_SIZE, output_buf_size: int = IO_BUFFER_SIZE) -> Iterator[memoryview]: + sj = StreamingJob(job, output_buf_size=output_buf_size) + input_buf = bytearray(input_buf_size) while not sj.finished: sz = f.readinto(input_buf) # type: ignore yield sj(memoryview(input_buf)[:sz]) @@ -61,13 +63,18 @@ def signature_of_file(path: str) -> Iterator[memoryview]: with open(path, 'rb') as f: f.seek(0, os.SEEK_END) fsz = f.tell() - job = begin_create_signature(fsz) + job, block_len, strong_len = begin_create_signature(fsz) + strong_len = max(0, strong_len) f.seek(0) - yield from drive_job_on_file(f, job) + # see whole.c in librsync source for size calculations + yield from drive_job_on_file(f, job, input_buf_size=4 * block_len, output_buf_size=12 + 4 * (4 + (strong_len or IO_BUFFER_SIZE))) class LoadSignature(StreamingJob): + # see whole.c in librsync source for size calculations + expected_input_size = 16 * 1024 + def __init__(self) -> None: job, self.signature = begin_load_signature() super().__init__(job, output_buf_size=0) @@ -79,15 +86,19 @@ class LoadSignature(StreamingJob): def delta_for_file(path: str, sig: 'SignatureCapsule') -> Iterator[memoryview]: job = begin_create_delta(sig) with open(path, 'rb') as f: - yield from drive_job_on_file(f, job) + # see whole.c in librsync source for size calculations + yield from drive_job_on_file(f, job, input_buf_size=8 * IO_BUFFER_SIZE, output_buf_size=4 * IO_BUFFER_SIZE) class PatchFile(StreamingJob): + # see whole.c in librsync source for size calculations + expected_input_size = IO_BUFFER_SIZE + def __init__(self, src_path: str): self.src_file = open(src_path, 'rb') job = begin_patch(self.read_from_src) - super().__init__(job) + super().__init__(job, output_buf_size=4 * IO_BUFFER_SIZE) def read_from_src(self, b: memoryview, pos: int) -> int: self.src_file.seek(pos) diff --git a/kittens/transfer/rsync.c b/kittens/transfer/rsync.c index fdf717930..0e1691bb4 100644 --- a/kittens/transfer/rsync.c +++ b/kittens/transfer/rsync.c @@ -64,7 +64,7 @@ begin_create_signature(PyObject *self UNUSED, PyObject *args) { } #endif CREATE_JOB(rs_sig_begin, NULL, block_len, strong_len, magic_number); - return job_capsule; + return Py_BuildValue("Nnn", job_capsule, (Py_ssize_t)block_len, (Py_ssize_t)strong_len); } #define GET_JOB_FROM_CAPSULE \ diff --git a/kittens/transfer/rsync.pyi b/kittens/transfer/rsync.pyi index 09be9b19b..a31915082 100644 --- a/kittens/transfer/rsync.pyi +++ b/kittens/transfer/rsync.pyi @@ -15,7 +15,7 @@ class RsyncError(Exception): pass -def begin_create_signature(file_size: int = -1, strong_len: int = 0) -> JobCapsule: +def begin_create_signature(file_size: int = -1, strong_len: int = 0) -> Tuple[JobCapsule, int, int]: pass