introduce PackFileObjectPool to reduce file open syscalls on reading chunked blobs

Author: Fallen-Breath

prime_backup/action/helpers/blob_exporter.py

@@ -10,6 +10,7 @@
 
 from prime_backup import logger
 from prime_backup.action.helpers.chunk_io import ChunkIO
+from prime_backup.action.helpers.pack_reader import PackFileObjectPool
 from prime_backup.compressors import CompressMethod
 from prime_backup.compressors import Compressor
 from prime_backup.db.session import DbSession
@@ -171,10 +172,10 @@ def __export_to_fs_chunked(self, output_path: Path):
 		blob_chunks = self.blob_chunks_getter.get(self.blob.id)
 		hash_method = chunk_utils.get_hash_method()
 
-		with open(output_path, 'wb') as f_out:
+		with open(output_path, 'wb') as f_out, PackFileObjectPool() as pack_file_obj_pool:
 			for oc in blob_chunks:
 				bypass_reader: Optional[BypassReader] = None
-				with ChunkIO(oc.chunk).open_decompressed() as f_in:
+				with ChunkIO(oc.chunk, pack_file_obj_pool=pack_file_obj_pool).open_decompressed() as f_in:
 					if self.verify_blob:
 						bypass_reader = BypassReader(f_in, calc_hash=True, hash_method=hash_method)
 						file_utils.copy_file_obj_fast(bypass_reader, f_out, estimate_read_size=oc.chunk.raw_size)
