Batch async worker Pub/Sub publishes - #45
BinaryFiddler wants to merge 4 commits into
Conversation
BinaryFiddler
left a comment
BinaryFiddler
left a comment
There was a problem hiding this comment.
Automated review pass (Claude Code), verified against the pinned google-cloud-pubsub==2.15.2 source and a live batching repro rather than from the diff alone. Six findings inline, all on code this PR introduces — nothing blocking-by-default, but #1 and #2 are worth resolving before merge.
- #1 the outer batch cap of 250 collides with the native cap of 250, so full batches split
[249, 1]— reproduced. - #2
_gapic_publishinstrumentation is unguarded, andBatch._commitonly catchesGoogleAPIError, so a metrics failure can strand a whole batch's futures. - #3 moving off
max_messages=1means one poison message can drop up to 249 good ones;PublishErroris missing from the transient set. - #4 the shutdown drain's
qsize()snapshot loses anything_requeueputs back. - #5
live_clientsgauge leaks ifclient.stop()raises. - #6 post-stop
publish()now raises where it previously enqueued, which thediscord_smitesinks turn into retries + sleeps.
Things I specifically checked and found correct: the _stopped-flag guard against wait_for swallowing cancellation after queue.get() completes (the 3.11 wait_for path matches the comment); the asyncio.TimeoutError catch, which fixes a real base-code bug where bare TimeoutError missed asyncio.TimeoutError on <=3.10; metrics.histogram (inherited from DogStatsd, and its _report(metric, "h", value, tags, sample_rate) call matches the _DogStatsd._report override); lease-count / _states bookkeeping and the pool's drain-before-_stop_client ordering; and _topic_metric_tags consistency with _PublisherState._metric_tags.
| self._client, | ||
| project_id, | ||
| topic_id, | ||
| self._settings.max_messages, |
There was a problem hiding this comment.
Outer batch cap equals the native cap, so every full outer batch emits two GAPIC calls (249 + 1).
_acquire_state passes self._settings.max_messages (250) as _PublisherState._max_messages, and _take_outer_batch fills up to that same 250. In google-cloud-pubsub 2.15.2, Batch.publish computes overflow before appending:
new_count = len(self._messages) + 1
overflow = new_size > size_limit or new_count >= self.settings.max_messages
if not self._messages or not overflow: # -> does NOT append the 250thSo the native batch commits at 249 and the 250th message opens a fresh batch. Reproduced against 2.15.2 with these exact settings: batch sizes came out [249, 1].
Consequences: transport_batch fires at ~2x the intended rate, messages_per_transport_batch averages ~125 rather than 250, and each outer flush blocks an extra transport_max_latency_seconds waiting for the straggler batch's commit timer, since _sync_flush awaits every future.
Fix: make the outer cap max_messages - 1, or set the native max_messages to outer + 1.
Aside: the 10 ms commit timer in 2.15.2 is a single client-level thread (_wait_and_commit_sequencers) shared across all topics, so it can also fire mid-flush and split an outer batch further. A slightly larger transport_max_latency_seconds would make batching more deterministic.
| class _InstrumentedPublisherClient(pubsub_v1.PublisherClient): # type: ignore[misc] | ||
| """Meters logical GAPIC batches, not internal physical retry attempts.""" | ||
|
|
||
| def _gapic_publish(self, *args: Any, **kwargs: Any) -> Any: |
There was a problem hiding this comment.
Unguarded instrumentation here can strand an entire batch's futures forever.
Batch._commit (2.15.2, _batch/thread.py:280) wraps the _gapic_publish call in except google.api_core.exceptions.GoogleAPIError only, and runs it on a dedicated Thread-CommitBatchPublisher. Any non-GoogleAPIError raised from this override escapes _commit, the thread dies, and none of the batch's futures are ever resolved.
Two ways that happens: kwargs['topic'] / kwargs['messages'] raise KeyError if a future library version passes those positionally, and anything raised out of metrics.increment / metrics.histogram.
Downstream effect: _sync_flush blocks the full future.result(timeout=35) for up to 250 messages, gets TimeoutError (which _TRANSIENT_PUBLISH_ERRORS treats as transient), and requeues the whole batch — an indefinite 35-seconds-per-attempt stall for that topic.
Fix: use kwargs.get(...) and wrap the metric emission in try/except Exception so instrumentation can never break the publish path.
| @dataclass(frozen=True) | ||
| class PublisherBatchSettings: | ||
| max_bytes: int = 2_000_000 | ||
| max_messages: int = 250 |
There was a problem hiding this comment.
Native batching turns a single poison message into up to 249 dropped good messages.
With the previous BatchSettings(max_messages=1), a server-side rejection failed exactly one message. Now Batch._commit sets the same exception on every future in the batch, so one message the server rejects with a permanent error (e.g. InvalidArgument on a malformed attribute) fails all of its co-batched messages. _is_transient_publish_error returns False for those, so _sync_flush logs and drops them.
The same applies to pubsub_v1.publisher.exceptions.PublishError, which _commit raises when len(response.message_ids) != len(futures) (partial publish, thread.py:299-315). It is a GoogleAPICallError but is absent from _TRANSIENT_PUBLISH_ERRORS, so a partial publish silently drops up to 250 messages that were never acknowledged.
Suggest adding PublishError (and probably Cancelled / Unknown) to the transient set, and/or bounding the outer batch so the blast radius of one poison message is explicit.
| except asyncio.CancelledError: | ||
| pass | ||
|
|
||
| pending_at_stop = self._queue.qsize() |
There was a problem hiding this comment.
The drain uses a queue-size snapshot, so anything requeued during shutdown is silently lost.
pending_at_stop is sampled once. _flush_batch -> _requeue puts transient failures back onto the queue, but their slots have already been decremented from pending_at_stop, so the while pending_at_stop > 0 loop exits with messages still queued; pool._stop_client() then closes the client and the process exits.
Scenario: SIGTERM during a Pub/Sub blip. The drain publishes batch 1, gets ServiceUnavailable for all of it, requeues 250 messages, and drops them.
This loop clearly intends a full drain (unlike the base code's single flush), so either re-check qsize() each iteration with a bounded attempt/deadline cap, or explicitly log and count the residual as shutdown loss rather than letting it vanish.
| async def _stop_client(self) -> None: | ||
| if self._client is None or self._client_stopped: | ||
| return | ||
| self._client_stopped = True |
There was a problem hiding this comment.
live_clients gauge leaks if client.stop() raises, and gauge bookkeeping can fail shutdown.
_client_stopped = True is set before await asyncio.to_thread(self._client.stop). PublisherClient.stop() raises RuntimeError("Cannot stop a publisher already stopped.") (and sequencer.stop() can raise too); on that path _change_live_client_count(-1) never runs, so the process gauge stays permanently inflated — defeating this PR's own "gauge falls from 8 to 1" guardrail — and the _client_stopped flag prevents any retry.
Two smaller points on _change_live_client_count:
- Raising
RuntimeError('live Pub/Sub client count became negative')escalates a pure metrics-accounting inconsistency into apool.stop()failure during shutdown. Logging would be safer than raising. metrics.gaugeis called while holding_live_client_count_lock. Moving it outside the lock keeps a statsd send from serializing client create/close, and avoids holding the lock across afork().
| self._stop_task: Optional[asyncio.Task[None]] = None | ||
|
|
||
| def _assert_running(self) -> None: | ||
| if self._stopped or self._state._stopped or self._pool._stopped: |
There was a problem hiding this comment.
publish / publish_bytes now raise after stop — a new exception surface for existing callers.
_assert_running raises RuntimeError('publisher is stopped'); the base implementation had no _stopped concept and simply enqueued.
Every discord_smite sink constructs AsyncPubSubPublisher directly (osprey_async_plugins/.../register_plugins.py), and AsyncOutputSink._push_one catches the exception, calls logger.exception, increments output_sink.error, and retries with await asyncio.sleep(0.5 * attempt). So a late publish during shutdown now costs retries plus sleeps per message, aborts the remainder of that sink's push() (e.g. async_event_effects_output_sink publishes analytics and webhooks in one push, so the webhook publish gets skipped), and generates error-metric noise.
Worth deciding deliberately whether a post-stop publish should raise or drop-and-count. For what it is worth, I checked that nothing downstream touches the internals that moved into _PublisherState (_client, _queue, _metric_tags, _sync_flush) outside of mocks, so the _PublisherState split itself is caller-safe.
1 participant
Summary
Production baseline
Dashboard: https://us5.datadoghq.com/dashboard/skb-sq3-y4i/osprey-coordinator--final
Window: 2026-08-12 through 2026-08-19 UTC.
Existing queries:
Planned metrics below are unavailable until a Discord worker pins this Osprey commit and rolls it out:
transport_batchcounts logical_gapic_publishcalls. It does not count physical Google API attempts made by the retry policy.live_clientsis an untagged process-local running value. The four stable workers should report the same value, so the pod query is representative per process; multiply it by four for the configured pod total.Target and guardrails
Verification
publish()calls produce one batch./run-tests.sh osprey_async_worker/src/osprey/async_worker/tests/test_publisher.pyunable to get image 'quay.io/coreos/etcd:v3.4.18': permission denied while trying to connect to the Docker daemon socket at unix:///var/run/docker.sock: Get "http://%2Fvar%2Frun%2Fdocker.sock/v1.51/images/quay.io/coreos/etcd:v3.4.18/json": dial unix /var/run/docker.sock: connect: operation not permittedRollback
Set
max_messages=1and construct an unsharedAsyncPubSubPublisher. The public constructor remains available, so rollback requires no payload, topic, or caller change.