py_hamt

55def blake3_hashfn(input_bytes: bytes) -> bytes:
56    """
57    This is the default blake3 hash function used for the `HAMT`, with a 32 byte hash size.
58
59    """
60    # 32 bytes is the recommended byte size for blake3 and the default, but multihash forces us to explicitly specify
61    digest: bytes = b3.digest(input_bytes, size=32)
62    raw_bytes: bytes = b3.unwrap(digest)
63    return raw_bytes

303class HAMT:
304    """
305    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.
306
307    Use this to store arbitrarily large key-value mappings in your CAS of choice.
308
309    For writing, this HAMT is async safe but NOT thread safe. Only write in an async event loop within the same thread.
310
311    When in read-only mode, the HAMT is both async and thread safe.
312
313    #### A note about memory management, read+write and read-only modes
314    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.
315
316    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`.
317
318    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.
319
320    #### IPFS HAMT Sample Code
321    ```python
322    kubo_cas = KuboCAS() # connects to a local kubo node with the default endpoints
323    hamt = await HAMT.build(cas=kubo_cas)
324    await hamt.set("foo", "bar")
325    assert (await hamt.get("foo")) == "bar"
326    await hamt.make_read_only()
327    cid = hamt.root_node_id # our root node CID
328    print(cid)
329    ```
330    """
331
332    def __init__(
333        self,
334        cas: ContentAddressedStore,
335        hash_fn: Callable[[bytes], bytes] = blake3_hashfn,
336        root_node_id: IPLDKind | None = None,
337        read_only: bool = False,
338        max_bucket_size: int = 4,
339        values_are_bytes: bool = False,
340    ):
341        """
342        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.
343        """
344
345        self.cas: ContentAddressedStore = cas
346        """The backing storage system. py-hamt provides an implementation `KuboCAS` for IPFS."""
347
348        self.hash_fn: Callable[[bytes], bytes] = hash_fn
349        """
350        This is the hash function used to place a key-value within the HAMT.
351
352        To provide your own hash function, create a function that takes in arbitrarily long bytes and returns the hash bytes.
353
354        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.
355
356        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.
357        """
358
359        self.lock: asyncio.Lock = asyncio.Lock()
360        """@private"""
361
362        self.values_are_bytes: bool = values_are_bytes
363        """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.
364
365        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.
366        """
367
368        if max_bucket_size < 1:
369            raise ValueError("Bucket size maximum must be a positive integer")
370        self.max_bucket_size: int = max_bucket_size
371        """
372        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.
373
374        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.
375
376        This must be a positive integer with a minimum of 1.
377        """
378
379        self.root_node_id: IPLDKind = root_node_id
380        """
381        This is type IPLDKind but the documentation generator pdoc mangles it a bit.
382
383        Read from this only when in read mode to get something valid!
384        """
385
386        self.read_only: bool = read_only
387        """Clients should NOT modify this.
388
389        This is here for checking whether the HAMT is in read only or read/write mode.
390
391        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.
392        """
393        self.node_store: NodeStore
394        """@private"""
395        if read_only:
396            self.node_store = ReadCacheStore(self)
397        else:
398            self.node_store = InMemoryTreeStore(self)
399
400    @classmethod
401    async def build(cls, *args: Any, **kwargs: Any) -> "HAMT":
402        """
403        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__.
404
405        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.
406        """
407        hamt = cls(*args, **kwargs)
408        if hamt.root_node_id is None:
409            hamt.root_node_id = await hamt.node_store.save(None, Node())
410        return hamt
411
412    # 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
413    async def make_read_only(self) -> None:
414        """
415        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.
416
417        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.
418        """
419        async with self.lock:
420            inmemory_tree: InMemoryTreeStore = cast(InMemoryTreeStore, self.node_store)
421            await inmemory_tree.vacate()
422
423            self.read_only = True
424            self.node_store = ReadCacheStore(self)
425
426    async def enable_write(self) -> None:
427        """
428        Enable both reads and writes. This creates an internal structure for performance optimizations which will result in the root node ID no longer being valid, in order to read that at the end of your operations you must first use `make_read_only`.
429        """
430        async with self.lock:
431            # The read cache has no writes that need to be sent upstream so we can remove it without vacating
432            self.read_only = False
433            self.node_store = InMemoryTreeStore(self)
434
435    async def cache_size(self) -> int:
436        """
437        Returns the memory used by some internal performance optimization tools in bytes.
438
439        This is async concurrency safe, so call it whenever. This does mean it will block and wait for other writes to finish however.
440
441        Be warned that this may take a while to run for large HAMTs.
442
443        For more on memory management, see the `HAMT` class documentation.
444        """
445        if self.read_only:
446            return self.node_store.size()
447        async with self.lock:
448            return self.node_store.size()
449
450    async def cache_vacate(self) -> None:
451        """
452        Vacate and completely empty out the internal read/write cache.
453
454        Be warned that this may take a while if there have been a lot of write operations.
455
456        For more on memory management, see the `HAMT` class documentation.
457        """
458        if self.read_only:
459            await self.node_store.vacate()
460        else:
461            async with self.lock:
462                await self.node_store.vacate()
463
464    async def _reserialize_and_link(
465        self, node_stack: list[tuple[IPLDKind, Node]]
466    ) -> None:
467        """
468        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
469        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
470        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
471        If a node ends up being empty, then it is deleted entirely, unless it is the root node
472        Modifies in place
473        """
474        # iterate in the reverse direction, this range goes from n-1 to 0, from the bottommost tree node to the root
475        for stack_index in range(len(node_stack) - 1, -1, -1):
476            old_id, node = node_stack[stack_index]
477
478            # If this node is empty, and it's not the root node, then we can delete it entirely from the list
479            is_root: bool = stack_index == 0
480            if node.is_empty() and not is_root:
481                # Unlink from the rest of the tree
482                _, prev_node = node_stack[stack_index - 1]
483                # When removing links, don't worry about two nodes having the same link since all nodes are guaranteed to be different by the removal of empty nodes after every single operation
484                for link_index in prev_node.iter_link_indices():
485                    link = prev_node.get_link(link_index)
486                    if link == old_id:
487                        # Delete the link by making it an empty bucket
488                        prev_node.data[link_index] = {}
489                        break
490
491                # Remove from our stack, continue reserializing up the tree
492                node_stack.pop(stack_index)
493                continue
494
495            # If not an empty node, just reserialize like normal and replace this one
496            new_store_id: IPLDKind = await self.node_store.save(old_id, node)
497            node_stack[stack_index] = (new_store_id, node)
498
499            # 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
500            if not is_root:
501                _, prev_node = node_stack[stack_index - 1]
502                prev_node.replace_link(old_id, new_store_id)
503
504    # automatically skip encoding if the value provided is of the bytes variety
505    async def set(self, key: str, val: IPLDKind) -> None:
506        """Write a key-value mapping."""
507        if self.read_only:
508            raise Exception("Cannot call set on a read only HAMT")
509
510        data: bytes
511        if self.values_are_bytes:
512            data = cast(
513                bytes, val
514            )  # let users get an exception if they pass in a non bytes when they want to skip encoding
515        else:
516            data = dag_cbor.encode(val)
517
518        pointer: IPLDKind = await self.cas.save(data, codec="raw")
519        await self._set_pointer(key, pointer)
520
521    async def _set_pointer(self, key: str, val_ptr: IPLDKind) -> None:
522        async with self.lock:
523            node_stack: list[tuple[IPLDKind, Node]] = []
524            root_node: Node = await self.node_store.load(self.root_node_id)
525            node_stack.append((self.root_node_id, root_node))
526
527            # FIFO queue to keep track of all the KVs we need to insert
528            # This is needed if any buckets overflow and so we need to reinsert all those KVs
529            kvs_queue: list[tuple[str, IPLDKind]] = []
530            kvs_queue.append((key, val_ptr))
531
532            while len(kvs_queue) > 0:
533                _, top_node = node_stack[-1]
534                curr_key, curr_val_ptr = kvs_queue[0]
535
536                raw_hash: bytes = self.hash_fn(curr_key.encode())
537                map_key: int = extract_bits(raw_hash, len(node_stack) - 1, 8)
538
539                item = top_node.data[map_key]
540                if isinstance(item, list):
541                    next_node_id: IPLDKind = item[0]
542                    next_node: Node = await self.node_store.load(next_node_id)
543                    node_stack.append((next_node_id, next_node))
544                elif isinstance(item, dict):
545                    bucket: dict[str, IPLDKind] = item
546
547                    # If this bucket already has this same key, or has space, just rewrite the value and then go work on the others in the queue
548                    if curr_key in bucket or len(bucket) < self.max_bucket_size:
549                        bucket[curr_key] = curr_val_ptr
550                        kvs_queue.pop(0)
551                        continue
552
553                    # The current key is not in the bucket and the bucket is too full, so empty KVs from the bucket and restart insertion
554                    for k in bucket:
555                        v_ptr = bucket[k]
556                        kvs_queue.append((k, v_ptr))
557
558                    # Create a new link to a new node so that we can reflow these KVs into a new subtree
559                    new_node = Node()
560                    new_node_id: IPLDKind = await self.node_store.save(None, new_node)
561                    link: list[IPLDKind] = [new_node_id]
562                    top_node.data[map_key] = link
563
564            # Finally, reserialize and fix all links, deleting empty nodes as needed
565            await self._reserialize_and_link(node_stack)
566            self.root_node_id = node_stack[0][0]
567
568    async def delete(self, key: str) -> None:
569        """Delete a key-value mapping."""
570
571        # Also deletes the pointer at the same time so this doesn't have a _delete_pointer duo
572        if self.read_only:
573            raise Exception("Cannot call delete on a read only HAMT")
574
575        async with self.lock:
576            raw_hash: bytes = self.hash_fn(key.encode())
577
578            node_stack: list[tuple[IPLDKind, Node]] = []
579            root_node: Node = await self.node_store.load(self.root_node_id)
580            node_stack.append((self.root_node_id, root_node))
581
582            created_change: bool = False
583            while True:
584                _, top_node = node_stack[-1]
585                map_key: int = extract_bits(raw_hash, len(node_stack) - 1, 8)
586
587                item = top_node.data[map_key]
588                if isinstance(item, dict):
589                    bucket = item
590                    if key in bucket:
591                        del bucket[key]
592                        created_change = True
593                    # 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
594                    break
595                elif isinstance(item, list):
596                    link: IPLDKind = item[0]
597                    next_node: Node = await self.node_store.load(link)
598                    node_stack.append((link, next_node))
599
600            # Finally, reserialize and fix all links, deleting empty nodes as needed
601            if created_change:
602                await self._reserialize_and_link(node_stack)
603                self.root_node_id = node_stack[0][0]
604            else:
605                # If we didn't make a change, then this key must not exist within the HAMT
606                raise KeyError
607
608    async def get(
609        self,
610        key: str,
611        offset: Optional[int] = None,
612        length: Optional[int] = None,
613        suffix: Optional[int] = None,
614    ) -> IPLDKind:
615        """Get a value."""
616        pointer: IPLDKind = await self.get_pointer(key)
617        data: bytes = await self.cas.load(
618            pointer, offset=offset, length=length, suffix=suffix
619        )
620        if self.values_are_bytes:
621            return data
622        else:
623            return dag_cbor.decode(data)
624
625    async def get_pointer(self, key: str) -> IPLDKind:
626        """
627        Get a store ID that points to the value for this key.
628
629        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.
630        """
631        # If read only, no need to acquire a lock
632        pointer: IPLDKind
633        if self.read_only:
634            pointer = await self._get_pointer(key)
635        else:
636            async with self.lock:
637                pointer = await self._get_pointer(key)
638
639        return pointer
640
641    # Callers MUST handle acquiring a lock
642    async def _get_pointer(self, key: str) -> IPLDKind:
643        with instrumentation.span(
644            "py_hamt.hamt.lookup", {"py_hamt.hamt.lookup.key": key}
645        ):
646            lookup_started_at = time.perf_counter()
647            raw_hash: bytes = self.hash_fn(key.encode())
648
649            current_id: IPLDKind = self.root_node_id
650            current_depth: int = 0
651            node_loads = 0
652            node_cache_hits = 0
653
654            # 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
655            result_ptr: IPLDKind = None
656            found_a_result: bool = False
657            try:
658                while True:
659                    top_id: IPLDKind = current_id
660                    if (
661                        isinstance(self.node_store, ReadCacheStore)
662                        and top_id in self.node_store.cache
663                    ):
664                        node_cache_hits += 1
665                    node_loads += 1
666                    top_node: Node = await self.node_store.load(top_id)
667                    map_key: int = extract_bits(raw_hash, current_depth, 8)
668
669                    # Check if this key is in one of the buckets
670                    item = top_node.data[map_key]
671                    if isinstance(item, dict):
672                        bucket = item
673                        if key in bucket:
674                            result_ptr = bucket[key]
675                            found_a_result = True
676                            break
677
678                    if isinstance(item, list):
679                        link: IPLDKind = item[0]
680                        current_id = link
681                        current_depth += 1
682                        continue
683
684                    # Nowhere left to go, stop walking down the tree
685                    break
686
687                if not found_a_result:
688                    raise KeyError
689
690                return result_ptr
691            finally:
692                instrumentation.record_hamt_lookup(
693                    key,
694                    depth=current_depth,
695                    node_loads=node_loads,
696                    node_cache_hits=node_cache_hits,
697                    found=found_a_result,
698                    seconds=time.perf_counter() - lookup_started_at,
699                )
700
701    # Callers MUST handle locking or not on their own
702    async def _iter_nodes(self) -> AsyncIterator[tuple[IPLDKind, Node]]:
703        node_id_stack: list[IPLDKind] = [self.root_node_id]
704        while len(node_id_stack) > 0:
705            top_id: IPLDKind = node_id_stack.pop()
706            node: Node = await self.node_store.load(top_id)
707            yield (top_id, node)
708            node_id_stack.extend(list(node.iter_links()))
709
710    async def keys(self) -> AsyncIterator[str]:
711        """
712        AsyncIterator returning all keys in the HAMT.
713
714        If the HAMT is write enabled, to maintain strong consistency this will obtain an async lock and not allow any other operations to proceed.
715
716        When the HAMT is in read only mode however, this can be run concurrently with get operations.
717        """
718        if self.read_only:
719            async for k in self._keys_no_locking():
720                yield k
721        else:
722            async with self.lock:
723                async for k in self._keys_no_locking():
724                    yield k
725
726    async def _keys_no_locking(self) -> AsyncIterator[str]:
727        async for _, node in self._iter_nodes():
728            for bucket in node.iter_buckets():
729                for key in bucket:
730                    yield key
731
732    async def len(self) -> int:
733        """
734        Return the number of key value mappings in this HAMT.
735
736        When the HAMT is write enabled, to maintain strong consistency it will acquire a lock and thus not allow any other operations to proceed until the length is fully done being calculated. If read only, then this can be run concurrently with other operations.
737        """
738        count: int = 0
739        async for _ in self.keys():
740            count += 1
741
742        return count

