copy_file_obj_fast attempt

Author: Fallen-Breath

prime_backup/action/export_backup_action_zip.py

@@ -1,6 +1,5 @@
 import contextlib
 import os
-import shutil
 import stat
 import time
 import zipfile
@@ -14,6 +13,7 @@
 from prime_backup.constants.constants import BACKUP_META_FILE_NAME
 from prime_backup.db import schema
 from prime_backup.types.export_failure import ExportFailures
+from prime_backup.utils import file_utils
 from prime_backup.utils.io_types import SupportsReadBytes
 
 
@@ -56,7 +56,7 @@ def __export_file(self, blob_chunks_getter: BlobChunksGetter, zipf: zipfile.ZipF
 
 			def reader_csm(reader: SupportsReadBytes):
 				with zipf.open(info, 'w') as zip_item:
-					shutil.copyfileobj(reader, zip_item)
+					file_utils.copy_file_obj_fast(reader, zip_item, estimate_read_size=file.blob_raw_size)
 
 			self._create_blob_exporter(blob_chunks_getter, file).export_as_reader(reader_csm)
 		elif stat.S_ISDIR(file.mode):

prime_backup/action/helpers/blob_allocator.py

@@ -342,15 +342,15 @@ def get_blob_path_for_write(h: str) -> Path:
 
 				blob_path = get_blob_path_for_write(blob_hash)
 				with self.__time_costs.measure_time_cost(CreateBackupTimeCostKey.kind_io_copy):
-					cr = compressor.copy_compressed(temp_file_path, blob_path, calc_hash=False)
+					cr = compressor.copy_compressed(temp_file_path, blob_path, calc_hash=False, estimate_read_size=st.st_size)
 				raw_size, stored_size = cr.read_size, cr.write_size
 
 		elif policy == _DirectBlobCreatePolicy.hash_once:
 			# read once, compress+hash to temp file, then move
 			misc_utils.assert_true(blob_hash is None, 'blob_hash should not be calculated')
 			with self.__make_temp_file(src_path_md5) as temp_file_path:
 				with self.__time_costs.measure_time_cost(CreateBackupTimeCostKey.kind_io_copy):
-					cr = compressor.copy_compressed(src_path, temp_file_path, calc_hash=True, open_r_func=SourceFileNotFoundWrapper.open_rb)
+					cr = compressor.copy_compressed(src_path, temp_file_path, calc_hash=True, estimate_read_size=st.st_size, open_r_func=SourceFileNotFoundWrapper.open_rb)
 				check_changes(cr.read_size, None)  # the size must be unchanged, to satisfy the uniqueness
 
 				raw_size, blob_hash, stored_size = cr.read_size, cr.read_hash, cr.write_size
@@ -391,7 +391,7 @@ def get_blob_path_for_write(h: str) -> Path:
 				else:
 					# copy+compress+hash to blob store
 					with self.__time_costs.measure_time_cost(CreateBackupTimeCostKey.kind_io_copy):
-						cr = compressor.copy_compressed(src_path, blob_path, calc_hash=True, open_r_func=SourceFileNotFoundWrapper.open_rb)
+						cr = compressor.copy_compressed(src_path, blob_path, calc_hash=True, estimate_read_size=st.st_size, open_r_func=SourceFileNotFoundWrapper.open_rb)
 					raw_size, stored_size = cr.read_size, cr.write_size
 					actual_sah = SizeAndHash(cr.read_size, cr.read_hash)
 				check_changes(actual_sah.size, actual_sah.hash)

prime_backup/action/helpers/blob_exporter.py

@@ -1,7 +1,6 @@
 import contextlib
 import dataclasses
 import functools
-import shutil
 import threading
 from abc import ABC, abstractmethod
 from pathlib import Path
@@ -152,9 +151,9 @@ def __export_to_fs_direct(self, output_path: Path):
 				with open(output_path, 'wb') as f_out:
 					if self.verify_blob:
 						bypass_reader = BypassReader(f_in, calc_hash=True)
-						shutil.copyfileobj(bypass_reader, f_out)
+						file_utils.copy_file_obj_fast(bypass_reader, f_out, estimate_read_size=self.blob.raw_size)
 					else:
-						shutil.copyfileobj(f_in, f_out)
+						file_utils.copy_file_obj_fast(f_in, f_out, estimate_read_size=self.blob.raw_size)
 			if self.verify_blob and bypass_reader is not None:
 				self.__verify_exported_blob(bypass_reader.get_read_len(), bypass_reader.get_hash())
 
@@ -169,9 +168,9 @@ def __export_to_fs_chunked(self, output_path: Path):
 				with compressor.open_decompressed(chunk_path) as f_in:
 					if self.verify_blob:
 						bypass_reader = BypassReader(f_in, calc_hash=True, hash_method=chunk_utils.get_hash_method())
-						shutil.copyfileobj(bypass_reader, f_out)
+						file_utils.copy_file_obj_fast(bypass_reader, f_out, estimate_read_size=oc.chunk.raw_size)
 					else:
-						shutil.copyfileobj(f_in, f_out)
+						file_utils.copy_file_obj_fast(f_in, f_out, estimate_read_size=oc.chunk.raw_size)
 
 				if self.verify_blob and bypass_reader is not None:
 					self.__verify_exported_chunk(oc.chunk, bypass_reader.get_read_len(), bypass_reader.get_hash())

prime_backup/action/import_backup_action.py

@@ -1,5 +1,4 @@
 import json
-import shutil
 from pathlib import Path
 from typing import IO, Optional, List, Dict, Tuple
 
@@ -24,7 +23,7 @@
 from prime_backup.types.standalone_backup_format import StandaloneBackupFormat
 from prime_backup.types.tar_format import TarFormat
 from prime_backup.types.units import ByteCount
-from prime_backup.utils import blob_utils, misc_utils, chunk_utils, collection_utils
+from prime_backup.utils import blob_utils, misc_utils, chunk_utils, collection_utils, file_utils
 from prime_backup.utils.hash_utils import SizeAndHash
 
 
@@ -68,7 +67,7 @@ def __create_blob_file(self, file_reader: IO[bytes], sah: SizeAndHash) -> Tuple[
 		compress_method: CompressMethod = self.config.backup.get_compress_method_from_size(sah.size)
 		compressor = Compressor.create(compress_method)
 		with compressor.open_compressed_bypassed(blob_path) as (writer, f):
-			shutil.copyfileobj(file_reader, f)
+			file_utils.copy_file_obj_fast(file_reader, f, estimate_read_size=sah.size)
 
 		return writer.get_write_len(), compress_method
 

prime_backup/action/migrate_compress_method_action.py

@@ -1,4 +1,3 @@
-import shutil
 import time
 from concurrent.futures import Future
 from pathlib import Path
@@ -13,7 +12,7 @@
 from prime_backup.db.session import DbSession
 from prime_backup.db.values import BlobStorageMethod
 from prime_backup.types.size_diff import SizeDiff
-from prime_backup.utils import blob_utils, chunk_utils
+from prime_backup.utils import blob_utils, chunk_utils, file_utils
 from prime_backup.utils.thread_pool import FailFastBlockingThreadPool
 
 _OLD_BLOB_SUFFIX = '_old'
@@ -64,7 +63,7 @@ def __migrate_single_direct_blob(self, blob: schema.Blob) -> bool:
 		blob_path.replace(old_trash_path)
 		with decompressor.open_decompressed(old_trash_path) as f_src:
 			with compressor.open_compressed_bypassed(blob_path) as (writer, f_dst):
-				shutil.copyfileobj(f_src, f_dst)
+				file_utils.copy_file_obj_fast(f_src, f_dst, estimate_read_size=blob.stored_size)
 
 		blob.compress = new_compress_method.name
 		blob.stored_size = writer.get_write_len()
@@ -101,7 +100,7 @@ def __migrate_single_chunk(self, chunk: schema.Chunk) -> bool:
 		chunk_path.replace(old_trash_path)
 		with decompressor.open_decompressed(old_trash_path) as f_src:
 			with compressor.open_compressed_bypassed(chunk_path) as (writer, f_dst):
-				shutil.copyfileobj(f_src, f_dst)
+				file_utils.copy_file_obj_fast(f_src, f_dst, estimate_read_size=chunk.stored_size)
 
 		chunk.compress = new_compress_method.name
 		chunk.stored_size = writer.get_write_len()

prime_backup/compressors.py

@@ -1,12 +1,12 @@
 import contextlib
 import dataclasses
 import enum
-import shutil
 from abc import abstractmethod, ABC
 from typing import BinaryIO, Union, ContextManager, Tuple, Callable, Literal, Generator
 
 from typing_extensions import Protocol, override
 
+from prime_backup.utils import file_utils
 from prime_backup.utils.bypass_io import BypassReader, BypassWriter
 from prime_backup.utils.path_like import PathLike
 
@@ -43,6 +43,7 @@ def ensure_lib(cls):
 	def copy_compressed(
 			self, source_path: PathLike, dest_path: PathLike, *,
 			calc_hash: bool = False,
+			estimate_read_size: int = 0,
 			open_r_func: Callable[[PathLike, "Literal['rb']"], BinaryIO] = open,
 			open_w_func: Callable[[PathLike, "Literal['wb']"], BinaryIO] = open,
 	) -> CopyCompressResult:
@@ -52,7 +53,7 @@ def copy_compressed(
 		with open_r_func(source_path, 'rb') as f_in, open_w_func(dest_path, 'wb') as f_out:
 			reader = BypassReader(f_in, calc_hash=calc_hash)
 			writer = BypassWriter(f_out)
-			self._copy_compressed(reader, writer)
+			self._copy_compressed(reader, writer, estimate_read_size=estimate_read_size)
 			return self.CopyCompressResult(reader.get_read_len(), reader.get_hash(), writer.get_write_len())
 
 	def copy_decompressed(
@@ -106,19 +107,19 @@ def open_decompressed_bypassed(self, source_path: PathLike) -> Generator[Tuple[B
 			with self.decompress_stream(reader) as f_decompressed:
 				yield reader, f_decompressed
 
-	def _copy_compressed(self, f_in: BinaryIO, f_out: BinaryIO):
+	def _copy_compressed(self, f_in: BinaryIO, f_out: BinaryIO, *, estimate_read_size: int = 0):
 		"""
 		(f_in) --[compress]--> (f_out)
 		"""
 		with self.compress_stream(f_out) as compressed_out:
-			shutil.copyfileobj(f_in, compressed_out)
+			file_utils.copy_file_obj_fast(f_in, compressed_out, estimate_read_size=estimate_read_size)
 
-	def _copy_decompressed(self, f_in: BinaryIO, f_out: BinaryIO):
+	def _copy_decompressed(self, f_in: BinaryIO, f_out: BinaryIO, *, estimate_read_size: int = 0):
 		"""
 		(f_in) --[decompress]--> (f_out)
 		"""
 		with self.decompress_stream(f_in) as compressed_in:
-			shutil.copyfileobj(compressed_in, f_out)
+			file_utils.copy_file_obj_fast(compressed_in, f_out, estimate_read_size=estimate_read_size)
 
 	@abstractmethod
 	def compress_stream(self, f_out: BinaryIO) -> ContextManager[BinaryIO]:

prime_backup/db/db_file_backup.py

@@ -1,8 +1,8 @@
 import contextlib
-import shutil
 from pathlib import Path
 from typing import Generator
 
+from prime_backup.utils import file_utils
 from prime_backup.utils.io_types import SupportsReadBytes, SupportsWriteBytes
 
 
@@ -46,11 +46,11 @@ def create(self, skip_existing: bool):
 		if skip_existing and self.backup_file.is_file():
 			return
 		with open(self.src_file, 'rb') as f_src, self.__open_backup_for_write() as f_dst:
-			shutil.copyfileobj(f_src, f_dst)
+			file_utils.copy_file_obj_fast(f_src, f_dst)
 
 	def restore(self):
 		with open(self.src_file, 'wb') as f_dst, self.__open_backup_for_read() as f_src:
-			shutil.copyfileobj(f_src, f_dst)
+			file_utils.copy_file_obj_fast(f_src, f_dst)
 
 	def delete_all(self):
 		self.backup_file.unlink(missing_ok=True)

prime_backup/utils/file_utils.py

@@ -1,13 +1,17 @@
 import errno
+import functools
 import os
+import queue
 import shutil
 import stat
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Optional, Literal, BinaryIO, Callable
 
 import psutil
 
 from prime_backup.utils import path_utils
+from prime_backup.utils.io_types import SupportsReadBytes, SupportsWriteBytes
 
 HAS_COPY_FILE_RANGE = callable(getattr(os, 'copy_file_range', None))
 
@@ -46,6 +50,54 @@ def copy_file_fast(
 	shutil.copyfile(src_path, dst_path, follow_symlinks=False)
 
 
+class _ThreadedFastFileObjCopier:
+	COPY_BUFSIZE = 1024 * 1024 if os.name == 'nt' else 4 * 1024  # == shutil.COPY_BUFSIZE
+	MEMORY_CACHE_SIZE = 8 * 1048576
+
+	def __init__(self, concurrency: int):
+		self.concurrency = concurrency
+		self.thread_pool = ThreadPoolExecutor(max_workers=concurrency)
+
+	def copy(self, src: SupportsReadBytes, dst: SupportsWriteBytes):
+		q: 'queue.Queue[Optional[bytes]]' = queue.Queue(maxsize=max(1, self.MEMORY_CACHE_SIZE // self.COPY_BUFSIZE))
+		func_exited = False
+
+		# reference: shutil.copyfileobj
+		q_put = q.put
+		q_get = q.get
+		buf_size = self.COPY_BUFSIZE
+		read_func = src.read
+		write_func = dst.write
+
+		def read_worker():
+			try:
+				while not func_exited and (read_buf := read_func(buf_size)):
+					q_put(read_buf)
+			finally:
+				q_put(None)
+
+		future = self.thread_pool.submit(read_worker)
+		try:
+			while (write_buf := q_get()) is not None:
+				write_func(write_buf)
+			future.result()
+		finally:
+			func_exited = True
+
+
+@functools.lru_cache(None)
+def __get_copier():
+	from prime_backup.config.config import Config
+	return _ThreadedFastFileObjCopier(Config.get().get_effective_concurrency())
+
+
+def copy_file_obj_fast(src: SupportsReadBytes, dst: SupportsWriteBytes, *, estimate_read_size: int = 0):
+	if estimate_read_size > 1048576 and False:  # TODO
+		__get_copier().copy(src, dst)
+	else:
+		shutil.copyfileobj(src, dst)
+
+
 def rm_rf(path: Path, *, missing_ok: bool = False):
 	"""
 	Does not follow symlink