refactor: separate pub/sub interface from durable task queue - #4480
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
fd7edd7 to
d239557
Compare
Benchmark resultsCompared against |
igor-kupczynski
left a comment
There was a problem hiding this comment.
💯 to have the PubSub interface and related primitives; i've raised some comments.
Can we split this PR? At the very least into the Hijack fix, but otherwise when ready you can prompt cursor/claude to split into 2-3 stacked PRs for an easier review.
| return poolConn.Conn(), nil | ||
| // Hijack removes the connection from the pool's accounting: pgxlisten | ||
| // owns it and closes it when the listener exits, so the pool can be | ||
| // closed cleanly and slots aren't leaked across listener reconnects. | ||
| return poolConn.Hijack(), nil | ||
| }, |
There was a problem hiding this comment.
@abelanger5 this deserves to be in its own PR imho; this is a bug in the current OSS version (bug as in, this can lead a to a very slow connection leak if I understand this correctly).
| // Topic identifies a best-effort pub/sub destination. | ||
| type Topic struct { | ||
| name string | ||
| tenantStream bool |
There was a problem hiding this comment.
nit, but can you rename to isTenantStream? in my mind is conveys the boolean-ness better
There was a problem hiding this comment.
also it's maybe a smell as it's closed for extension
an enum would be preferred imho
| // semantics: handler errors are logged, never redelivered. It returns a | ||
| // cleanup function that should be called when the subscription is no | ||
| // longer needed. | ||
| Sub(topic Topic, handler AckHook) (func() error, error) |
There was a problem hiding this comment.
I've scratched my head why do we need ACK hook for a best-effort pubsub, but this is just a re-use of an existing class right? Maybe define an alias or refine calling it Handler?
There was a problem hiding this comment.
Yeah, this is because it's reusing AckHook from the old implementation. We can rename AckHook to MsgHandler, it'll be more clear that way.
There was a problem hiding this comment.
Possibly missing a tenant-stream publish for task-cancelled on the API cancellation path?
Before this PR, this SendMessage was implicitly mirrored to the tenant fanout exchange (as the old pubMessage mirrored every durable message with a non-nil TenantID). Now that mirroring is explicit, this path only publishes to the durable queue.
There was a problem hiding this comment.
I had Fable enumerate the messages consumed from the tenant exchange:
┌───────────────────────────────┬───────────────────────┬──────────────────────────┐
│ Msg ID │ Handler │ Becomes │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ created-task │ msgsToWorkflowEvent │ STEP_RUN / STARTED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ task-completed │ msgsToWorkflowEvent │ STEP_RUN / COMPLETED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ task-failed │ msgsToWorkflowEvent │ STEP_RUN / FAILED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ task-cancelled │ msgsToWorkflowEvent │ STEP_RUN / CANCELLED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ task-stream-event │ msgsToWorkflowEvent │ STEP_RUN / STREAM │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ workflow-run-finished │ both handlers │ WORKFLOW_RUN / COMPLETED │
│ │ │ |FAILED|CANCELLED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ workflow-run-finished-candida │ isMatchingWorkflowRun │ wakes the run-finished │
│ te │ V1 only │ poll │
└───────────────────────────────┴───────────────────────┴──────────────────────────┘
So yes, task-cancelled is consumed - nice catch! Will audit these as well.
There was a problem hiding this comment.
Instead of doing each Pub manually, we can probably just pull this enum into a helper and automatically pub these messages - seems like it'd be cleaner than the manual pubs everywhere.
There was a problem hiding this comment.
I had Fable enumerate the messages consumed from the tenant exchange:
┌───────────────────────────────┬───────────────────────┬──────────────────────────┐
│ Msg ID │ Handler │ Becomes │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ created-task │ msgsToWorkflowEvent │ STEP_RUN / STARTED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ task-completed │ msgsToWorkflowEvent │ STEP_RUN / COMPLETED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ task-failed │ msgsToWorkflowEvent │ STEP_RUN / FAILED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ task-cancelled │ msgsToWorkflowEvent │ STEP_RUN / CANCELLED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ task-stream-event │ msgsToWorkflowEvent │ STEP_RUN / STREAM │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ workflow-run-finished │ both handlers │ WORKFLOW_RUN / COMPLETED │
│ │ │ |FAILED|CANCELLED │
├───────────────────────────────┼───────────────────────┼──────────────────────────┤
│ workflow-run-finished-candida │ isMatchingWorkflowRun │ wakes the run-finished │
│ te │ V1 only │ poll │
└───────────────────────────────┴───────────────────────┴──────────────────────────┘
So yes, task-cancelled is consumed - nice catch! Will audit these as well.
| switch strings.ToLower(pubsubKind) { | ||
| case "postgres": | ||
| // never dc.Pool and never the pgbouncer URL: the pub/sub owns its pool, | ||
| // and LISTEN does not survive transaction pooling | ||
| pubsubPoolConfig, err := pgxpool.ParseConfig(dc.DirectDatabaseURL) | ||
|
|
||
| if err != nil { | ||
| return nil, nil, fmt.Errorf("could not parse direct database url for pubsub: %w", err) | ||
| } | ||
|
|
||
| setPgxApplicationName(pubsubPoolConfig, cf.OpenTelemetry.ServiceName+":pubsub") | ||
|
|
||
| pubsubPoolConfig.MaxConns = cf.MessageQueue.PubSub.Postgres.MaxConns | ||
| pubsubPoolConfig.MinConns = cf.MessageQueue.PubSub.Postgres.MinConns | ||
| pubsubPoolConfig.ConnConfig.Tracer = newOTelPgxTracer() | ||
|
|
||
| pubsubPool, err := pgxpool.NewWithConfig(context.Background(), pubsubPoolConfig) | ||
|
|
||
| if err != nil { | ||
| return nil, nil, fmt.Errorf("could not create pubsub pool: %w", err) | ||
| } | ||
|
|
||
| pubsubRepo, cleanupPubSubRepo := repov1.NewMessageQueueRepositoryWithPool(&l, pubsubPool) | ||
|
|
||
| cleanupPg, ps, err := pgmq.NewPubSub( | ||
| pubsubRepo, | ||
| pgmq.WithPubSubLogger(&l), | ||
| ) | ||
|
|
||
| if err != nil { | ||
| pubsubPool.Close() | ||
| return nil, nil, fmt.Errorf("could not init postgres pubsub: %w", err) | ||
| } | ||
|
|
||
| pubsubv1 = ps | ||
| cleanupPubSub = func() error { | ||
| if err := cleanupPg(); err != nil { | ||
| return err | ||
| } | ||
|
|
||
| if err := cleanupPubSubRepo(); err != nil { | ||
| return err | ||
| } | ||
|
|
||
| pubsubPool.Close() | ||
| return nil | ||
| } | ||
| case "rabbitmq": | ||
| if pubsubURL == "" { | ||
| return nil, nil, fmt.Errorf("using RabbitMQ as pubsub requires a URL to be set") | ||
| } | ||
|
|
||
| cleanupRmq, ps, err := rabbitmq.NewPubSub( | ||
| rabbitmq.WithPubSubURL(pubsubURL), | ||
| rabbitmq.WithPubSubLogger(&l), | ||
| rabbitmq.WithPubSubMaxPubChannels(cf.MessageQueue.PubSub.RabbitMQ.MaxPubChans), | ||
| rabbitmq.WithPubSubMaxSubChannels(cf.MessageQueue.PubSub.RabbitMQ.MaxSubChans), | ||
| // compression settings inherit from the durable queue: both sides | ||
| // publish to the same tenant exchanges, so they must agree | ||
| rabbitmq.WithPubSubGzip( | ||
| cf.MessageQueue.RabbitMQ.CompressionEnabled, | ||
| cf.MessageQueue.RabbitMQ.CompressionThreshold, | ||
| ), | ||
| ) | ||
|
|
||
| if err != nil { | ||
| return nil, nil, fmt.Errorf("could not init rabbitmq pubsub: %w", err) | ||
| } | ||
|
|
||
| pubsubv1 = ps | ||
| cleanupPubSub = cleanupRmq | ||
| default: | ||
| return nil, nil, fmt.Errorf("invalid pubsub kind %q, must be 'rabbitmq' or 'postgres'", pubsubKind) | ||
| } |
There was a problem hiding this comment.
nit: should be refactored into a helper
There was a problem hiding this comment.
Instead of doing each Pub manually, we can probably just pull this enum into a helper and automatically pub these messages - seems like it'd be cleaner than the manual pubs everywhere.
…helper, loader helper - Replace Topic's tenantStream bool with a TopicKind enum (tenant-stream / scheduler-partition) - Rename AckHook to MsgHandler: on the durable queue it acts as a pre/post-ack hook, on the pub/sub it's the sole handler with no ack semantics - Add msgqueue.PubTenantMessage(ctx, l, mq, ps, queue, msg): single source of truth for the stream-consumed message IDs; does the durable send (when a queue is given) plus the tenant-stream publish (when the ID is one the dispatcher's streams consume), collapsing every producer call site to one call. The API cancellation path (handleCancelTasks) now goes through it too, restoring exact parity with the old implicit mirror for task-cancelled - Extract the pub/sub construction in the loader into createPubSubV1
|
|
| } | ||
|
|
||
| func (p *PubSub) IsReady() bool { | ||
| return p.pubChannels.hasActiveConnection() && p.subChannels.hasActiveConnection() |
There was a problem hiding this comment.
@abelanger5 let's discuss this, it feeds into the readiness check and liveness check, so if rabbit is disconnected, the pod will be marked as not ready;
readiness check --> can pod serve traffic (arguably with pubsub, it can even if the pubsub is disconnected)
liveness check --> is pod broken (if yes, kubelet will restart it); so this means that we will restart the pod on rabbit disconnect; is it what we want?
There was a problem hiding this comment.
I think we should likely continue to expose IsReady but not call it from the probes; I agree it doesn't make sense to reject all API work if our pub/sub is down
| handleMsg := func(task *msgqueue.Message, ackId *int64) { | ||
| if ackId != nil { | ||
| if err := p.repo.AckMessage(subscribeCtx, *ackId); err != nil { | ||
| p.l.Error().Err(err).Msg("error acking message") | ||
| } | ||
| } | ||
|
|
||
| if err := handler(task); err != nil { | ||
| p.l.Error().Err(err).Msgf("error handling pubsub message %s", task.ID) | ||
| } | ||
| } |
There was a problem hiding this comment.
-
Handlers can receive gzipped payloads here when the durable queue is RabbitMQ with
compressionEnabled: trueaspubMessagecompresses in place (msg.Payloads = compressed; msg.Compressed = true) andPubTenantMessagecallsSendMessagebefore handing the same pointer toPub--> the Postgres pub/sub publishes the mutated message. -
The RabbitMQ pub/sub decompresses on its Sub side; this backend never does, so
handler(task)gets gzip bytes withtask.Compressed == truewhich breaks subscribers.
Check for .Compressed also in the postgres handler and decompress as required.
There was a problem hiding this comment.
should be resolved by a825a7b#diff-0a18c7ff5e3579de9defe982a69865977cb787a893184ce4e1ea2369e56f00ad
…end payloads - Health probes no longer gate on the pub/sub's IsReady: a disconnected pub/sub degrades stream notifications but the pod can still serve traffic, and restarting it on a broker disconnect would make things worse. The PubSub interface keeps IsReady for observability. - Fix compressed payloads leaking across backends: when the durable queue is rabbitmq with compression enabled, pubMessage compressed msg.Payloads in place before PubTenantMessage handed the same pointer to the pub/sub, so a postgres pub/sub published gzip bytes its subscribers never decompressed. Fixed at both ends: pubMessage now compresses a shallow copy (compression is a rabbitmq wire concern), and the postgres subscriber decompresses Compressed messages via the new msgqueue.DecompressPayloads (moved from the rabbitmq package so all backends share it).
| handleMsg := func(task *msgqueue.Message, ackId *int64) { | ||
| if ackId != nil { | ||
| if err := p.repo.AckMessage(subscribeCtx, *ackId); err != nil { | ||
| p.l.Error().Err(err).Msg("error acking message") | ||
| } | ||
| } | ||
|
|
||
| if err := handler(task); err != nil { | ||
| p.l.Error().Err(err).Msgf("error handling pubsub message %s", task.ID) | ||
| } | ||
| } |
There was a problem hiding this comment.
should be resolved by a825a7b#diff-0a18c7ff5e3579de9defe982a69865977cb787a893184ce4e1ea2369e56f00ad
…gres IsReady, constants, comment cleanups - Move the gzip Compressor (and its tests) from the rabbitmq package into msgqueue so compression is available to all implementations; rabbitmq embeds msgqueue.Compressor. Making compression a msgqueue-level config option is deferred: enabling it for postgres today would break mixed-version fleets, since pre-split engines' postgres consumers never decompress. - Add Message.Clone and use it for the compress-copy and chunk-split paths. - Pull magic values into constants: maxPayloadSize (16MB), consumerTimeout (5m), msgqueue.DefaultCompressionThreshold (5KB) — applied to both the durable rabbitmq impl and the pubsub. - Postgres PubSub.IsReady now pings the dedicated pool (2s timeout) instead of returning true; informational only, still not wired into health probes. - Reword the postgres PubSub type comment, document the direct (non-pgbouncer) URL requirement on NewPubSub, and drop stale comments in health.New and the status-update processors.
|
|
|
|
The test failed deterministically on CI with an absurd growth figure (17592186044415.85 MB = 2^64 bytes). Two compounding problems: 1. The baseline was systematically inflated: tests added to this package by the pub/sub split (the moved gzip tests and pubsub tests) run before this test in the same binary and leave heap state behind — notably sync.Pool contents, which survive one GC in the victim cache and are only freed by a second GC. The test's baseline GC was that first GC, its final GC the second, so the heap consistently ended up ~0.3-0.8MB *below* baseline on constrained runners. 2. Alloc is a uint64, so the negative growth underflowed to ~2^64. Fixed by double-GC'ing before the baseline (clears victim caches from preceding tests) and subtracting as floats so a genuinely shrinking heap reads as negative instead of failing. Verified stable (0.47-0.54 MB) across repeated full-package runs at GOMAXPROCS=2, which reproduced the CI failure before the fix.
There was a problem hiding this comment.
Note
Igor: This was posted by AI a bit too eagerly; keeping as is though and will comment below
Did a full pass over the final state of the branch (checked out locally, read the changed files in full, audited producer/consumer call sites, ran go build ./... and the msgqueue/loader unit tests).
What I verified:
- Producer parity with the old implicit mirror: every producer of the seven stream-consumed message IDs (
created-task,task-completed,task-failed,task-cancelled,task-stream-event,workflow-run-finished,workflow-run-finished-candidate) now publishes explicitly viaPubTenantMessageorps.Pub— dispatcher completed/failed, scheduler cancel/internal-retry/dead-letter, task controller cancel + run-finished-candidate, OLAP signaler created-task + republish + run-finished, ingestor stream events. The one implicit mirror that was dropped without replacement (sendTaskCancellationsToDispatcher, which sendstask-cancelledwithSignalTaskCancelledPayloadto dispatcher queues) previously decoded to zero-valuedCancelledTaskPayloadjunk on the stream, so dropping it is an improvement. - Wire compatibility: tenant topic (
<uuid>_v1) and scheduler partition (<p>_scheduler_v1) names, exchange/queue declare args, and postgres NOTIFY channel names all match legacy, so mixed-version fleets interoperate. - Cross-backend compression: durable
pubMessagenow compresses a clone (no caller mutation), and the postgres subscriber decompressesCompressedmessages — the earlier review finding is properly fixed at both ends. SharedBufferedTenantReadercorrectly passesWithKind(PostAck), so the single-handler pubsub Sub maps onto the buffer's active hook.- Pool isolation: postgres pubsub gets its own small pool on the direct (non-pgbouncer) URL; rabbitmq pubsub owns its connections. Health probes no longer gate on pubsub readiness, per the review discussion.
- Fallback-row hygiene: >8KB postgres fallback rows are acked by subscribers and reaped by the retention controller's existing
CleanupMessageQueueItemsloop regardless of msgqueue kind. - Config inheritance is backwards compatible (incl. legacy
SERVER_TASKQUEUE_*aliases) and covered byloader_pubsub_test.go; docs updated; CI fully green on the latest commit.
Non-blocking follow-ups (none of these should hold up the merge):
rabbitmq/pubsub.goSubleavesmsg.Compressed = trueafter decompressing (the postgres pubsub resets it tofalse; the durable rabbitmq subscriber has the same pre-existing quirk). Harmless today since no handler re-publishes the same*Message, but it's a footgun if one ever does — worth normalizing.tenantStreamMsgIDsinpubsub.gois a hand-maintained mirror of the dispatcher's stream consumers (msgsToWorkflowEvent/isMatchingWorkflowRunV1). A small drift test asserting the two stay in sync would prevent a silently dead stream event type when a new msg ID is added.- Rabbitmq pubsub
PubsetsExpiration: "0"on all publishes, whereas the old tenant-exchange mirror published without a per-message TTL — messages now drop when a subscriber queue momentarily has no ready consumer. Consistent with the at-most-once contract, just noting the subtle semantic change. - Postgres >8KB fallback rows deliver to only one subscriber per message (
ReadMessages+Ack), unlike the rabbitmq fanout. This is pre-existing postgres behavior, but it's now a per-implementation divergence of thePubSubcontract — worth documenting on the interface.
|
Only 1 and 4 make sense to me; I'll push a commit. I'm considering 2 as well, but only if that's small. |
…cument postgres fanout divergence - The rabbitmq PubSub subscriber decompressed payloads but left msg.Compressed = true, unlike the postgres subscriber; handlers now never see a decompressed message still claiming to be compressed. - Document on PubSub.Sub and the postgres implementation that >8KB fallback rows are delivered to exactly one subscriber, unlike inline NOTIFY payloads and the rabbitmq fanout, which reach every subscriber.
igor-kupczynski
left a comment
There was a problem hiding this comment.
Amazing stuff Alexander, thank you!
|
I'm going to merge to unblock some of my work; given the 7 commit history, I'm going to squash. |
Selectable per deployment with SERVER_MSGQUEUE_PUBSUB_KIND=nats and SERVER_MSGQUEUE_PUBSUB_NATS_URL. NATS implements only the best-effort pub/sub side of the split introduced in #4480; durable queues stay on RabbitMQ and Postgres. Core NATS, no JetStream, so delivery is at-most-once. Connect is synchronous and validates the server's max_payload against the 16MiB the durable RabbitMQ path allows, because task stream events routinely exceed the NATS default of 1MiB and a startup failure is easier to diagnose than a later oversized publish. ReconnectBufSize(-1) makes publishes fail fast while disconnected rather than silently buffering messages a best-effort subscriber will never wait for. The other backends namespace installations at the connection level (a RabbitMQ vhost, a Postgres database); NATS namespaces by subject, so SERVER_MSGQUEUE_PUBSUB_NATS_SUBJECT_PREFIX is the parity knob for installations that share a NATS server. Username and password are passed as connect options rather than embedded in the URL so that reconnects to gossip-discovered cluster peers authenticate too.
Selectable per deployment with SERVER_MSGQUEUE_PUBSUB_KIND=nats and SERVER_MSGQUEUE_PUBSUB_NATS_URL. NATS implements only the best-effort pub/sub side of the split introduced in #4480; durable queues stay on RabbitMQ and Postgres. Core NATS, no JetStream, so delivery is at-most-once. Connect is synchronous and validates the server's max_payload against the 16MiB the durable RabbitMQ path allows, because task stream events routinely exceed the NATS default of 1MiB and a startup failure is easier to diagnose than a later oversized publish. ReconnectBufSize(-1) makes publishes fail fast while disconnected rather than silently buffering messages a best-effort subscriber will never wait for. The other backends namespace installations at the connection level (a RabbitMQ vhost, a Postgres database); NATS namespaces by subject, so SERVER_MSGQUEUE_PUBSUB_NATS_SUBJECT_PREFIX is the parity knob for installations that share a NATS server. Username and password are passed as connect options rather than embedded in the URL so that reconnects to gossip-discovered cluster peers authenticate too.
* test(dispatcher): guard tenant-stream msg IDs against drift The pub/sub split (#4480) made the implicit durable-send mirror to tenant streams an explicit allowlist (msgqueue.tenantStreamMsgIDs), which nearly dropped task-cancelled on the API cancellation path during review. Nothing kept that allowlist in sync with what the dispatcher's streams consume. Make the dispatcher the checkable source of truth: msgsToWorkflowEvent's switch becomes a lookup in a package-level workflowEventConverters map and isMatchingWorkflowRunV1's switch a lookup in workflowRunMatchers, preserving behavior. msgqueue exports the allowlist via TenantStreamMsgIDs(), and TestTenantStreamMsgIDsInSync asserts set equality between the two, with failure messages saying exactly what to update on either side. * refactor(dispatcher): extract tenant-stream event handling into its own file Move workflowEventConverters, workflowRunMatchers, msgsToWorkflowEvent and isMatchingWorkflowRunV1 out of server.go into tenant_stream_events.go, next to the drift test that asserts their keys against msgqueue's allowlist. Collapse the repeated decode/loop/append scaffolding into two small generic adapters (eventConverter, runMatcher), so each map entry is just its payload type and field mapping; the workflow-run-finished status switch stays inside its closure. Drop the unused DispatcherImpl receivers from both functions. No behavior change: map misses and empty match results return exactly what they did before.
Description
The
msgqueueinterface currently has two responsibilities:These use-cases are long-term divergent and we'd like to experiment with other pub/sub technology, so this PR splits the pub/sub implementation from the durable task queue.
Type of change
What's Changed
PubSubimplementationChecklist
Changes have been:
🤖 AI Disclosure