16class ContentAddressedStore(ABC):
17    """
18    Abstract class that represents a content addressed storage that the `HAMT` can use for keeping data.
19
20    Note that the return type of save and input to load is really type `IPLDKind`, but the documentation generator pdoc mangles it unfortunately.
21
22    #### A note on the IPLDKind return types
23    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:
24    1. No lists or dicts, since python does not classify these as immutable.
25    2. No `None` values since this is used in HAMT's `__init__` to indicate that an empty HAMT needs to be initialized.
26    """
27
28    CodecInput = Literal["raw", "dag-cbor"]
29
30    @abstractmethod
31    async def save(self, data: bytes, codec: CodecInput) -> IPLDKind:
32        """Save data to a storage mechanism, and return an ID for the data in the IPLDKind type.
33
34        `codec` will be set to "dag-cbor" if this data should be marked as special linked data a la IPLD data model.
35        """
36
37    @abstractmethod
38    async def load(
39        self,
40        id: IPLDKind,
41        offset: Optional[int] = None,
42        length: Optional[int] = None,
43        suffix: Optional[int] = None,
44    ) -> bytes:
45        """Retrieve data."""
46
47    async def pin_cid(self, id: IPLDKind, target_rpc: str) -> None:
48        """Pin a CID in the storage."""
49        pass  # pragma: no cover
50
51    async def unpin_cid(self, id: IPLDKind, target_rpc: str) -> None:
52        """Unpin a CID in the storage."""
53        pass  # pragma: no cover
54
55    async def pin_update(
56        self, old_id: IPLDKind, new_id: IPLDKind, target_rpc: str
57    ) -> None:
58        """Update the pinned CID in the storage."""
59        pass  # pragma: no cover
60
61    async def pin_ls(self, target_rpc: str) -> list[Dict[str, Any]]:
62        """List all pinned CIDs in the storage."""
63        return []  # pragma: no cover

 66class InMemoryCAS(ContentAddressedStore):
 67    """Used mostly for faster testing, this is why this is not exported. It hashes all inputs and uses that as a key to an in-memory python dict, mimicking a content addressed storage system. The hash bytes are the ID that `save` returns and `load` takes in."""
 68
 69    store: dict[bytes, bytes]
 70    hash_alg: Multihash
 71
 72    def __init__(self):
 73        self.store = dict()
 74        self.hash_alg = multihash.get("blake3")
 75
 76    async def save(self, data: bytes, codec: ContentAddressedStore.CodecInput) -> bytes:
 77        hash: bytes = self.hash_alg.digest(data, size=32)
 78        self.store[hash] = data
 79        return hash
 80
 81    async def load(
 82        self,
 83        id: IPLDKind,
 84        offset: Optional[int] = None,
 85        length: Optional[int] = None,
 86        suffix: Optional[int] = None,
 87    ) -> bytes:
 88        """
 89        `ContentAddressedStore` allows any IPLD scalar key.  For the in-memory
 90        backend we *require* a `bytes` hash; anything else is rejected at run
 91        time. In OO type-checking, a subclass may widen (make more general) argument types,
 92        but it must never narrow them; otherwise callers that expect the base-class contract can break.
 93        Mypy enforces this contra-variance rule and emits the "violates Liskov substitution principle" error.
 94        This is why we use `cast` here, to tell mypy that we know what we are doing.
 95        h/t https://stackoverflow.com/questions/75209249/overriding-a-method-mypy-throws-an-incompatible-with-super-type-error-when-ch
 96        """
 97        key = cast(bytes, id)
 98        if not isinstance(key, (bytes, bytearray)):  # defensive guard
 99            raise TypeError(
100                f"InMemoryCAS only supports byte‐hash keys; got {type(id).__name__}"
101            )
102        data: bytes
103        try:
104            data = self.store[key]
105        except KeyError as exc:
106            raise KeyError("Object not found in in-memory store") from exc
107
108        if offset is not None:
109            start = offset
110            if length is not None:
111                end = start + length
112                return data[start:end]
113            else:
114                return data[start:]
115        elif suffix is not None:  # If only length is given, assume start from 0
116            return data[-suffix:]
117        else:  # Full load
118            return data