Skip to content

Commit e160426

Browse files
committed
Let one append take plain values and links as one batch.
An agent writes an event and the link that goes with it in one append, so a resend dedupes them together.
1 parent c5ae004 commit e160426

4 files changed

Lines changed: 133 additions & 25 deletions

File tree

‎temporalio/contrib/streams/README.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,9 @@ async for item in stream_handle(client, agent_id, topic=EVENTS).follow():
186186
`join(*, name, until=None)` writes a `JOIN`. It says the producer waited
187187
for the link up to `until`, or to the linked stream's end. A link names a
188188
Workflow's stream only. Any other ref raises `ValueError`.
189+
- **Mixed batches.** `append` also takes a `StreamLink` next to plain values,
190+
so `append(event, StreamLink(LinkEvent.OPEN, ref, name))` writes the event
191+
and its link as one batch, with one dedupe identity.
189192
- **Following.** `StreamHandle.follow()` and `ActivityStreamHandle.follow()`
190193
read a stream and the streams it links, from client or Activity code.
191194
Workflow code can't follow.

‎temporalio/contrib/streams/_handles.py‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -122,10 +122,14 @@ def attempt(self) -> int:
122122
"""The attempt this producer writes."""
123123
return self._attempt
124124

125-
async def append(self, *values: T) -> Cursor:
125+
async def append(self, *values: T | StreamLink) -> Cursor:
126126
"""Append ``values`` as one batch and return the cursor of its last record.
127127
128-
The batch lands whole and in order. The producer moves its sequence
128+
A :class:`temporalio.contrib.streams.StreamLink` among the values is a
129+
``LINK`` record, as :meth:`link` and :meth:`join` write it, and the
130+
other values are ``DATA`` records. So one batch can hold a value and
131+
the link that goes with it, under one dedupe identity. The batch lands
132+
whole and in order. The producer moves its sequence
129133
past the batch only when the append succeeds, so a call that raised
130134
:class:`temporalio.contrib.streams.StreamOutcomeUnknownError` can be
131135
repeated with the same values: the store answers a repeat of the
@@ -151,11 +155,19 @@ async def append(self, *values: T) -> Cursor:
151155
StreamRefusedError: The store refused the batch, for example
152156
because it is out of memory. Nothing was written.
153157
StreamClosedError: The stream refuses appends.
158+
ValueError: A link names nothing, or names a stream that no
159+
Workflow owns. Nothing was written.
154160
"""
161+
items: list[tuple[RecordKind, Any, StreamLink | None]] = [
162+
(RecordKind.LINK, None, v)
163+
if isinstance(v, StreamLink)
164+
else (RecordKind.DATA, v, None)
165+
for v in values
166+
]
155167
async with self._lock:
156-
if not values:
168+
if not items:
157169
return self._last
158-
return await self._write(RecordKind.DATA, values)
170+
return await self._write(items)
159171

160172
async def finish(self) -> Cursor:
161173
"""Write ``FINISH`` for this producer on this topic and return its cursor.
@@ -164,7 +176,7 @@ async def finish(self) -> Cursor:
164176
work behind it succeeded, and it does not end anyone's read.
165177
"""
166178
async with self._lock:
167-
return await self._write(RecordKind.FINISH, [None])
179+
return await self._write([(RecordKind.FINISH, None, None)])
168180

