OpportunisticStreamingCarBlockStore

Author: DavidBuchanan314

src/atmst/blockstore/car_file.py

@@ -105,6 +105,77 @@ def get_block(self, key: bytes) -> bytes:
 	def del_block(self, key: bytes) -> None:
 		raise NotImplementedError("ReadOnlyCARBlockStore does not support delete()")
 
+class OptimisticRetryError(KeyError):
+	"""Raised when an optimistic streaming read fails and a retry with
+	optimistic=False may succeed."""
+	pass
+
+class OpportunisticStreamingCarBlockStore(BlockStore):
+	"""
+	A BlockStore backed by a CarStreamReader. Optimistically reads blocks
+	sequentially from the stream, assuming preorder traversal order. If a
+	block is received out of order, falls back to slurping the rest of
+	the stream into memory.
+
+	This should work in 99.999% of cases. The fallback case could fail if the CAR
+	starts off in canonical order, but has duplicate record CIDs where the second
+	record is *not* duplicated into the CAR.
+
+	In such a case, an OptimisticRetryError is raised, and the caller should
+	retry from scratch with optimistic=False.
+	"""
+
+	car_root: CID
+
+	def __init__(self, carstream: "CarStreamReader", optimistic: bool=True) -> None:
+		self.car_root = carstream.car_root
+		self._car_iter = iter(carstream)
+		self._blocks: Dict[bytes, bytes] = {}
+		self._optimistic = optimistic
+		self._was_optimistic = optimistic
+		if not optimistic:
+			self._slurp(self._car_iter)
+
+	def _slurp(self, car_iter) -> None:
+		for k, v in car_iter:
+			self._blocks[bytes(k)] = v
+
+	def get_block(self, key: bytes) -> bytes:
+		if self._optimistic:
+			try:
+				k, v = next(self._car_iter)
+			except StopIteration:
+				raise KeyError(f"block not found: {key!r}")
+			if bytes(k) == key:
+				return v
+			# CAR is not canonically ordered, slurp the rest into memory
+			self._optimistic = False
+			self._blocks[bytes(k)] = v
+			self._slurp(self._car_iter)
+		try:
+			return self._blocks[key]
+		except KeyError:
+			if self._was_optimistic:
+				raise OptimisticRetryError(key)
+			raise
+
+	def is_canonical(self) -> bool:
+		if not self._optimistic:
+			return False
+		# check that the stream has no remaining blocks
+		try:
+			next(self._car_iter)
+			return False
+		except StopIteration:
+			return True
+
+	def put_block(self, key: bytes, value: bytes) -> None:
+		raise NotImplementedError
+
+	def del_block(self, key: bytes) -> None:
+		raise NotImplementedError
+
+
 class CarStreamReader:
 	"""
 	Rather than pre-indexing the block offsets, this lets you iterate over the k/v pairs

src/atmst/cartool.py

@@ -6,12 +6,11 @@
 
 from cbrrr import encode_dag_cbor, decode_dag_cbor, CID
 
-from .blockstore.car_file import ReadOnlyCARBlockStore, encode_varint, CarStreamReader
+from .blockstore.car_file import ReadOnlyCARBlockStore, encode_varint, CarStreamReader, OpportunisticStreamingCarBlockStore, OptimisticRetryError
 from .blockstore import OverlayBlockStore
 from .mst.node_store import NodeStore
 from .mst.node_walker import NodeWalker
 from .mst.diff import mst_diff, record_diff
-from .mst.node import MSTNode
 
 
 def prettify_record(record) -> str:
@@ -110,48 +109,28 @@ def print_record_diff(car_a: str, car_b: str):
 	for delta in record_diff(ns, mst_created, mst_deleted):
 		print(delta)
 
-def verify_car_streaming(carstream: CarStreamReader):
-	blocks = {} # for a preorder-traversal-ordered CAR, this never grows beyond 0
-	optimistic = [True]
-	car_iter = iter(carstream)
-	def lazy_get(key: CID) -> bytes:
-		print("len", len(blocks))
-		if optimistic[0]:
-			try:
-				k, v = next(car_iter)
-			except StopIteration:
-				raise ValueError(f"lookup failed for {key}")
-			if k == key:
-				return v
-			# if we reached here the CAR is not canonically ordered
-			optimistic[0] = False
-			blocks[k] = v
-			for k, v in car_iter: # slurp the entire rest of CAR into RAM
-				blocks[k] = v
-			# fall thru
-		return blocks[key] # TODO: reopen input and re-slurp if this fails
-	commit = decode_dag_cbor(lazy_get(carstream.car_root))
+def verify_car_streaming(carstream: CarStreamReader, optimistic: bool=True):
+	bs = OpportunisticStreamingCarBlockStore(carstream, optimistic=optimistic)
+	commit = decode_dag_cbor(bs.get_block(bytes(bs.car_root)))
 	assert isinstance(commit, dict)
-	root_cid = commit["data"]
-	assert isinstance(root_cid, CID)
-	def verify_mst(node_cid: CID):
-		node = MSTNode.deserialise(lazy_get(node_cid))
-		if node.subtrees[0] is not None:
-			verify_mst(node.subtrees[0])
-		for k, v, subtree in zip(node.keys, node.vals, node.subtrees[1:]):
-			print(k)
-			rv = lazy_get(v)
-			print(k, len(rv))
-			if subtree is not None:
-				verify_mst(subtree)
-
-	verify_mst(root_cid)
-	print(carstream.file.tell()) # should be at EOF now
+	root = commit["data"]
+	assert isinstance(root, CID)
+	ns = NodeStore(bs)
+	count = 0
+	for k, v in NodeWalker(ns, root).iter_kv():
+		bs.get_block(bytes(v)) # force-read every record block
+		count += 1
+	print(f"Verified {count} records (assuming commit signature is valid)")
+	print("canonical:", bs.is_canonical())
 
 def verify_car(car_path: str):
-	with open(car_path, "rb") as carfile:
-		carstream = CarStreamReader(carfile)
-		verify_car_streaming(carstream)
+	try:
+		with open(car_path, "rb") as carfile:
+			verify_car_streaming(CarStreamReader(carfile))
+	except OptimisticRetryError:
+		print("Optimistic streaming failed, retrying with full buffering...")
+		with open(car_path, "rb") as carfile:
+			verify_car_streaming(CarStreamReader(carfile), optimistic=False)
 
 COMMANDS = {
 	"info": (print_info, "print CAR header and repo info"),