diff --git a/temporalio/contrib/streams/_plugin.py b/temporalio/contrib/streams/_plugin.py index 0fcfd1d77..61f24acac 100644 --- a/temporalio/contrib/streams/_plugin.py +++ b/temporalio/contrib/streams/_plugin.py @@ -84,13 +84,12 @@ def __init__(self, name: str, config: StreamStoreConfig) -> None: # Imported here, since that module imports this one. from temporalio.contrib.streams._replay_reads import ( BatchCache, - ByteCache, ReadStates, ) self._batches = BatchCache() self._read_states = ReadStates() - self._follows = ByteCache() + self._follows = BatchCache() def configure_client(self, config: ClientConfig) -> ClientConfig: """Register this store on the client. diff --git a/temporalio/contrib/streams/_replay_reads.py b/temporalio/contrib/streams/_replay_reads.py index a61992512..e1b99b2e1 100644 --- a/temporalio/contrib/streams/_replay_reads.py +++ b/temporalio/contrib/streams/_replay_reads.py @@ -96,73 +96,46 @@ def before(self, chain: ReadChain) -> bytes | None: return state if seq == chain.seq - 1 else None +def range_key(range: ReadRangeRequest) -> bytes: + """The cache key of one read: its holder, its range and its digest.""" + return b"\0".join( + [ + range.floor_holder.encode(), + range.after.encode(), + range.to.encode(), + range.records_digest, + ] + ) + + class BatchCache: - """The newest batches positions-only reads read in this process.""" + """The newest answers reads read in this process, by key.""" def __init__(self, size: int = _CACHE_SIZE, max_bytes: int = _CACHE_BYTES) -> None: - """Keep up to ``size`` batches of no more than ``max_bytes`` in all.""" + """Keep up to ``size`` answers of no more than ``max_bytes`` in all.""" self._size = size self._max_bytes = max_bytes self._bytes = 0 - self._batches: OrderedDict[bytes, bytes] = OrderedDict() - - @staticmethod - def key(range: ReadRangeRequest) -> bytes: - """The key of one read: its holder, its range and its digest.""" - return b"\0".join( - [ - range.floor_holder.encode(), - range.after.encode(), - range.to.encode(), - range.records_digest, - ] - ) - - def put(self, range: ReadRangeRequest, response: ReadResponse) -> None: - """Keep ``response``, whose bodies are already decoded.""" - key = self.key(range) - data = response.SerializeToString() - self._bytes += len(data) - len(self._batches.get(key, b"")) - self._batches[key] = data - self._batches.move_to_end(key) - while self._batches and ( - len(self._batches) > self._size or self._bytes > self._max_bytes - ): - self._bytes -= len(self._batches.popitem(last=False)[1]) - - def take(self, range: ReadRangeRequest) -> ReadResponse | None: - """The batch of ``range``, which stays cached for a later replay.""" - data = self._batches.get(self.key(range)) - return ReadResponse.FromString(data) if data is not None else None - - def clear(self) -> None: - """Drop every batch, so the next decode reads from the store.""" - self._batches.clear() - self._bytes = 0 - - -class ByteCache: - """The newest values put in this process, by key.""" - - def __init__(self, size: int = _CACHE_SIZE) -> None: - """Keep up to ``size`` values.""" - self._size = size self._values: OrderedDict[bytes, bytes] = OrderedDict() def put(self, key: bytes, value: bytes) -> None: """Keep ``value`` under ``key``.""" + self._bytes += len(value) - len(self._values.get(key, b"")) self._values[key] = value self._values.move_to_end(key) - while len(self._values) > self._size: - self._values.popitem(last=False) + while self._values and ( + len(self._values) > self._size or self._bytes > self._max_bytes + ): + self._bytes -= len(self._values.popitem(last=False)[1]) def get(self, key: bytes) -> bytes | None: - """The value under ``key``, which stays for a later replay.""" + """The value under ``key``, which stays cached for a later replay.""" return self._values.get(key) def clear(self) -> None: - """Drop every value.""" + """Drop every value, so the next decode reads from the store.""" self._values.clear() + self._bytes = 0 def positions_payload( @@ -316,8 +289,8 @@ async def _batch(self, payload: Payload) -> Payload: # A read with no stored record has no record at all, since a # SUPERSEDED record comes only before a stored one. batch = answer - elif (cached := self._plugin._batches.take(range)) is not None: - batch = cached + elif (cached := self._plugin._batches.get(range_key(range))) is not None: + batch = ReadResponse.FromString(cached) else: if chain is not None and _STATE_ELIDED_KEY in payload.metadata: state = self._plugin._read_states.before(chain) diff --git a/temporalio/contrib/streams/_workflow.py b/temporalio/contrib/streams/_workflow.py index 6b764c476..5b2653a8d 100644 --- a/temporalio/contrib/streams/_workflow.py +++ b/temporalio/contrib/streams/_workflow.py @@ -69,6 +69,7 @@ is_chained, is_positions, positions_payload, + range_key, ) from temporalio.contrib.streams._topic import StreamTopic, resolve_topic from temporalio.contrib.streams._wakes import register_wake_context, wake_state @@ -1003,7 +1004,7 @@ async def _read_batch(input: _ReadInput) -> RawValue: ) # The Worker turns the result back into this batch before Workflow code # sees it, from here while the run stays in this process. - store._batches.put(range, response) + store._batches.put(range_key(range), response.SerializeToString()) chain = ReadChain(f"{info.workflow_run_id}/{input.reader_id}", input.read_seq) if input.chained: # The input is the previous read's answer state, which that read's diff --git a/tests/contrib/streams/test_replay_reads.py b/tests/contrib/streams/test_replay_reads.py index 0df92fd36..e5d81d1aa 100644 --- a/tests/contrib/streams/test_replay_reads.py +++ b/tests/contrib/streams/test_replay_reads.py @@ -35,6 +35,7 @@ ReadStates, ReplayReadCodec, is_positions, + range_key, ) from temporalio.contrib.streams._workflow import _ReadInput, read_activity from temporalio.contrib.streams.memory import MemoryStreams @@ -378,9 +379,10 @@ def test_the_batch_cache_keeps_its_bytes_bounded(): ranges = [ReadRangeRequest(after=str(i), to=str(i + 1)) for i in range(3)] big = ReadResponse(cursor="x" * 1400) for read in ranges: - cache.put(read, big) - assert cache.take(ranges[0]) is None - assert cache.take(ranges[1]) == big and cache.take(ranges[2]) == big + cache.put(range_key(read), big.SerializeToString()) + assert cache.get(range_key(ranges[0])) is None + assert cache.get(range_key(ranges[1])) == big.SerializeToString() + assert cache.get(range_key(ranges[2])) == big.SerializeToString() async def test_a_first_range_trimmed_past_its_hold_is_lost_not_other_records(