169181
async def link(
170182
self,
@@ -189,7 +201,7 @@ async def link(
189201
"""
190202
link = StreamLink(LinkEvent.OPEN, ref, name, after=after)
191203
async with self._lock:
192-
return await self._write(RecordKind.LINK, [value], link)
204+
return await self._write([(RecordKind.LINK, value, link)])
193205

194206
async def join(self, *, name: str, until: Cursor | None = None) -> Cursor:
195207
"""Write a ``LINK`` that says this producer waited for link ``name``.
@@ -203,10 +215,10 @@ async def join(self, *, name: str, until: Cursor | None = None) -> Cursor:
203215
"""
204216
link = StreamLink(LinkEvent.JOIN, None, name, until=until)
205217
async with self._lock:
206-
return await self._write(RecordKind.LINK, [None], link)
218+
return await self._write([(RecordKind.LINK, None, link)])
207219

208220
async def _write(
209-
self, kind: RecordKind, values: Any, link: StreamLink | None = None
221+
self, items: list[tuple[RecordKind, Any, StreamLink | None]]
210222
) -> Cursor:
211223
converter = self._handle._converter
212224
wires = [
@@ -220,7 +232,7 @@ async def _write(
220232
sequence=self._sequence + index,
221233
link=link,
222234
)
223-
for index, value in enumerate(values)
235+
for index, (kind, value, link) in enumerate(items)
224236
]
225237
request = AppendRequest(
226238
stream=self._handle._address(self._topic),

‎temporalio/contrib/streams/_workflow_producer.py‎

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -118,39 +118,53 @@ def topic(self) -> str:
118118
"""The name of the topic this producer is bound to."""
119119
return self._topic
120120

121-
def append(self, value: T) -> None:
122-
"""Append ``value`` on this topic and return at once.
123-
124-
The value is converted here, with the Workflow's payload converter. A
125-
:class:`temporalio.common.RawValue` passes through pre-encoded. The
126-
record is visible when the activation that appended it completes,
121+
def append(self, *values: T | StreamLink) -> None:
122+
"""Append ``values`` on this topic, in order, and return at once.
123+
124+
A :class:`temporalio.contrib.streams.StreamLink` among the values is a
125+
``LINK`` record, as :meth:`link` and :meth:`join` write it, and the
126+
other values are ``DATA`` records. Every value is converted here, with
127+
the Workflow's payload converter, before any record is added, so a
128+
bad value adds none of them. A :class:`temporalio.common.RawValue`
129+
passes through pre-encoded. The records go with the activation, in
130+
one batch with the activation's other records. A record is visible
131+
when the activation that appended it completes,
127132
which is before the Workflow Task completes. If the task then fails or
128133
times out, the record stays, and the task's rerun appends again under
129134
a higher attempt, which readers see as ``SUPERSEDED``.
130135
131136
Raises:
132-
ValueError: The topic was already finished in this run.
133-
temporalio.contrib.streams.StreamError: The value's payload is
137+
ValueError: The topic was already finished in this run, or a link
138+
names nothing or names a stream that no Workflow owns.
139+
temporalio.contrib.streams.StreamError: A value's payload is
134140
over 1.5 MiB. Nothing is appended.
135141
temporalio.workflow.ReadOnlyContextError: Called from a query or
136142
an update validator.
137143
"""
138144
_check_writable("append to")
139145
if self._topic in self._producers.finished:
140146
raise ValueError(f"topic {self._topic!r} was already finished")
147+
records = [
148+
self._link_record(value, None)
149+
if isinstance(value, StreamLink)
150+
else self._data_record(value)
151+
for value in values
152+
]
153+
for record in records:
154+
self._producers.add(record)
155+
156+
def _data_record(self, value: Any) -> OutputRecord:
141157
payload = workflow.payload_converter().to_payloads([value])[0]
142158
if payload.ByteSize() > _MAX_RECORD_BYTES:
143159
raise StreamError(
144160
f"a record of {payload.ByteSize()} bytes on {self._topic!r} is over "
145161
f"the per-record limit of {_MAX_RECORD_BYTES} bytes"
146162
)
147-
self._producers.add(
148-
OutputRecord(
149-
topic=self._topic,
150-
kind=StreamRecordKind.STREAM_RECORD_KIND_DATA,
151-
body=payload,
152-
content_hash=content_hash(payload),
153-
)
163+
return OutputRecord(
164+
topic=self._topic,
165+
kind=StreamRecordKind.STREAM_RECORD_KIND_DATA,
166+
body=payload,
167+
content_hash=content_hash(payload),
154168
)
155169

156170
def link(
@@ -197,6 +211,9 @@ def _add_link(self, link: StreamLink, value: Any) -> None:
197211
_check_writable("link on")
198212
if self._topic in self._producers.finished:
199213
raise ValueError(f"topic {self._topic!r} was already finished")
214+
self._producers.add(self._link_record(link, value))
215+
216+
def _link_record(self, link: StreamLink, value: Any) -> OutputRecord:
200217
record = OutputRecord(
201218
topic=self._topic,
202219
kind=StreamRecordKind.STREAM_RECORD_KIND_LINK,
@@ -211,7 +228,7 @@ def _add_link(self, link: StreamLink, value: Any) -> None:
211228
)
212229
record.body.CopyFrom(payload)
213230
record.content_hash = content_hash(payload)
214-
self._producers.add(record)
231+
return record
215232

216233
async def flush(self) -> None:
217234
"""Return at once, since each record goes with its activation.

‎tests/contrib/streams/test_workflow_producer.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -725,3 +725,79 @@ async def test_a_producer_links_and_joins(client: Client):
725725
assert records[1].value == {"label": "search"}
726726
assert records[2].link == StreamLink(LinkEvent.JOIN, None, "sub-1")
727727
assert [r.sequence for r in records] == [1, 2, 3]
728+
729+
730+
@workflow.defn
731+
class AppendsMixed:
732+
@workflow.run
733+
async def run(self) -> str:
734+
steps = stream_producer(STEPS)
735+
child = StreamRef.for_workflow("child")
736+
steps.append("d1", StreamLink(LinkEvent.OPEN, child, "sub-1"), "d2")
737+
try:
738+
steps.append("lost", StreamLink(LinkEvent.OPEN, child, ""))
739+
except ValueError as error:
740+
return str(error)
741+
return ""
742+
743+
744+
async def test_a_workflow_appends_values_and_links_in_one_call(client: Client):
745+
streams = await connect_with(client, MemoryStreams())
746+
workflow_id = new_workflow_id()
747+
async with new_worker(streams, AppendsMixed) as worker:
748+
refused = await streams.execute_workflow(
749+
AppendsMixed.run, id=workflow_id, task_queue=worker.task_queue
750+
)
751+
records = await read_all(stream_handle(streams, workflow_id).read(topic=STEPS))
752+
assert [(r.kind, r.value) for r in records] == [
753+
(RecordKind.DATA, "d1"),
754+
(RecordKind.LINK, None),
755+
(RecordKind.DATA, "d2"),
756+
]
757+
assert records[1].link == StreamLink(
758+
LinkEvent.OPEN, StreamRef.for_workflow("child"), "sub-1"
759+
)
760+
# A bad link adds none of the call's records.
761+
assert "1 to 256" in refused
762+
763+
764+
async def mixed_batch_appends_once(streams: Client) -> None:
765+
async with new_worker(streams, Waits) as worker:
766+
workflow_id = await running_owner(streams, worker.task_queue)
767+
stream = stream_handle(streams, workflow_id, topic=STEPS)
768+
child = StreamRef.for_workflow(f"{workflow_id}-child", topic=STEPS)
769+
batch = (
770+
"sent",
771+
StreamLink(LinkEvent.OPEN, child, "sub-1"),
772+
StreamLink(LinkEvent.JOIN, None, "sub-1"),
773+
"received",
774+
)
775+
first = stream.producer(producer_id="agent", attempt=1)
776+
last = await first.append(*batch)
777+
# A resend of the same batch from a fresh producer object dedupes.
778+
again = stream.producer(producer_id="agent", attempt=1)
779+
assert await again.append(*batch) == last
780+
with pytest.raises(ValueError, match="1 to 256"):
781+
await first.append("x", StreamLink(LinkEvent.OPEN, child, ""))
782+
await first.append("after")
783+
await streams.get_workflow_handle(workflow_id).signal(Waits.finish)
784+
records = await read_all(stream.read())
785+
assert [r.kind for r in records] == [
786+
RecordKind.DATA,
787+
RecordKind.LINK,
788+
RecordKind.LINK,
789+
RecordKind.DATA,
790+
RecordKind.DATA,
791+
]
792+
assert records[3].cursor == last
793+
assert records[1].link == StreamLink(LinkEvent.OPEN, child, "sub-1")
794+
assert records[2].link == StreamLink(LinkEvent.JOIN, None, "sub-1")
795+
assert [r.sequence for r in records] == [1, 2, 3, 4, 5]
796+
797+
798+
async def test_a_producer_appends_values_and_links_as_one_batch(client: Client):
799+
await mixed_batch_appends_once(await connect_with(client, MemoryStreams()))
800+
801+
802+
async def test_a_resent_mixed_batch_dedupes_on_redis(client: Client):
803+
await mixed_batch_appends_once(await connect_with(client, redis_store()))

0 commit comments

Comments
 (0)