py_hamt

54def blake3_hashfn(input_bytes: bytes) -> bytes:
55    """Return the HAMT's default raw 32-byte BLAKE3 digest."""
56    return blake3(input_bytes).digest(length=32)

361class HAMT:
362    """
363    An implementation of a Hash Array Mapped Trie for an arbitrary Content Addressed Storage (CAS) system, e.g. IPFS. This uses the IPLD data model.
364
365    Use this to store arbitrarily large key-value mappings in your CAS of choice.
366
367    For writing, this HAMT is async safe but NOT thread safe. Only write in an async event loop within the same thread.
368
369    When in read-only mode, the HAMT is both async and thread safe.
370
371    #### A note about memory management, read+write and read-only modes
372    The HAMT can be in either read+write mode or read-only mode. For either of these modes, the HAMT has some internal performance optimizations.
373
374    Note that in read+write, the real root node id IS NOT VALID. You should call `make_read_only()` to convert to read only mode and then read `root_node_id`.
375
376    These optimizations also trade off performance for memory use. Use `cache_size` to monitor the approximate memory usage. Be warned that for large key-value mapping sets this may take a bit to run. Use `cache_vacate` if you are over your memory limits.
377
378    #### IPFS HAMT Sample Code
379    ```python
380    kubo_cas = KuboCAS() # connects to a local kubo node with the default endpoints
381    hamt = await HAMT.build(cas=kubo_cas)
382    await hamt.set("foo", "bar")
383    assert (await hamt.get("foo")) == "bar"
384    await hamt.make_read_only()
385    cid = hamt.root_node_id # our root node CID
386    print(cid)
387    ```
388    """
389
390    def __init__(
391        self,
392        cas: ContentAddressedStore,
393        hash_fn: Callable[[bytes], bytes] = blake3_hashfn,
394        root_node_id: IPLDKind | None = None,
395        read_only: bool = False,
396        max_bucket_size: int = 4,
397        values_are_bytes: bool = False,
398    ):
399        """
400        Use `build` if you need to create a completely empty HAMT, as this requires some async operations with the CAS. For what each of the constructor input variables refer to, check the documentation with the matching names below.
401        """
402
403        self.cas: ContentAddressedStore = cas
404        """The backing storage system. py-hamt provides an implementation `KuboCAS` for IPFS."""
405
406        self.hash_fn: Callable[[bytes], bytes] = hash_fn
407        """
408        This is the hash function used to place a key-value within the HAMT.
409
410        To provide your own hash function, create a function that takes in arbitrarily long bytes and returns the hash bytes.
411
412        It's important to note that the resulting hash must must always be a multiple of 8 bits since python bytes object can only represent in segments of bytes, and thus 8 bits.
413
414        Theoretically your hash size must only be a minimum of 1 byte, and there can be less than or the same number of hash collisions as the bucket size. Any more and the HAMT will most likely throw errors.
415        """
416
417        self.lock: asyncio.Lock = asyncio.Lock()
418        """@private"""
419
420        self.values_are_bytes: bool = values_are_bytes
421        """Set this to true if you are only going to be storing python bytes objects into the hamt. This will improve performance by skipping a serialization step from IPLDKind.
422
423        This is theoretically safe to change in between operations, but this has not been verified in testing, so only do this at your own risk.
424        """
425
426        if max_bucket_size < 1:
427            raise ValueError("Bucket size maximum must be a positive integer")
428        self.max_bucket_size: int = max_bucket_size
429        """
430        This is only important for tuning performance when writing! For reading a HAMT that was written with a different max bucket size, this does not need to match and can be left unprovided.
431
432        This is an internal detail that has been exposed for performance tuning. The HAMT handles large key-value mapping sets even on a content addressed system by essentially sharding all the mappings across many smaller Nodes. The memory footprint of each of these Nodes footprint is a linear function of the maximum bucket size. Larger bucket sizes will result in larger Nodes, but more time taken to retrieve and decode these nodes from your backing CAS.
433
434        This must be a positive integer with a minimum of 1.
435        """
436
437        self.root_node_id: IPLDKind = root_node_id
438        """
439        This is type IPLDKind but the documentation generator pdoc mangles it a bit.
440
441        Read from this only when in read mode to get something valid!
442        """
443
444        self.read_only: bool = read_only
445        """Clients should NOT modify this.
446
447        This is here for checking whether the HAMT is in read only or read/write mode.
448
449        The distinction is made for performance and correctness reasons. In read only mode, the HAMT has an internal read cache that can speed up operations. In read/write mode, for reads the HAMT maintains strong consistency for reads by using async locks, and for writes the HAMT writes to an in memory buffer rather than performing (possibly) network calls to the underlying CAS.
450        """
451        self.node_store: NodeStore
452        """@private"""
453        if read_only:
454            self.node_store = ReadCacheStore(self)
455        else:
456            self.node_store = InMemoryTreeStore(self)
457
458    @classmethod
459    async def build(cls, *args: Any, **kwargs: Any) -> "HAMT":
460        """
461        Use this if you are initializing a completely empty HAMT! That means passing in None for the root_node_id. Method arguments are the exact same as `__init__`. If the root_node_id is not None, this will have no difference than creating a HAMT instance with __init__.
462
463        This separate async method is required since initializing an empty HAMT means sending some internal objects to the underlying CAS, which requires async operations. python does not allow for an async __init__, so this method is separately provided.
464        """
465        hamt = cls(*args, **kwargs)
466        if hamt.root_node_id is None:
467            hamt.root_node_id = await hamt.node_store.save(None, Node())
468        return hamt
469
470    # This is typically a massive blocking operation, you dont want to be running this concurrently with a bunch of other operations, so it's ok to have it not be async
471    async def make_read_only(self) -> None:
472        """
473        Makes the HAMT read only, which allows for more parallel read operations. The HAMT also needs to be in read only mode to get the real root node ID.
474
475        In read+write mode, the HAMT normally has to block separate get calls to enable strong consistency in case a set/delete operation falls in between.
476        """
477        async with self.lock:
478            inmemory_tree: InMemoryTreeStore = cast(InMemoryTreeStore, self.node_store)
479            await inmemory_tree.vacate()
480
481            self.read_only = True
482            self.node_store = ReadCacheStore(self)
483
484    async def enable_write(self) -> None:
485        """
486        Enable both reads and writes. Calling this while writes are already enabled is a no-op that preserves any buffered changes. The read-only to writable transition creates an internal structure for performance optimizations which will result in the root node ID no longer being valid; to read it at the end of your operations, first use `make_read_only`.
487        """
488        async with self.lock:
489            if not self.read_only:
490                return
491
492            # The read cache has no writes that need to be sent upstream so we can remove it without vacating
493            self.read_only = False
494            self.node_store = InMemoryTreeStore(self)
495
496    async def cache_size(self) -> int:
497        """
498        Returns the memory used by some internal performance optimization tools in bytes.
499
500        This is async concurrency safe, so call it whenever. This does mean it will block and wait for other writes to finish however.
501
502        Be warned that this may take a while to run for large HAMTs.
503
504        For more on memory management, see the `HAMT` class documentation.
505        """
506        if self.read_only:
507            return self.node_store.size()
508        async with self.lock:
509            return self.node_store.size()
510
511    async def cache_vacate(self) -> None:
512        """
513        Vacate and completely empty out the internal read/write cache.
514
515        Be warned that this may take a while if there have been a lot of write operations.
516
517        For more on memory management, see the `HAMT` class documentation.
518        """
519        if self.read_only:
520            await self.node_store.vacate()
521        else:
522            async with self.lock:
523                await self.node_store.vacate()
524
525    async def _reserialize_and_link(
526        self,
527        node_stack: list[tuple[IPLDKind, Node]],
528        link_path: list[int],
529    ) -> None:
530        """
531        This function starts from the node at the end of the list and reserializes so that each node holds valid new IDs after insertion into the store
532        Takes a stack of nodes, we represent a stack with a list where the first element is the root element and the last element is the top of the stack
533        Each element in the list is a tuple where the first element is the ID from the store and the second element is the Node in python
534        `link_path[i]` is the index in `node_stack[i - 1]` through which
535        `node_stack[i]` is linked. The root entry at index zero is a sentinel.
536        If a node ends up being empty, then it is deleted entirely, unless it is the root node
537        Modifies in place
538        """
539        # iterate in the reverse direction, this range goes from n-1 to 0, from the bottommost tree node to the root
540        for stack_index in range(len(node_stack) - 1, -1, -1):
541            old_id, node = node_stack[stack_index]
542
543            # If this node is empty, and it's not the root node, then we can delete it entirely from the list
544            is_root: bool = stack_index == 0
545            if node.is_empty() and not is_root:
546                # Unlink from the rest of the tree using the recorded parent slot.
547                _, prev_node = node_stack[stack_index - 1]
548                prev_node.data[link_path[stack_index]] = {}
549
550                # Remove from our stack, continue reserializing up the tree
551                node_stack.pop(stack_index)
552                link_path.pop(stack_index)
553                continue
554
555            # If not an empty node, just reserialize like normal and replace this one
556            new_store_id: IPLDKind = await self.node_store.save(old_id, node)
557            node_stack[stack_index] = (new_store_id, node)
558
559            # If this is not the last i.e. root node, we need to change the linking of the node prior in the list since we just reserialized
560            if not is_root:
561                _, prev_node = node_stack[stack_index - 1]
562                prev_node.set_link(link_path[stack_index], new_store_id)
563
564    async def _collect_subtree_entries(
565        self, node: Node, limit: int
566    ) -> dict[str, IPLDKind] | None:
567        """Collect a subtree's entries when they fit in a single bucket."""
568        entries: dict[str, IPLDKind] = {}
569
570        for bucket in node.iter_buckets():
571            if len(entries) + len(bucket) > limit:
572                return None
573            entries.update(bucket)
574
575        for link in node.iter_links():
576            child = await self.node_store.load(link)
577            child_entries = await self._collect_subtree_entries(
578                child, limit - len(entries)
579            )
580            if child_entries is None:
581                return None
582            entries.update(child_entries)
583
584        return entries
585
586    async def _collapse_delete_path(
587        self,
588        node_stack: list[tuple[IPLDKind, Node]],
589        link_path: list[int],
590    ) -> None:
591        """Collapse small subtrees into parent buckets after a deletion.
592
593        Applying this bottom-up restores the same shape produced by a fresh build:
594        every linked subtree contains more entries than ``max_bucket_size``. This
595        can change CIDs only for trees that previously retained a non-canonical
596        post-delete shape.
597        """
598        node_store = cast(InMemoryTreeStore, self.node_store)
599        for stack_index in range(len(node_stack) - 1, 0, -1):
600            old_id, node = node_stack[stack_index]
601            entries = await self._collect_subtree_entries(node, self.max_bucket_size)
602            if entries is None:
603                continue
604
605            _, parent = node_stack[stack_index - 1]
606            parent.data[link_path[stack_index]] = entries
607            node_store.remove_clean_node(old_id)
608            node_stack.pop(stack_index)
609            link_path.pop(stack_index)
610
611    # automatically skip encoding if the value provided is of the bytes variety
612    async def set(self, key: str, val: IPLDKind) -> None:
613        """Write a key-value mapping."""
614        if self.read_only:
615            raise Exception("Cannot call set on a read only HAMT")
616
617        data: bytes
618        if self.values_are_bytes:
619            data = cast(
620                bytes, val
621            )  # let users get an exception if they pass in a non bytes when they want to skip encoding
622        else:
623            data = dag_cbor.encode(val)
624
625        pointer: IPLDKind = await self.cas.save(data, codec="raw")
626        await self._set_pointer(key, pointer)
627
628    async def _build_overflow_subtree(
629        self,
630        kvs_queue: list[tuple[str, IPLDKind, bytes]],
631        depth: int,
632        key_hashes: dict[str, bytes],
633    ) -> IPLDKind:
634        """Build a detached subtree for a full bucket.
635
636        Every node touched here is new to the current write. Hash exhaustion can
637        therefore fail without exposing a partial mutation in the existing tree.
638        """
639        root_node = Node()
640        root_node_id = await self.node_store.save(None, root_node)
641        node_stack: list[tuple[IPLDKind, Node]] = [(root_node_id, root_node)]
642
643        while kvs_queue:
644            _, top_node = node_stack[-1]
645            curr_key, curr_val_ptr, raw_hash = kvs_queue[0]
646            map_key = extract_bits(raw_hash, depth + len(node_stack) - 1, 8)
647
648            item = top_node.data[map_key]
649            if isinstance(item, list):
650                next_node_id = item[0]
651                next_node = await self.node_store.load(next_node_id)
652                node_stack.append((next_node_id, next_node))
653                continue
654
655            bucket = cast(dict[str, IPLDKind], item)
656            if curr_key in bucket or len(bucket) < self.max_bucket_size:
657                bucket[curr_key] = curr_val_ptr
658                kvs_queue.pop(0)
659                continue
660
661            for bucket_key, bucket_value in bucket.items():
662                kvs_queue.append((bucket_key, bucket_value, key_hashes[bucket_key]))
663
664            new_node = Node()
665            new_node_id = await self.node_store.save(None, new_node)
666            top_node.set_link(map_key, new_node_id)
667
668        return root_node_id
669
670    async def _set_pointer(self, key: str, val_ptr: IPLDKind) -> None:
671        """Set a value pointer without exposing partial tree mutations on failure."""
672        async with self.lock:
673            node_store = cast(InMemoryTreeStore, self.node_store)
674            node_stack: list[tuple[IPLDKind, Node]] = []
675            link_path: list[int] = [-1]
676            root_node: Node = await self.node_store.load(self.root_node_id)
677            node_stack.append((self.root_node_id, root_node))
678
679            raw_hash = self.hash_fn(key.encode())
680            while True:
681                _, top_node = node_stack[-1]
682                map_key = extract_bits(raw_hash, len(node_stack) - 1, 8)
683
684                item = top_node.data[map_key]
685                if isinstance(item, list):
686                    next_node_id = item[0]
687                    next_node = await self.node_store.load(next_node_id)
688                    node_stack.append((next_node_id, next_node))
689                    link_path.append(map_key)
690                    continue
691
692                bucket = cast(dict[str, IPLDKind], item)
693                if key in bucket or len(bucket) < self.max_bucket_size:
694                    bucket[key] = val_ptr
695                    break
696
697                key_hashes = {key: raw_hash}
698                kvs_queue = [(key, val_ptr, raw_hash)]
699                for bucket_key, bucket_value in bucket.items():
700                    key_hashes[bucket_key] = self.hash_fn(bucket_key.encode())
701                    kvs_queue.append((bucket_key, bucket_value, key_hashes[bucket_key]))
702
703                original_buffer_ids = set(node_store.buffer)
704                try:
705                    new_node_id = await self._build_overflow_subtree(
706                        kvs_queue, len(node_stack), key_hashes
707                    )
708                except BaseException:
709                    for buffer_id in set(node_store.buffer) - original_buffer_ids:
710                        del node_store.buffer[buffer_id]
711                    raise
712
713                top_node.set_link(map_key, new_node_id)
714                break
715
716            # Finally, reserialize and fix all links.
717            await self._reserialize_and_link(node_stack, link_path)
718            self.root_node_id = node_stack[0][0]
719
720    async def delete(self, key: str) -> None:
721        """Delete a key-value mapping.
722
723        Failure-atomic with respect to storage errors: all fallible CAS loads
724        happen before any node is mutated, so if deletion raises, the
725        observable tree remains unchanged.
726        """
727
728        # Also deletes the pointer at the same time so this doesn't have a _delete_pointer duo
729        if self.read_only:
730            raise Exception("Cannot call delete on a read only HAMT")
731
732        async with self.lock:
733            raw_hash: bytes = self.hash_fn(key.encode())
734
735            node_stack: list[tuple[IPLDKind, Node]] = []
736            link_path: list[int] = [-1]
737            root_node: Node = await self.node_store.load(self.root_node_id)
738            node_stack.append((self.root_node_id, root_node))
739
740            created_change: bool = False
741            while True:
742                _, top_node = node_stack[-1]
743                map_key: int = extract_bits(raw_hash, len(node_stack) - 1, 8)
744
745                item = top_node.data[map_key]
746                if isinstance(item, dict):
747                    bucket = item
748                    if key in bucket:
749                        # Collapse may inspect sibling subtrees after the bucket is
750                        # changed. Load them first so a CAS failure cannot leave a
751                        # shared cached node partially mutated. The pre-delete tree
752                        # has one extra entry along this path, hence the +1 budget.
753                        for _, path_node in node_stack[1:]:
754                            await self._collect_subtree_entries(
755                                path_node, self.max_bucket_size + 1
756                            )
757                        del bucket[key]
758                        created_change = True
759                    # Break out since whether or not the key is in the bucket, it should have been here so either now reserialize or raise a KeyError
760                    break
761                elif isinstance(item, list):
762                    link: IPLDKind = item[0]
763                    next_node: Node = await self.node_store.load(link)
764                    node_stack.append((link, next_node))
765                    link_path.append(map_key)
766
767            # Finally, restore the canonical shape and fix all remaining links.
768            if created_change:
769                await self._collapse_delete_path(node_stack, link_path)
770                await self._reserialize_and_link(node_stack, link_path)
771                self.root_node_id = node_stack[0][0]
772            else:
773                # If we didn't make a change, then this key must not exist within the HAMT
774                raise KeyError
775
776    async def get(
777        self,
778        key: str,
779        offset: Optional[int] = None,
780        length: Optional[int] = None,
781        suffix: Optional[int] = None,
782    ) -> IPLDKind:
783        """Get a value."""
784        pointer: IPLDKind = await self.get_pointer(key)
785        data: bytes = await self.cas.load(
786            pointer, offset=offset, length=length, suffix=suffix
787        )
788        if self.values_are_bytes:
789            return data
790        else:
791            return dag_cbor.decode(data)
792
793    async def get_pointer(self, key: str) -> IPLDKind:
794        """
795        Get a store ID that points to the value for this key.
796
797        This is useful for some applications that want to implement a read cache. Due to the restrictions of `ContentAddressedStore` on IDs, pointers are regarded as immutable by python so they can be easily used as IDs for read caches. This is utilized in `ZarrHAMTStore` for example.
798        """
799        # If read only, no need to acquire a lock
800        pointer: IPLDKind
801        if self.read_only:
802            pointer = await self._get_pointer(key)
803        else:
804            async with self.lock:
805                pointer = await self._get_pointer(key)
806
807        return pointer
808
809    # Callers MUST handle acquiring a lock
810    async def _get_pointer(self, key: str) -> IPLDKind:
811        with instrumentation.span(
812            "py_hamt.hamt.lookup", {"py_hamt.hamt.lookup.key": key}
813        ):
814            lookup_started_at = time.perf_counter()
815            raw_hash: bytes = self.hash_fn(key.encode())
816
817            current_id: IPLDKind = self.root_node_id
818            current_depth: int = 0
819            node_loads = 0
820            node_cache_hits = 0
821
822            # Don't check if result is none but use a boolean to indicate finding something, this is because None is a possible value of IPLDKind
823            result_ptr: IPLDKind = None
824            found_a_result: bool = False
825            try:
826                while True:
827                    top_id: IPLDKind = current_id
828                    if (
829                        isinstance(self.node_store, ReadCacheStore)
830                        and top_id in self.node_store.cache
831                    ):
832                        node_cache_hits += 1
833                    node_loads += 1
834                    top_node: Node = await self.node_store.load(top_id)
835                    map_key: int = extract_bits(raw_hash, current_depth, 8)
836
837                    # Check if this key is in one of the buckets
838                    item = top_node.data[map_key]
839                    if isinstance(item, dict):
840                        bucket = item
841                        if key in bucket:
842                            result_ptr = bucket[key]
843                            found_a_result = True
844                            break
845
846                    if isinstance(item, list):
847                        link: IPLDKind = item[0]
848                        current_id = link
849                        current_depth += 1
850                        continue
851
852                    # Nowhere left to go, stop walking down the tree
853                    break
854
855                if not found_a_result:
856                    raise KeyError
857
858                return result_ptr
859            finally:
860                instrumentation.record_hamt_lookup(
861                    key,
862                    depth=current_depth,
863                    node_loads=node_loads,
864                    node_cache_hits=node_cache_hits,
865                    found=found_a_result,
866                    seconds=time.perf_counter() - lookup_started_at,
867                )
868
869    # Callers MUST handle locking or not on their own
870    async def _iter_nodes(self) -> AsyncIterator[tuple[IPLDKind, Node]]:
871        node_id_stack: list[IPLDKind] = [self.root_node_id]
872        while len(node_id_stack) > 0:
873            top_id: IPLDKind = node_id_stack.pop()
874            node: Node = await self.node_store.load(top_id)
875            yield (top_id, node)
876            node_id_stack.extend(list(node.iter_links()))
877
878    async def keys(self) -> AsyncIterator[str]:
879        """
880        AsyncIterator returning all keys in the HAMT.
881
882        If the HAMT is write enabled, the keys present when iteration starts are
883        copied while holding the async lock. The lock is released before any key
884        is yielded, so reads and mutations are safe between iterations and do not
885        affect the keys returned by an iteration already in progress.
886
887        When the HAMT is in read only mode however, this can be run concurrently with get operations.
888        """
889        if self.read_only:
890            async for k in self._keys_no_locking():
891                yield k
892        else:
893            # Buffered nodes are mutable, so copying only the root ID would not
894            # isolate iteration from mutations made after a caller-visible yield.
895            async with self.lock:
896                keys_snapshot = [key async for key in self._keys_no_locking()]
897            for key in keys_snapshot:
898                yield key
899
900    async def _keys_no_locking(self) -> AsyncIterator[str]:
901        async for _, node in self._iter_nodes():
902            for bucket in node.iter_buckets():
903                for key in bucket:
904                    yield key
905
906    async def len(self) -> int:
907        """
908        Return the number of key value mappings in this HAMT.
909
910        When the HAMT is write enabled, keys are counted directly while holding
911        the async lock, without materializing a snapshot. If read only, counting
912        can run concurrently with other operations.
913        """
914        count: int = 0
915        if self.read_only:
916            async for _ in self._keys_no_locking():
917                count += 1
918        else:
919            async with self.lock:
920                async for _ in self._keys_no_locking():
921                    count += 1
922
923        return count

