|
2 | 2 |
|
3 | 3 | import asyncio |
4 | 4 | import pickle |
| 5 | +import time |
5 | 6 | from abc import abstractmethod |
6 | 7 | from typing import TYPE_CHECKING, Self |
7 | 8 |
|
|
10 | 11 | from zarr.storage import WrapperStore |
11 | 12 |
|
12 | 13 | if TYPE_CHECKING: |
| 14 | + from collections.abc import AsyncIterator, Iterable, Sequence |
13 | 15 | from typing import Any |
14 | 16 |
|
15 | 17 | from zarr.core.buffer.core import BufferPrototype |
@@ -653,7 +655,10 @@ def set_latency(self) -> float: |
653 | 655 | return max(0.0, np.random.normal(loc=self._set_latency[0], scale=self._set_latency[1])) |
654 | 656 |
|
655 | 657 | def _with_store(self, store: Store) -> Self: |
656 | | - return type(self)(store, get_latency=self.get_latency, set_latency=self.set_latency) |
| 658 | + # Pass the raw latency config, not the sampled `get_latency`/`set_latency` |
| 659 | + # properties — sampling would freeze a `(loc, scale)` distribution into |
| 660 | + # one fixed float on derived stores (e.g. via `with_read_only`). |
| 661 | + return type(self)(store, get_latency=self._get_latency, set_latency=self._set_latency) |
657 | 662 |
|
658 | 663 | async def set(self, key: str, value: Buffer) -> None: |
659 | 664 | """ |
@@ -698,3 +703,76 @@ async def get( |
698 | 703 | """ |
699 | 704 | await asyncio.sleep(self.get_latency) |
700 | 705 | return await self._store.get(key, prototype=prototype, byte_range=byte_range) |
| 706 | + |
| 707 | + def get_sync( |
| 708 | + self, |
| 709 | + key: str, |
| 710 | + *, |
| 711 | + prototype: BufferPrototype | None = None, |
| 712 | + byte_range: ByteRequest | None = None, |
| 713 | + ) -> Buffer | None: |
| 714 | + """Add latency to `get_sync`. |
| 715 | +
|
| 716 | + Sleeps `self.get_latency` on the calling thread (the sync path runs on |
| 717 | + worker threads, not the event loop) before delegating to the wrapped |
| 718 | + store. |
| 719 | + """ |
| 720 | + time.sleep(self.get_latency) |
| 721 | + return super().get_sync(key, prototype=prototype, byte_range=byte_range) |
| 722 | + |
| 723 | + def set_sync(self, key: str, value: Buffer) -> None: |
| 724 | + """Add latency to `set_sync`. |
| 725 | +
|
| 726 | + Sleeps `self.set_latency` on the calling thread (the sync path runs on |
| 727 | + worker threads, not the event loop) before delegating to the wrapped |
| 728 | + store. |
| 729 | + """ |
| 730 | + time.sleep(self.set_latency) |
| 731 | + super().set_sync(key, value) |
| 732 | + |
| 733 | + async def get_ranges( |
| 734 | + self, |
| 735 | + key: str, |
| 736 | + byte_ranges: Sequence[ByteRequest | None], |
| 737 | + *, |
| 738 | + prototype: BufferPrototype, |
| 739 | + max_concurrency: int | None = None, |
| 740 | + max_gap_bytes: int | None = None, |
| 741 | + max_coalesced_bytes: int | None = None, |
| 742 | + ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]: |
| 743 | + """Byte-range reads built on `self.get`, so each fetch pays latency. |
| 744 | +
|
| 745 | + Routes through the coalescing `Store.get_ranges` default instead of the |
| 746 | + `WrapperStore` delegation, which would bypass this wrapper's `get` and |
| 747 | + therefore the synthetic latency. `None` for a coalescing kwarg means |
| 748 | + "use the `Store` default". |
| 749 | + """ |
| 750 | + kwargs: dict[str, int] = {} |
| 751 | + if max_concurrency is not None: |
| 752 | + kwargs["max_concurrency"] = max_concurrency |
| 753 | + if max_gap_bytes is not None: |
| 754 | + kwargs["max_gap_bytes"] = max_gap_bytes |
| 755 | + if max_coalesced_bytes is not None: |
| 756 | + kwargs["max_coalesced_bytes"] = max_coalesced_bytes |
| 757 | + async for group in Store.get_ranges(self, key, byte_ranges, prototype=prototype, **kwargs): |
| 758 | + yield group |
| 759 | + |
| 760 | + async def get_partial_values( |
| 761 | + self, |
| 762 | + prototype: BufferPrototype, |
| 763 | + key_ranges: Iterable[tuple[str, ByteRequest | None]], |
| 764 | + ) -> list[Buffer | None]: |
| 765 | + """Partial-value reads built on `self.get`, so each fetch pays latency. |
| 766 | +
|
| 767 | + Issues one `self.get` per `(key, byte_range)` pair instead of the |
| 768 | + `WrapperStore` delegation, which would bypass this wrapper's `get` and |
| 769 | + therefore the synthetic latency. |
| 770 | + """ |
| 771 | + return list( |
| 772 | + await asyncio.gather( |
| 773 | + *( |
| 774 | + self.get(key, prototype=prototype, byte_range=byte_range) |
| 775 | + for key, byte_range in key_ranges |
| 776 | + ) |
| 777 | + ) |
| 778 | + ) |
0 commit comments