@@ -217,18 +218,21 @@ def __export_as_reader_direct(self, reader_csm: Callable[[SupportsReadBytes], An
 	def __export_as_reader_chunked(self, reader_csm: Callable[[SupportsReadBytes], Any]):
 		blob_chunks = self.blob_chunks_getter.get(self.blob.id)
 		exit_flag = False
+		gen_yielded = False
 
 		def open_chunk_gen() -> Generator[_OpenedChunk, None, None]:
 			hash_method = chunk_utils.get_hash_method()
 			for oc in blob_chunks:
-				with ChunkIO(oc.chunk).open_decompressed() as f_in:
+				with ChunkIO(oc.chunk, pack_file_obj_pool=pack_file_obj_pool).open_decompressed() as f_in:
 					try:
 						peek_reader = _PeekReader(f_in, 32 * 1024)
 						peek_reader.peek()
 					except Exception as e:
 						self.logger.error(f'Failed to peek-read chunk {oc.chunk.hash} in pack entry {oc.chunk.pack_entry.pack_id}@{oc.chunk.pack_entry.offset} for {oc}: {e}')
 						raise
 
+					nonlocal gen_yielded
+					gen_yielded = True
 					if self.verify_blob:
 						def verify_callback(ck: ChunkInfo, br: BypassReader):
 							self.__verify_exported_chunk(ck, br.get_read_len(), br.get_hash())
@@ -242,14 +246,16 @@ def verify_callback(ck: ChunkInfo, br: BypassReader):
 				if exit_flag:
 					break
 
-		chunk_gen = open_chunk_gen()
-		try:
-			reader_csm(_CombinedChunksReader(chunk_gen))
-		finally:
-			# ensure possible opened chunk reader is closed
-			exit_flag = True
-			with contextlib.suppress(StopIteration):
-				next(chunk_gen)
+		with PackFileObjectPool() as pack_file_obj_pool:
+			chunk_gen = open_chunk_gen()
+			try:
+				reader_csm(_CombinedChunksReader(chunk_gen))
+			finally:
+				# ensure possible opened ChunkIO is closed
+				exit_flag = True
+				if gen_yielded:
+					with contextlib.suppress(StopIteration):
+						next(chunk_gen)
 
 	def __verify_exported_blob(self, written_size: int, written_hash: str):
 		self.__verify_exported_data(lambda: 'blob', self.blob.raw_size, self.blob.hash, written_size, written_hash)

prime_backup/action/helpers/chunk_io.py

@@ -1,16 +1,17 @@
 import contextlib
-from typing import Generator, Tuple
+from typing import Generator, Tuple, Optional
 
-from prime_backup.action.helpers.pack_reader import PackReader
+from prime_backup.action.helpers.pack_reader import PackReader, PackFileObjectPool
 from prime_backup.compressors import Compressor
 from prime_backup.types.chunk_info import ChunkInfo
 from prime_backup.utils.bypass_io import BypassReader
 from prime_backup.utils.io_types import SupportsReadBytes, SupportsReadAndSeek
 
 
 class ChunkIO:
-	def __init__(self, chunk: ChunkInfo):
+	def __init__(self, chunk: ChunkInfo, *, pack_file_obj_pool: Optional[PackFileObjectPool] = None):
 		self.chunk = chunk
+		self.pack_file_obj_pool = pack_file_obj_pool
 
 	def __get_pack_id(self) -> int:
 		if self.chunk.pack_entry.pack_id <= 0:
@@ -19,7 +20,12 @@ def __get_pack_id(self) -> int:
 
 	@contextlib.contextmanager
 	def open_raw(self) -> Generator[SupportsReadAndSeek, None, None]:
-		with PackReader.open_entry(self.__get_pack_id(), self.chunk.pack_entry.offset, self.chunk.stored_size) as reader:
+		with PackReader.open_entry(
+				pack_id=self.__get_pack_id(),
+				offset=self.chunk.pack_entry.offset,
+				length=self.chunk.stored_size,
+				file_obj_pool=self.pack_file_obj_pool
+		) as reader:
 			yield reader
 
 	@contextlib.contextmanager

prime_backup/action/helpers/pack_reader.py

@@ -1,13 +1,81 @@
+import collections
 import contextlib
 import os
-from typing import BinaryIO, Generator
+import threading
+from typing import BinaryIO, Generator, Optional, List
 
 from typing_extensions import Final
 
 from prime_backup.utils import pack_utils
 from prime_backup.utils.io_types import SupportsReadAndSeek
 
 
+class PackFileObjectPool:
+	def __init__(self, max_size: int = 4):
+		if max_size <= 0:
+			raise ValueError('max_size should be positive, got {}'.format(max_size))
+		self.__max_size = max_size
+		self.__files: 'collections.OrderedDict[int, BinaryIO]' = collections.OrderedDict()
+		self.__lock = threading.Lock()
+		self.__closed = False
+
+	def __enter__(self) -> 'PackFileObjectPool':
+		return self
+
+	def __exit__(self, exc_type, exc_val, exc_tb):
+		self.close()
+
+	@classmethod
+	def __close_files(cls, files: List[BinaryIO]):
+		errors: List[Exception] = []
+		for file in files:
+			try:
+				file.close()
+			except Exception as e:
+				errors.append(e)
+		if len(errors) > 0:
+			raise Exception(f'Failed to close {len(errors)} files: {errors}')
+
+	@contextlib.contextmanager
+	def acquire(self, pack_id: int) -> Generator[BinaryIO, None, None]:
+		with self.__lock:
+			if self.__closed:
+				raise RuntimeError('file object pool is closed')
+			existing_file: Optional[BinaryIO] = self.__files.pop(pack_id, None)
+
+		file: BinaryIO
+		if existing_file is None:
+			file = open(pack_utils.get_pack_path(pack_id), 'rb')
+		else:
+			file = existing_file
+
+		try:
+			yield file
+		finally:
+			to_close_files: List[BinaryIO] = []
+			with self.__lock:
+				if self.__closed:
+					to_close_files.append(file)
+				else:
+					old_file = self.__files.pop(pack_id, None)
+					self.__files[pack_id] = file
+					self.__files.move_to_end(pack_id)
+
+					if old_file is not None and old_file is not file:
+						to_close_files.append(old_file)
+					while len(self.__files) > self.__max_size:
+						_, old_file = self.__files.popitem(last=False)
+						to_close_files.append(old_file)
+			self.__close_files(to_close_files)
+
+	def close(self):
+		with self.__lock:
+			self.__closed = True
+			files = list(self.__files.values())
+			self.__files.clear()
+		self.__close_files(files)
+
+
 class PackEntryReader:
 	def __init__(self, file_obj: BinaryIO, offset: int, length: int):
 		self.__file_obj: Final[BinaryIO] = file_obj
@@ -52,7 +120,12 @@ def seek(self, offset: int, whence: int = 0):
 class PackReader:
 	@classmethod
 	@contextlib.contextmanager
-	def open_entry(cls, pack_id: int, offset: int, length: int) -> Generator[SupportsReadAndSeek, None, None]:
-		pack_path = pack_utils.get_pack_path(pack_id)
-		with open(pack_path, 'rb') as fh:
-			yield PackEntryReader(fh, offset, length)
+	def open_entry(cls, pack_id: int, offset: int, length: int, *, file_obj_pool: Optional[PackFileObjectPool] = None) -> Generator[SupportsReadAndSeek, None, None]:
+		with contextlib.ExitStack() as es:
+			file: BinaryIO
+			if file_obj_pool is None:
+				pack_path = pack_utils.get_pack_path(pack_id)
+				file = es.enter_context(open(pack_path, 'rb'))
+			else:
+				file = es.enter_context(file_obj_pool.acquire(pack_id))
+			yield PackEntryReader(file, offset, length)