395class ContentAddressedStore(ABC):
396    """
397    Abstract class that represents a content addressed storage that the `HAMT` can use for keeping data.
398
399    Note that the return type of save and input to load is really type `IPLDKind`, but the documentation generator pdoc mangles it unfortunately.
400
401    #### A note on the IPLDKind return types
402    Save and load return the type IPLDKind and not just a CID. As long as python regards the underlying type as immutable it can be used, allowing for more flexibility. There are two exceptions:
403    1. No lists or dicts, since python does not classify these as immutable.
404    2. No `None` values since this is used in HAMT's `__init__` to indicate that an empty HAMT needs to be initialized.
405    """
406
407    CodecInput = Literal["raw", "dag-cbor"]
408
409    @abstractmethod
410    async def save(self, data: bytes, codec: CodecInput) -> IPLDKind:
411        """Save data to a storage mechanism, and return an ID for the data in the IPLDKind type.
412
413        `codec` will be set to "dag-cbor" if this data should be marked as special linked data a la IPLD data model.
414        """
415
416    @abstractmethod
417    async def load(
418        self,
419        id: IPLDKind,
420        offset: Optional[int] = None,
421        length: Optional[int] = None,
422        suffix: Optional[int] = None,
423    ) -> bytes:
424        """Retrieve data."""
425
426    async def pin_cid(self, id: IPLDKind, target_rpc: str) -> None:
427        """Pin a CID in the storage."""
428        pass  # pragma: no cover
429
430    async def unpin_cid(self, id: IPLDKind, target_rpc: str) -> None:
431        """Unpin a CID in the storage."""
432        pass  # pragma: no cover
433
434    async def pin_update(
435        self, old_id: IPLDKind, new_id: IPLDKind, target_rpc: str
436    ) -> None:
437        """Update the pinned CID in the storage."""
438        pass  # pragma: no cover
439
440    async def pin_ls(self, target_rpc: str) -> list[Dict[str, Any]]:
441        """List all pinned CIDs in the storage."""
442        return []  # pragma: no cover