Receive channel notifications with asyncio.Protocol instead of streams

[cboulay]

Draft — not for merge yet. Opening this so @griffinmilsap can look it over before it goes anywhere.

What and why

Profiling a two-process SHM hop showed the time is not where I expected. Serialization is ~14us of a ~97us hop; the dominant stage is 50us between the publisher finishing broadcast() and the subscriber's channel waking up. That turned out to be the streams API, not the transport and not ezmsg.

asyncio.open_connection delivers bytes by resolving a future and scheduling the reading task, costing a full event-loop iteration per message before any work happens. An asyncio.Protocol has data_received() called synchronously from the transport's read callback, so the notification is cached and subscribers notified right there.

Measured

Two-process SHM hop, 8 KB AxisArray, medians over 2400 messages, interleaved A/B on this base:

stage	streams	Protocol
pub exit → rx wakeup	50.6 us	15.6 us	3.2x
end-to-end	97.0 us	69.0 us	−29%
throughput	17,337 msg/s	18,697 msg/s	+8%

Protocol won every round. Corroborating detail: latency is flat from 16 B to 512 KB payloads, which is what pointed at fixed scheduling cost rather than serialization in the first place.

uvloop

Written to public asyncio API only — no reader._buffer, no transport internals — so uvloop stays a drop-in on POSIX. Verified end-to-end under uvloop with no code changes: all messages delivered, clean, ~50 us end-to-end. So this gets stock asyncio most of the way, and uvloop later would roughly halve the original baseline.

Design

frameproto.FramedProtocol runs a connection in two phases:

• handshake — sequential, driven by await read_exactly(). One task wakeup per read, which is fine since it happens once.
• dispatch — start_dispatch() flips to frames_available() being called inline from data_received.

SHM resize is the one path needing an await, which a read callback can't do. On a generation change the protocol pauses reading and parks the frame; a task then re-dispatches it against the new segment, or drops it and releases backpressure if the attach fails. Covered by test_shm_grow and a rewritten reattach unit test.

Things worth a reviewer's attention

• Behaviour change: a frame-dispatch exception is now logged with traceback and the transport closed, rather than propagating out of the channel task and potentially surfacing through wait_closed() during teardown. Happy to make it propagate if that's preferred.
• Now atomic: because dispatch never awaits, delivery is atomic w.r.t. other tasks on the loop. The task-based version wasn't. This is a behavioural difference even though it looks like a pure win.
• Test rewrite: test_channel_preserves_cached_message_during_shm_reattach drove _publisher_connection with a StreamReader, which no longer exists. Rewritten against a real ChannelProtocol with a fake transport, and strengthened — it now also asserts the segment swapped, the parked frame was re-dispatched, the client was notified, and the buffer fully consumed.

Deliberately not included

The publisher's ack receive was built and passed the full suite, but measured consistently ~12% worse end-to-end across 5 interleaved rounds. Localized by stage: broadcast() itself and the subscriber side were unchanged; the entire cost was bytes leaving the publisher later. Working hypothesis is that inline ack handling delays the transport's write scheduling, where the ack task previously competed on equal footing in the ready queue — unverified, flagging it as a hypothesis. Reverted rather than shipped.

The graph server connections stay on streams: control plane, not per-message.

Testing

Full suite green on this base: 424 passed, 1 skipped. Lint clean.

Measurements are single-machine (Darwin, arm64, Python 3.13). Worth replicating on Linux before anyone leans on the numbers.

🤖 Generated with Claude Code