From d239557313e03b2c68a39b6657ae3a5d048cf891 Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Tue, 21 Jul 2026 17:23:22 -0700 Subject: [PATCH 1/6] refactor: pub/sub separate from durable task queue --- cmd/hatchet-engine/engine/run.go | 18 +- .../self-hosting/configuration-options.mdx | 30 +- internal/msgqueue/mq_pub_buffer_test.go | 5 +- internal/msgqueue/mq_sub_buffer.go | 23 +- internal/msgqueue/msgqueue.go | 88 ---- internal/msgqueue/postgres/exchange.go | 59 --- internal/msgqueue/postgres/msgqueue.go | 35 +- internal/msgqueue/postgres/pubsub.go | 257 ++++++++++ internal/msgqueue/postgres/pubsub_test.go | 280 +++++++++++ internal/msgqueue/pubsub.go | 102 ++++ internal/msgqueue/pubsub_test.go | 73 +++ internal/msgqueue/rabbitmq/gzip.go | 12 +- internal/msgqueue/rabbitmq/gzip_test.go | 8 +- internal/msgqueue/rabbitmq/pubsub.go | 473 ++++++++++++++++++ internal/msgqueue/rabbitmq/pubsub_test.go | 236 +++++++++ internal/msgqueue/rabbitmq/rabbitmq.go | 172 ++----- internal/msgqueue/rabbitmq/rabbitmq_test.go | 29 +- internal/msgqueue/shared_tenant_reader.go | 37 +- internal/services/admin/admin.go | 15 +- internal/services/admin/server.go | 6 +- internal/services/admin/v1/admin.go | 15 +- internal/services/admin/v1/server.go | 6 +- .../services/controllers/olap/controller.go | 19 + .../olap/process_dag_status_updates.go | 10 +- .../olap/process_task_status_updates.go | 10 +- .../controllers/olap/signal/signal.go | 16 +- .../services/controllers/task/controller.go | 33 +- .../controllers/task/durable_callbacks.go | 2 +- .../controllers/task/trigger/trigger.go | 4 +- internal/services/dispatcher/dispatcher.go | 25 +- internal/services/dispatcher/dispatcher_v1.go | 4 +- internal/services/dispatcher/server.go | 18 +- internal/services/health/health.go | 8 +- internal/services/ingestor/ingestor.go | 15 +- internal/services/ingestor/server_v1.go | 5 +- internal/services/scheduler/v1/scheduler.go | 45 +- .../ticker/schedule_workflow_v1_test.go | 3 +- pkg/config/database/config.go | 5 + pkg/config/loader/loader.go | 123 ++++- pkg/config/loader/loader_pubsub_test.go | 88 ++++ pkg/config/server/server.go | 41 ++ pkg/repository/mq.go | 14 + pkg/repository/multiplexer.go | 5 +- 43 files changed, 2003 insertions(+), 469 deletions(-) delete mode 100644 internal/msgqueue/postgres/exchange.go create mode 100644 internal/msgqueue/postgres/pubsub.go create mode 100644 internal/msgqueue/postgres/pubsub_test.go create mode 100644 internal/msgqueue/pubsub.go create mode 100644 internal/msgqueue/pubsub_test.go create mode 100644 internal/msgqueue/rabbitmq/pubsub.go create mode 100644 internal/msgqueue/rabbitmq/pubsub_test.go create mode 100644 pkg/config/loader/loader_pubsub_test.go diff --git a/cmd/hatchet-engine/engine/run.go b/cmd/hatchet-engine/engine/run.go index 139e6feb7b..52bb95c4f4 100644 --- a/cmd/hatchet-engine/engine/run.go +++ b/cmd/hatchet-engine/engine/run.go @@ -160,7 +160,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. var healthCleanup func() error healthProbes := sc.HasService("health") if healthProbes { - h = health.New(sc.V1.Health(), sc.MessageQueueV1, sc.Version, l) + h = health.New(sc.V1.Health(), sc.MessageQueueV1, sc.PubSubV1, sc.Version, l) cleanupHealth, err := h.Start(sc.Runtime.HealthcheckPort) if err != nil { return fmt.Errorf("could not start health: %w", err) @@ -197,6 +197,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. sv1, err := schedulerv1.New( schedulerv1.WithAlerter(sc.Alerter), schedulerv1.WithMessageQueue(sc.MessageQueueV1), + schedulerv1.WithPubSub(sc.PubSubV1), schedulerv1.WithRepository(sc.V1), schedulerv1.WithLogger(sc.Logger), schedulerv1.WithPartition(p), @@ -249,6 +250,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. tasks, err := task.New( task.WithAlerter(sc.Alerter), task.WithMessageQueue(sc.MessageQueueV1), + task.WithPubSub(sc.PubSubV1), task.WithV1Repository(sc.V1), task.WithLogger(sc.Logger), task.WithPartition(p), @@ -281,6 +283,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. olap, err := olap.New( olap.WithAlerter(sc.Alerter), olap.WithMessageQueue(sc.MessageQueueV1), + olap.WithPubSub(sc.PubSubV1), olap.WithRepository(sc.V1), olap.WithLogger(sc.Logger), olap.WithPartition(p), @@ -343,6 +346,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. d, err := dispatcher.New( dispatcher.WithAlerter(sc.Alerter), dispatcher.WithMessageQueueV1(sc.MessageQueueV1), + dispatcher.WithPubSub(sc.PubSubV1), dispatcher.WithRepositoryV1(sc.V1), dispatcher.WithLogger(sc.Logger), dispatcher.WithCache(cacheInstance), @@ -367,6 +371,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. // create the event ingestor ei, err := ingestor.NewIngestor( ingestor.WithMessageQueueV1(sc.MessageQueueV1), + ingestor.WithPubSub(sc.PubSubV1), ingestor.WithRepositoryV1(sc.V1), ingestor.WithLogIngestionEnabled(sc.Runtime.LogIngestionEnabled), ingestor.WithLocalScheduler(localScheduler), @@ -385,6 +390,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. adminSvc, err := admin.NewAdminService( admin.WithRepositoryV1(sc.V1), admin.WithMessageQueueV1(sc.MessageQueueV1), + admin.WithPubSub(sc.PubSubV1), admin.WithLocalScheduler(localScheduler), admin.WithLocalDispatcher(d), admin.WithOptimisticSchedulingEnabled(sc.Runtime.OptimisticSchedulingEnabled), @@ -400,6 +406,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. adminv1Svc, err := adminv1.NewAdminService( adminv1.WithRepository(sc.V1), adminv1.WithMessageQueue(sc.MessageQueueV1), + adminv1.WithPubSub(sc.PubSubV1), adminv1.WithAnalytics(sc.Analytics), adminv1.WithLocalScheduler(localScheduler), adminv1.WithLocalDispatcher(d), @@ -565,7 +572,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. var healthCleanup func() error if healthProbes { - h = health.New(sc.V1.Health(), sc.MessageQueueV1, sc.Version, l) + h = health.New(sc.V1.Health(), sc.MessageQueueV1, sc.PubSubV1, sc.Version, l) clean, err := h.Start(sc.Runtime.HealthcheckPort) @@ -631,6 +638,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. tasks, err := task.New( task.WithAlerter(sc.Alerter), task.WithMessageQueue(sc.MessageQueueV1), + task.WithPubSub(sc.PubSubV1), task.WithV1Repository(sc.V1), task.WithLogger(sc.Logger), task.WithPartition(p), @@ -666,6 +674,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. olap, err := olap.New( olap.WithAlerter(sc.Alerter), olap.WithMessageQueue(sc.MessageQueueV1), + olap.WithPubSub(sc.PubSubV1), olap.WithRepository(sc.V1), olap.WithLogger(sc.Logger), olap.WithPartition(p), @@ -748,6 +757,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. sv1, err := schedulerv1.New( schedulerv1.WithAlerter(sc.Alerter), schedulerv1.WithMessageQueue(sc.MessageQueueV1), + schedulerv1.WithPubSub(sc.PubSubV1), schedulerv1.WithRepository(sc.V1), schedulerv1.WithLogger(sc.Logger), schedulerv1.WithPartition(p), @@ -781,6 +791,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. d, err := dispatcher.New( dispatcher.WithAlerter(sc.Alerter), dispatcher.WithMessageQueueV1(sc.MessageQueueV1), + dispatcher.WithPubSub(sc.PubSubV1), dispatcher.WithRepositoryV1(sc.V1), dispatcher.WithLogger(sc.Logger), dispatcher.WithCache(cacheInstance), @@ -806,6 +817,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. // create the event ingestor ei, err := ingestor.NewIngestor( ingestor.WithMessageQueueV1(sc.MessageQueueV1), + ingestor.WithPubSub(sc.PubSubV1), ingestor.WithRepositoryV1(sc.V1), ingestor.WithLogIngestionEnabled(sc.Runtime.LogIngestionEnabled), ingestor.WithLocalScheduler(localScheduler), @@ -824,6 +836,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. adminSvc, err := admin.NewAdminService( admin.WithRepositoryV1(sc.V1), admin.WithMessageQueueV1(sc.MessageQueueV1), + admin.WithPubSub(sc.PubSubV1), admin.WithLocalScheduler(localScheduler), admin.WithLocalDispatcher(d), admin.WithOptimisticSchedulingEnabled(sc.Runtime.OptimisticSchedulingEnabled), @@ -840,6 +853,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. adminv1Svc, err := adminv1.NewAdminService( adminv1.WithRepository(sc.V1), adminv1.WithMessageQueue(sc.MessageQueueV1), + adminv1.WithPubSub(sc.PubSubV1), adminv1.WithAnalytics(sc.Analytics), adminv1.WithLocalScheduler(localScheduler), adminv1.WithLocalDispatcher(d), diff --git a/frontend/docs/pages/self-hosting/configuration-options.mdx b/frontend/docs/pages/self-hosting/configuration-options.mdx index 37ad08d197..a5107d3ec0 100644 --- a/frontend/docs/pages/self-hosting/configuration-options.mdx +++ b/frontend/docs/pages/self-hosting/configuration-options.mdx @@ -291,18 +291,24 @@ Variables marked with ⚠️ are conditionally required when specific features a ## Task Queue Configuration -| Variable | Description | Default Value | -| --------------------------------------------------- | -------------------------------------------------------------- | ------------- | -| `SERVER_MSGQUEUE_KIND` | Message queue kind (`rabbitmq` or `postgres`) | `rabbitmq` | -| `SERVER_MSGQUEUE_RABBITMQ_URL` | RabbitMQ URL | | -| `SERVER_MSGQUEUE_RABBITMQ_QOS` | RabbitMQ QoS (prefetch count) | `100` | -| `SERVER_MSGQUEUE_RABBITMQ_MAX_PUB_CHANS` | Max RabbitMQ publish channels | `20` | -| `SERVER_MSGQUEUE_RABBITMQ_MAX_SUB_CHANS` | Max RabbitMQ subscribe channels | `100` | -| `SERVER_MSGQUEUE_RABBITMQ_COMPRESSION_ENABLED` | Enable gzip compression of messages | `false` | -| `SERVER_MSGQUEUE_RABBITMQ_COMPRESSION_THRESHOLD` | Byte size above which messages are compressed | `5120` | -| `SERVER_MSGQUEUE_RABBITMQ_ENABLE_MESSAGE_REJECTION` | Reject (rather than requeue) messages past the max death count | `false` | -| `SERVER_MSGQUEUE_RABBITMQ_MAX_DEATH_COUNT` | Max redeliveries before a message is dead-lettered | `1000` | -| `SERVER_SINGLE_QUEUE_LIMIT` | Single queue limit | `100` | +| Variable | Description | Default Value | +| --------------------------------------------------- | ---------------------------------------------------------------------------------------------------------- | ------------- | +| `SERVER_MSGQUEUE_KIND` | Message queue kind (`rabbitmq` or `postgres`) | `rabbitmq` | +| `SERVER_MSGQUEUE_RABBITMQ_URL` | RabbitMQ URL | | +| `SERVER_MSGQUEUE_RABBITMQ_QOS` | RabbitMQ QoS (prefetch count) | `100` | +| `SERVER_MSGQUEUE_RABBITMQ_MAX_PUB_CHANS` | Max RabbitMQ publish channels | `20` | +| `SERVER_MSGQUEUE_RABBITMQ_MAX_SUB_CHANS` | Max RabbitMQ subscribe channels | `100` | +| `SERVER_MSGQUEUE_RABBITMQ_COMPRESSION_ENABLED` | Enable gzip compression of messages | `false` | +| `SERVER_MSGQUEUE_RABBITMQ_COMPRESSION_THRESHOLD` | Byte size above which messages are compressed | `5120` | +| `SERVER_MSGQUEUE_RABBITMQ_ENABLE_MESSAGE_REJECTION` | Reject (rather than requeue) messages past the max death count | `false` | +| `SERVER_MSGQUEUE_RABBITMQ_MAX_DEATH_COUNT` | Max redeliveries before a message is dead-lettered | `1000` | +| `SERVER_MSGQUEUE_PUBSUB_KIND` | Pub/sub kind for best-effort notifications (inherits `SERVER_MSGQUEUE_KIND` when unset) | | +| `SERVER_MSGQUEUE_PUBSUB_RABBITMQ_URL` | Pub/sub RabbitMQ URL (inherits `SERVER_MSGQUEUE_RABBITMQ_URL` when unset; always uses its own connections) | | +| `SERVER_MSGQUEUE_PUBSUB_RABBITMQ_MAX_PUB_CHANS` | Pub/sub RabbitMQ max publish channels | `10` | +| `SERVER_MSGQUEUE_PUBSUB_RABBITMQ_MAX_SUB_CHANS` | Pub/sub RabbitMQ max subscribe channels | `20` | +| `SERVER_MSGQUEUE_PUBSUB_POSTGRES_MAX_CONNS` | Pub/sub Postgres pool max connections (pool is built from the direct `DATABASE_URL`) | `5` | +| `SERVER_MSGQUEUE_PUBSUB_POSTGRES_MIN_CONNS` | Pub/sub Postgres pool min connections | `1` | +| `SERVER_SINGLE_QUEUE_LIMIT` | Single queue limit | `100` | ### Message Queue Buffers diff --git a/internal/msgqueue/mq_pub_buffer_test.go b/internal/msgqueue/mq_pub_buffer_test.go index 28d3ef74ae..91d9617ffa 100644 --- a/internal/msgqueue/mq_pub_buffer_test.go +++ b/internal/msgqueue/mq_pub_buffer_test.go @@ -5,8 +5,6 @@ import ( "sync/atomic" "testing" "time" - - "github.com/google/uuid" ) // mockMessageQueue satisfies MessageQueue for tests. Only SendMessage is wired up. @@ -27,8 +25,7 @@ func (m *mockMessageQueue) SendMessage(ctx context.Context, q Queue, msg *Messag func (m *mockMessageQueue) Subscribe(_ Queue, _ AckHook, _ AckHook) (func() error, error) { return func() error { return nil }, nil } -func (m *mockMessageQueue) RegisterTenant(_ context.Context, _ uuid.UUID) error { return nil } -func (m *mockMessageQueue) IsReady() bool { return true } +func (m *mockMessageQueue) IsReady() bool { return true } // TestPubBufferFlushesWhenFull verifies that when the channel is at capacity the // capacityRelease mechanism breaks the post-flush interval wait and triggers an diff --git a/internal/msgqueue/mq_sub_buffer.go b/internal/msgqueue/mq_sub_buffer.go index cb79056f6d..cc7603715e 100644 --- a/internal/msgqueue/mq_sub_buffer.go +++ b/internal/msgqueue/mq_sub_buffer.go @@ -39,11 +39,14 @@ const ( PreAck SubBufferKind = "preAck" ) +// SubscribeFunc subscribes to a message source with pre-ack and post-ack +// hooks, returning a cleanup function. +type SubscribeFunc func(preAck AckHook, postAck AckHook) (func() error, error) + // MQSubBuffer buffers messages coming out of the task queue, groups them by tenantId and msgId, and then flushes them // to the task handler as necessary. type MQSubBuffer struct { - queue Queue - mq MessageQueue + sub SubscribeFunc // buffers is keyed on a composite (tenantId, msgId) and contains a buffer of messages for that tenantId and msgId. buffers syncx.Map[string, *msgIdBuffer] @@ -100,13 +103,21 @@ func defaultMQSubBufferOpts() *mqSubBufferOpts { } func NewMQSubBuffer(queue Queue, mq MessageQueue, dst DstFunc, fs ...mqSubBufferOptFunc) *MQSubBuffer { + return NewSubBufferFromSubscribe(func(preAck AckHook, postAck AckHook) (func() error, error) { + return mq.Subscribe(queue, preAck, postAck) + }, dst, fs...) +} + +// NewSubBufferFromSubscribe creates a sub buffer over an arbitrary subscribe +// function, so the buffer can sit on top of either the durable MessageQueue +// or a best-effort PubSub subscription. +func NewSubBufferFromSubscribe(sub SubscribeFunc, dst DstFunc, fs ...mqSubBufferOptFunc) *MQSubBuffer { opts := defaultMQSubBufferOpts() for _, f := range fs { f(opts) } return &MQSubBuffer{ - queue: queue, - mq: mq, + sub: sub, dst: dst, kind: opts.kind, flushInterval: opts.flushInterval, @@ -128,9 +139,9 @@ func (m *MQSubBuffer) Start() (func() error, error) { switch m.kind { case PreAck: - cleanupQueue, err = m.mq.Subscribe(m.queue, f, NoOpHook) + cleanupQueue, err = m.sub(f, NoOpHook) case PostAck: - cleanupQueue, err = m.mq.Subscribe(m.queue, NoOpHook, f) + cleanupQueue, err = m.sub(NoOpHook, f) } if err != nil { diff --git a/internal/msgqueue/msgqueue.go b/internal/msgqueue/msgqueue.go index 570e59580f..375ea9b3f7 100644 --- a/internal/msgqueue/msgqueue.go +++ b/internal/msgqueue/msgqueue.go @@ -22,13 +22,6 @@ type Queue interface { // Exclusive returns true if this queue should only be accessed by the current connection. Exclusive() bool - // FanoutExchangeKey returns which exchange the queue should be subscribed to. This is only currently relevant - // to tenant pub/sub queues. - // - // In RabbitMQ terminology, the existence of a subscriber key means that the queue is bound to a fanout - // exchange, and a new random queue is generated for each connection when connections are retried. - FanoutExchangeKey() string - // DLQ returns the queue's dead letter queue, if it exists. DLQ() Queue @@ -69,10 +62,6 @@ func (s staticQueue) Exclusive() bool { return false } -func (s staticQueue) FanoutExchangeKey() string { - return "" -} - func (s staticQueue) DLQ() Queue { name := fmt.Sprintf("%s_dlq", s) @@ -138,10 +127,6 @@ func (d dispatcherQueue) Exclusive() bool { return true } -func (d dispatcherQueue) FanoutExchangeKey() string { - return "" -} - func (d dispatcherQueue) DLQ() Queue { return dlq{ staticQueue: DISPATCHER_DEAD_LETTER_QUEUE, @@ -166,73 +151,6 @@ func QueueTypeFromDispatcherID(d uuid.UUID) dispatcherQueue { return dispatcherQueue(d.String() + "_dispatcher_v1") } -type consumerQueue string - -func (s consumerQueue) Name() string { - return string(s) -} - -func (n consumerQueue) Durable() bool { - return false -} - -func (n consumerQueue) AutoDeleted() bool { - return true -} - -func (n consumerQueue) Exclusive() bool { - return true -} - -func (n consumerQueue) FanoutExchangeKey() string { - return "" -} - -func (n consumerQueue) DLQ() Queue { - return nil -} - -func (n consumerQueue) IsDLQ() bool { - return false -} - -func (n consumerQueue) IsAutoDLQ() bool { - return false -} - -func (n consumerQueue) IsExpirable() bool { - // since exclusive and auto-deleted, it's not expirable - return false -} - -const ( - Scheduler = "scheduler" -) - -func QueueTypeFromPartitionIDAndController(p, controller string) consumerQueue { - return consumerQueue(fmt.Sprintf("%s_%s_v1", p, controller)) -} - -type fanoutQueue struct { - consumerQueue -} - -// The fanout exchange key for a consumer is the name of the consumer queue. -func (f fanoutQueue) FanoutExchangeKey() string { - return f.Name() -} - -func TenantEventConsumerQueue(t uuid.UUID) fanoutQueue { - // generate a unique queue name for the tenant - return fanoutQueue{ - consumerQueue: consumerQueue(GetTenantExchangeName(t)), - } -} - -func GetTenantExchangeName(t uuid.UUID) string { - return t.String() + "_v1" -} - type AckHook func(task *Message) error type MessageQueue interface { @@ -249,12 +167,6 @@ type MessageQueue interface { // subscription is no longer needed. Subscribe(queue Queue, preAck AckHook, postAck AckHook) (func() error, error) - // RegisterTenant registers a new pub/sub mechanism for a tenant. This should be called when a - // new tenant is created. If this is not called, implementors should ensure that there's a check - // on the first message to a tenant to ensure that the tenant is registered, and store the tenant - // in an LRU cache which lives in-memory. - RegisterTenant(ctx context.Context, tenantId uuid.UUID) error - // IsReady returns true if the task queue is ready to accept tasks. IsReady() bool } diff --git a/internal/msgqueue/postgres/exchange.go b/internal/msgqueue/postgres/exchange.go deleted file mode 100644 index 88e91eb3e5..0000000000 --- a/internal/msgqueue/postgres/exchange.go +++ /dev/null @@ -1,59 +0,0 @@ -package postgres - -import ( - "context" - "encoding/json" - - "golang.org/x/sync/errgroup" - - "github.com/google/uuid" - - "github.com/hatchet-dev/hatchet/internal/msgqueue" -) - -func (p *PostgresMessageQueue) addTenantExchangeMessage(ctx context.Context, q msgqueue.Queue, msg *msgqueue.Message) error { - tenantId := msg.TenantID - - if tenantId == uuid.Nil { - return nil - } - - err := p.RegisterTenant(ctx, tenantId) - - if err != nil { - p.l.Error().Ctx(ctx).Msgf("error registering tenant exchange: %v", err) - return err - } - - queueName := msgqueue.GetTenantExchangeName(tenantId) - - // if the queue name does not equal the tenant exchange name, publish the message to the queue - if queueName != q.Name() { - return p.pubNonDurableMessages(ctx, queueName, msg) - } - - return nil -} - -func (p *PostgresMessageQueue) pubNonDurableMessages(ctx context.Context, queueName string, msg *msgqueue.Message) error { - eg := errgroup.Group{} - - for _, payload := range msg.Payloads { - msgCp := msg - msgCp.Payloads = [][]byte{payload} - - msgBytes, err := json.Marshal(msgCp) - - if err == nil { - eg.Go(func() error { - // Notify will automatically fall back to database storage if the - // wrapped message exceeds pg_notify's 8KB limit - return p.repo.Notify(ctx, queueName, string(msgBytes)) - }) - } else { - p.l.Error().Ctx(ctx).Err(err).Msg("error marshalling message") - } - } - - return eg.Wait() -} diff --git a/internal/msgqueue/postgres/msgqueue.go b/internal/msgqueue/postgres/msgqueue.go index a756d8f637..ec61b84e7a 100644 --- a/internal/msgqueue/postgres/msgqueue.go +++ b/internal/msgqueue/postgres/msgqueue.go @@ -150,11 +150,7 @@ func (p *PostgresMessageQueue) addMessage(ctx context.Context, queue msgqueue.Qu return err } - if !queue.Durable() { - err = p.pubNonDurableMessages(ctx, queue.Name(), task) - } else { - err = p.repo.AddMessage(ctx, queue.Name(), msgBytes) - } + err = p.repo.AddMessage(ctx, queue.Name(), msgBytes) if err != nil { p.l.Error().Err(err).Msgf("error adding message for queue %s", queue.Name()) @@ -168,10 +164,6 @@ func (p *PostgresMessageQueue) addMessage(ctx context.Context, queue msgqueue.Qu p.l.Error().Err(err).Msgf("error notifying queue %s", queue.Name()) } - if task.TenantID != uuid.Nil { - return p.addTenantExchangeMessage(ctx, queue, task) - } - return nil } @@ -285,20 +277,6 @@ func (p *PostgresMessageQueue) Subscribe(queue msgqueue.Queue, preAck msgqueue.A // start the listener go func() { err := p.repo.Listen(subscribeCtx, queue.Name(), func(ctx context.Context, notification *v1.PubSubMessage) error { - // if this is not a durable queue, and the message starts with JSON '{', then we process the message directly - if !queue.Durable() && len(notification.Payload) >= 1 && notification.Payload[0] == '{' { - var task msgqueue.Message - - err := json.Unmarshal([]byte(notification.Payload), &task) - - if err != nil { - p.l.Error().Err(err).Msg("error unmarshalling message") - return err - } - - return doTask(task, nil) - } - newMsgCh <- struct{}{} return nil }) @@ -334,10 +312,6 @@ func (p *PostgresMessageQueue) Subscribe(queue msgqueue.Queue, preAck msgqueue.A }, nil } -func (p *PostgresMessageQueue) RegisterTenant(ctx context.Context, tenantId uuid.UUID) error { - return p.upsertQueue(ctx, msgqueue.TenantEventConsumerQueue(tenantId)) -} - func (p *PostgresMessageQueue) IsReady() bool { return true } @@ -349,13 +323,6 @@ func (p *PostgresMessageQueue) upsertQueue(ctx context.Context, queue msgqueue.Q exclusive := queue.Exclusive() - // If the queue is a fanout exchange, then it is not exclusive. This is different from the RabbitMQ - // implementation, where a fanout exchange will map to an exclusively bound queue which has a random - // suffix appended to the queue name. In this implementation, there is no concept of an exchange. - if queue.FanoutExchangeKey() != "" { - exclusive = false - } - var consumer *string if exclusive { diff --git a/internal/msgqueue/postgres/pubsub.go b/internal/msgqueue/postgres/pubsub.go new file mode 100644 index 0000000000..0347539024 --- /dev/null +++ b/internal/msgqueue/postgres/pubsub.go @@ -0,0 +1,257 @@ +package postgres + +import ( + "context" + "encoding/json" + "time" + + "github.com/rs/zerolog" + "golang.org/x/sync/errgroup" + + "github.com/hatchet-dev/hatchet/internal/cache" + "github.com/hatchet-dev/hatchet/internal/msgqueue" + "github.com/hatchet-dev/hatchet/internal/queueutils" + "github.com/hatchet-dev/hatchet/pkg/logger" + v1 "github.com/hatchet-dev/hatchet/pkg/repository" +) + +// fallbackPollBatchSize is the max number of >8KB fallback rows drained per poll. +const fallbackPollBatchSize = 100 + +// PubSub implements msgqueue.PubSub over Postgres LISTEN/NOTIFY. Topic names +// match the legacy non-durable queue / NOTIFY names, and all traffic +// multiplexes over the single "hatchet_listener" channel, so mixed-version +// fleets interoperate. +// +// INVARIANT: the repo passed to NewPubSub must be built on a dedicated pool +// from the direct (non-pgbouncer) database URL — never the shared repository +// pool. Pub can be called from within durable-write paths, and LISTEN does not +// survive transaction pooling. +type PubSub struct { + repo v1.MessageQueueRepository + l *zerolog.Logger + + // ttlCache dedupes queue-row upserts, which are needed so >8KB payloads can + // fall back to durable rows + ttlCache *cache.TTLCache[string, bool] +} + +type PubSubOpt func(*PubSubOpts) + +type PubSubOpts struct { + l *zerolog.Logger +} + +func defaultPubSubOpts() *PubSubOpts { + l := logger.NewDefaultLogger("postgres-pubsub") + + return &PubSubOpts{ + l: &l, + } +} + +func WithPubSubLogger(l *zerolog.Logger) PubSubOpt { + return func(opts *PubSubOpts) { + opts.l = l + } +} + +// NewPubSub creates a new Postgres-backed PubSub over the given message queue +// repository. +func NewPubSub(repo v1.MessageQueueRepository, fs ...PubSubOpt) (func() error, *PubSub, error) { + opts := defaultPubSubOpts() + + for _, f := range fs { + f(opts) + } + + c := cache.NewTTL[string, bool]() + + p := &PubSub{ + repo: repo, + l: opts.l, + ttlCache: c, + } + + return func() error { + c.Stop() + return nil + }, p, nil +} + +func (p *PubSub) IsReady() bool { + return true +} + +// Pub publishes a message to the topic via NOTIFY. Payloads whose wrapped +// message exceeds pg_notify's 8KB limit fall back to a short-lived message +// queue row (see MessageQueueRepository.Notify), which subscribers drain with +// a ~1s poll — task-stream-event payloads routinely exceed 8KB and stream +// delivery must not regress. +func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Message) error { + // upsert the queue row so >8KB fallback rows have a queue to land on + err := p.ensureQueue(ctx, topic) + + if err != nil { + return err + } + + eg := errgroup.Group{} + + for _, payload := range msg.Payloads { + msgCp := *msg + msgCp.Payloads = [][]byte{payload} + + msgBytes, err := json.Marshal(&msgCp) + + if err == nil { + eg.Go(func() error { + // Notify will automatically fall back to database storage if the + // wrapped message exceeds pg_notify's 8KB limit + return p.repo.Notify(ctx, topic.Name(), string(msgBytes)) + }) + } else { + p.l.Error().Ctx(ctx).Err(err).Msg("error marshalling message") + } + } + + return eg.Wait() +} + +// Sub subscribes to a topic. Inline NOTIFY payloads are handled directly; +// non-JSON notifications and a 1s ticker wake a poll that drains >8KB fallback +// rows. Delivery is at-most-once: handler errors are logged, never redelivered. +func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.AckHook) (func() error, error) { + err := p.ensureQueue(context.Background(), topic) + + if err != nil { + return nil, err + } + + subscribeCtx, cancel := context.WithCancel(context.Background()) + + // update the lastActive time on the queue row every 60 seconds so the + // auto-delete reaper doesn't collect it while we're subscribed + go func() { + ticker := time.NewTicker(60 * time.Second) + defer ticker.Stop() + + for { + select { + case <-subscribeCtx.Done(): + return + case <-ticker.C: + err := p.repo.UpdateQueueLastActive(subscribeCtx, topic.Name()) + + if err != nil { + p.l.Error().Err(err).Msg("error updating lastActive time") + } + } + } + }() + + 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) + } + } + + // poll for >8KB fallback rows + op := queueutils.NewOperationPool(p.l, 60*time.Second, "postgres-pubsub", queueutils.OpMethod(func(ctx context.Context, _ string) (bool, error) { + messages, err := p.repo.ReadMessages(subscribeCtx, topic.Name(), fallbackPollBatchSize) + + if err != nil { + p.l.Error().Err(err).Msg("error reading fallback messages") + return false, err + } + + for _, message := range messages { + var task msgqueue.Message + + if err := json.Unmarshal(message.Payload, &task); err != nil { + p.l.Error().Err(err).Msg("error unmarshalling fallback message") + continue + } + + handleMsg(&task, &message.ID) + } + + return len(messages) == fallbackPollBatchSize, nil + })) + + newMsgCh := make(chan struct{}, 1) + + go func() { + err := p.repo.Listen(subscribeCtx, topic.Name(), func(ctx context.Context, notification *v1.PubSubMessage) error { + // messages small enough for pg_notify arrive inline as JSON + if len(notification.Payload) >= 1 && notification.Payload[0] == '{' { + var task msgqueue.Message + + if err := json.Unmarshal(notification.Payload, &task); err != nil { + p.l.Error().Err(err).Msg("error unmarshalling message") + return err + } + + handleMsg(&task, nil) + return nil + } + + // anything else is a wake-up signal to drain fallback rows + select { + case newMsgCh <- struct{}{}: + default: + } + + return nil + }) + + if err != nil && subscribeCtx.Err() == nil { + p.l.Error().Err(err).Msg("error listening for new messages") + } + }() + + ticker := time.NewTicker(time.Second) + + go func() { + defer ticker.Stop() + + for { + select { + case <-subscribeCtx.Done(): + return + case <-ticker.C: + op.RunOrContinue(topic.Name()) + case <-newMsgCh: + op.RunOrContinue(topic.Name()) + } + } + }() + + return func() error { + cancel() + return nil + }, nil +} + +func (p *PubSub) ensureQueue(ctx context.Context, topic msgqueue.Topic) error { + if valid, exists := p.ttlCache.Get(topic.Name()); valid && exists { + return nil + } + + err := p.repo.BindQueue(ctx, topic.Name(), false, true, false, nil) + + if err != nil { + p.l.Error().Err(err).Msg("error binding queue") + return err + } + + p.ttlCache.Set(topic.Name(), true, time.Second*15) + + return nil +} diff --git a/internal/msgqueue/postgres/pubsub_test.go b/internal/msgqueue/postgres/pubsub_test.go new file mode 100644 index 0000000000..ff3fb89858 --- /dev/null +++ b/internal/msgqueue/postgres/pubsub_test.go @@ -0,0 +1,280 @@ +//go:build integration + +package postgres + +import ( + "context" + "crypto/rand" + "encoding/json" + "os" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/hatchet-dev/hatchet/internal/msgqueue" + "github.com/hatchet-dev/hatchet/pkg/logger" + v1 "github.com/hatchet-dev/hatchet/pkg/repository" +) + +const defaultTestDatabaseURL = "postgresql://hatchet:hatchet@127.0.0.1:5431/hatchet" + +func testDatabaseURL() string { + if v := os.Getenv("DATABASE_URL"); v != "" { + return v + } + + return defaultTestDatabaseURL +} + +func newTestPool(t *testing.T) *pgxpool.Pool { + t.Helper() + + config, err := pgxpool.ParseConfig(testDatabaseURL()) + require.NoError(t, err) + + config.MaxConns = 5 + config.MinConns = 1 + + pool, err := pgxpool.NewWithConfig(context.Background(), config) + require.NoError(t, err) + + t.Cleanup(pool.Close) + + return pool +} + +func newTestPubSub(t *testing.T) (*PubSub, *pgxpool.Pool) { + t.Helper() + + pool := newTestPool(t) + + l := logger.NewDefaultLogger("postgres-pubsub-test") + + repo, cleanupRepo := v1.NewMessageQueueRepositoryWithPool(&l, pool) + + cleanup, ps, err := NewPubSub(repo, WithPubSubLogger(&l)) + require.NoError(t, err) + + t.Cleanup(func() { + if err := cleanup(); err != nil { + t.Errorf("error cleaning up pubsub: %v", err) + } + if err := cleanupRepo(); err != nil { + t.Errorf("error cleaning up repo: %v", err) + } + }) + + return ps, pool +} + +func TestPostgresPubSubNotifyRoundtrip(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + ps, _ := newTestPubSub(t) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + wg := &sync.WaitGroup{} + wg.Add(1) + + msg, err := msgqueue.NewTenantMessage(tenantId, "task-completed", true, false, map[string]interface{}{"key": "value"}) + require.NoError(t, err) + + cleanupSub, err := ps.Sub(topic, func(received *msgqueue.Message) error { + defer wg.Done() + assert.Equal(t, msg.ID, received.ID) + return nil + }) + require.NoError(t, err) + + // give the LISTEN connection time to establish + time.Sleep(1 * time.Second) + + require.NoError(t, ps.Pub(ctx, topic, msg)) + + wg.Wait() + + require.NoError(t, cleanupSub()) +} + +func TestPostgresPubSubLargePayloadFallback(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + ps, _ := newTestPubSub(t) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + // well beyond pg_notify's 8000-byte limit: must take the fallback-row path + payload := make([]byte, 32*1024) + _, err := rand.Read(payload) + require.NoError(t, err) + + msg, err := msgqueue.NewTenantMessage(tenantId, "task-stream-event", true, false, map[string]interface{}{"data": payload}) + require.NoError(t, err) + + originalPayload := make([]byte, len(msg.Payloads[0])) + copy(originalPayload, msg.Payloads[0]) + + received := make(chan *msgqueue.Message, 1) + + cleanupSub, err := ps.Sub(topic, func(m *msgqueue.Message) error { + received <- m + return nil + }) + require.NoError(t, err) + + time.Sleep(1 * time.Second) + + start := time.Now() + require.NoError(t, ps.Pub(ctx, topic, msg)) + + select { + case m := <-received: + assert.Equal(t, msg.ID, m.ID) + require.Len(t, m.Payloads, 1) + assert.Equal(t, originalPayload, m.Payloads[0]) + // the fallback poll runs at 1s intervals; allow some slack + assert.Less(t, time.Since(start), 5*time.Second, "fallback row should be drained within a few poll intervals") + case <-ctx.Done(): + t.Fatal("timed out waiting for >8KB fallback delivery") + } + + require.NoError(t, cleanupSub()) +} + +func TestPostgresPubSubTwoSubscribers(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + ps, _ := newTestPubSub(t) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + wg := &sync.WaitGroup{} + wg.Add(2) + + msg, err := msgqueue.NewTenantMessage(tenantId, "task-completed", true, false, map[string]interface{}{"key": "value"}) + require.NoError(t, err) + + handler := func(received *msgqueue.Message) error { + defer wg.Done() + assert.Equal(t, msg.ID, received.ID) + return nil + } + + cleanupSub1, err := ps.Sub(topic, handler) + require.NoError(t, err) + + cleanupSub2, err := ps.Sub(topic, handler) + require.NoError(t, err) + + time.Sleep(1 * time.Second) + + require.NoError(t, ps.Pub(ctx, topic, msg)) + + wg.Wait() + + require.NoError(t, cleanupSub1()) + require.NoError(t, cleanupSub2()) +} + +// TestPostgresPubSubDisjointPools asserts the core invariant: a full pub/sub +// cycle never acquires from any pool other than the PubSub's own dedicated pool. +func TestPostgresPubSubDisjointPools(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + ps, _ := newTestPubSub(t) + + // stands in for the shared repository pool + sharedPool := newTestPool(t) + + // let min conns settle before snapshotting + time.Sleep(500 * time.Millisecond) + + acquiresBefore := sharedPool.Stat().AcquireCount() + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + wg := &sync.WaitGroup{} + wg.Add(1) + + msg, err := msgqueue.NewTenantMessage(tenantId, "task-completed", true, false, map[string]interface{}{"key": "value"}) + require.NoError(t, err) + + cleanupSub, err := ps.Sub(topic, func(received *msgqueue.Message) error { + defer wg.Done() + return nil + }) + require.NoError(t, err) + + time.Sleep(1 * time.Second) + + require.NoError(t, ps.Pub(ctx, topic, msg)) + + wg.Wait() + + require.NoError(t, cleanupSub()) + + assert.Equal(t, acquiresBefore, sharedPool.Stat().AcquireCount(), "pub/sub cycle must not acquire from the shared pool") +} + +// TestPostgresPubSubMixedVersionInterop simulates an old-version engine +// mirroring a tenant message via the legacy repo.Notify path on a different +// pool: the new PubSub subscriber must receive it. +func TestPostgresPubSubMixedVersionInterop(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + ps, _ := newTestPubSub(t) + + // the "old engine" publishes through its own repository on a separate pool + oldEnginePool := newTestPool(t) + l := logger.NewDefaultLogger("old-engine") + oldRepo, cleanupOldRepo := v1.NewMessageQueueRepositoryWithPool(&l, oldEnginePool) + t.Cleanup(func() { + if err := cleanupOldRepo(); err != nil { + t.Errorf("error cleaning up old repo: %v", err) + } + }) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + wg := &sync.WaitGroup{} + wg.Add(1) + + msg, err := msgqueue.NewTenantMessage(tenantId, "task-completed", true, false, map[string]interface{}{"key": "value"}) + require.NoError(t, err) + + cleanupSub, err := ps.Sub(topic, func(received *msgqueue.Message) error { + defer wg.Done() + assert.Equal(t, msg.ID, received.ID) + return nil + }) + require.NoError(t, err) + + time.Sleep(1 * time.Second) + + // legacy pubNonDurableMessages wire format: the marshalled Message as the + // NOTIFY payload on the "_v1" channel + msgBytes, err := json.Marshal(msg) + require.NoError(t, err) + + require.NoError(t, oldRepo.Notify(ctx, topic.Name(), string(msgBytes))) + + wg.Wait() + + require.NoError(t, cleanupSub()) +} diff --git a/internal/msgqueue/pubsub.go b/internal/msgqueue/pubsub.go new file mode 100644 index 0000000000..5267213aed --- /dev/null +++ b/internal/msgqueue/pubsub.go @@ -0,0 +1,102 @@ +package msgqueue + +import ( + "context" + "fmt" + + "github.com/google/uuid" +) + +// Topic identifies a best-effort pub/sub destination. +type Topic struct { + name string + tenantStream bool +} + +func (t Topic) Name() string { + return t.name +} + +// IsTenantStream returns true if this topic feeds the dispatcher's per-tenant +// gRPC streams (as opposed to a scheduler partition wake-up). +func (t Topic) IsTenantStream() bool { + return t.tenantStream +} + +// TenantTopic feeds the dispatcher's gRPC streams. The wire name matches the +// legacy tenant fanout exchange / NOTIFY name ("_v1") so mixed-version +// fleets interoperate. +func TenantTopic(tenantId uuid.UUID) Topic { + return Topic{ + name: tenantId.String() + "_v1", + tenantStream: true, + } +} + +// SchedulerPartitionTopic wakes the scheduler owning a partition. The wire +// name matches the legacy consumer queue ("_scheduler_v1") so +// mixed-version fleets interoperate. +func SchedulerPartitionTopic(partitionId string) Topic { + return Topic{ + name: fmt.Sprintf("%s_scheduler_v1", partitionId), + } +} + +// PubSub is a best-effort, non-durable, at-most-once pub/sub mechanism. +// +// INVARIANT: implementations own their pooled resources (AMQP connections and +// channel pools, pgx pools) and never share them with the durable MessageQueue +// or the repository layer, since Pub can be called from within durable-write +// paths and sharing pools is a deadlock risk. +type PubSub interface { + // Pub publishes a message to the topic. Delivery is best-effort: if no + // subscriber is listening, the message is dropped. + Pub(ctx context.Context, topic Topic, msg *Message) error + + // Sub subscribes to a topic. Delivery is at-most-once with no ack + // 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) + + // IsReady returns true if the pub/sub mechanism is ready to accept messages. + IsReady() bool +} + +// gatedPubSub suppresses tenant-stream publishes when disabled, except +// task-stream-event (worker stream output must always flow — it has no other +// delivery path). Scheduler partition wake-ups are never gated. +type gatedPubSub struct { + inner PubSub + disableTenantStreamPubs bool +} + +// NewGatedPubSub wraps a PubSub so that tenant-stream publishes are dropped +// when disableTenantStreamPubs is true, except task-stream-event messages. +// Scheduler partition wake-ups always pass through. +func NewGatedPubSub(inner PubSub, disableTenantStreamPubs bool) PubSub { + if !disableTenantStreamPubs { + return inner + } + + return &gatedPubSub{ + inner: inner, + disableTenantStreamPubs: disableTenantStreamPubs, + } +} + +func (g *gatedPubSub) Pub(ctx context.Context, topic Topic, msg *Message) error { + if g.disableTenantStreamPubs && topic.IsTenantStream() && msg.ID != MsgIDTaskStreamEvent { + return nil + } + + return g.inner.Pub(ctx, topic, msg) +} + +func (g *gatedPubSub) Sub(topic Topic, handler AckHook) (func() error, error) { + return g.inner.Sub(topic, handler) +} + +func (g *gatedPubSub) IsReady() bool { + return g.inner.IsReady() +} diff --git a/internal/msgqueue/pubsub_test.go b/internal/msgqueue/pubsub_test.go new file mode 100644 index 0000000000..7e956023ab --- /dev/null +++ b/internal/msgqueue/pubsub_test.go @@ -0,0 +1,73 @@ +package msgqueue + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" +) + +func TestTopicNames(t *testing.T) { + tenantId := uuid.MustParse("707d0855-80ab-4e1f-a156-f1c4546cbf52") + + tenantTopic := TenantTopic(tenantId) + assert.Equal(t, "707d0855-80ab-4e1f-a156-f1c4546cbf52_v1", tenantTopic.Name()) + assert.True(t, tenantTopic.IsTenantStream()) + + schedulerTopic := SchedulerPartitionTopic("mypartition") + assert.Equal(t, "mypartition_scheduler_v1", schedulerTopic.Name()) + assert.False(t, schedulerTopic.IsTenantStream()) +} + +type recordingPubSub struct { + pubs []string +} + +func (r *recordingPubSub) Pub(ctx context.Context, topic Topic, msg *Message) error { + r.pubs = append(r.pubs, topic.Name()+"/"+msg.ID) + return nil +} + +func (r *recordingPubSub) Sub(topic Topic, handler AckHook) (func() error, error) { + return func() error { return nil }, nil +} + +func (r *recordingPubSub) IsReady() bool { + return true +} + +func TestGatedPubSub(t *testing.T) { + tenantId := uuid.New() + + cases := []struct { + name string + disabled bool + topic Topic + msgId string + delivered bool + }{ + {"enabled tenant other", false, TenantTopic(tenantId), MsgIDTaskCompleted, true}, + {"enabled tenant stream", false, TenantTopic(tenantId), MsgIDTaskStreamEvent, true}, + {"enabled scheduler", false, SchedulerPartitionTopic("p"), MsgIDCheckTenantQueue, true}, + {"disabled tenant other", true, TenantTopic(tenantId), MsgIDTaskCompleted, false}, + {"disabled tenant stream", true, TenantTopic(tenantId), MsgIDTaskStreamEvent, true}, + {"disabled scheduler", true, SchedulerPartitionTopic("p"), MsgIDCheckTenantQueue, true}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + inner := &recordingPubSub{} + ps := NewGatedPubSub(inner, tc.disabled) + + err := ps.Pub(context.Background(), tc.topic, &Message{ID: tc.msgId, TenantID: tenantId}) + assert.NoError(t, err) + + if tc.delivered { + assert.Len(t, inner.pubs, 1) + } else { + assert.Empty(t, inner.pubs) + } + }) + } +} diff --git a/internal/msgqueue/rabbitmq/gzip.go b/internal/msgqueue/rabbitmq/gzip.go index da7cda3ca1..3ae77b14e5 100644 --- a/internal/msgqueue/rabbitmq/gzip.go +++ b/internal/msgqueue/rabbitmq/gzip.go @@ -9,6 +9,14 @@ import ( "sync" ) +// compressor holds gzip compression settings shared by the durable +// MessageQueueImpl and the best-effort PubSub (compression must agree on both +// sides for wire compatibility). +type compressor struct { + compressionEnabled bool + compressionThreshold int +} + type CompressionResult struct { Payloads [][]byte WasCompressed bool @@ -40,7 +48,7 @@ func getPayloadSize(payloads [][]byte) int { // compressPayloads compresses message payloads using gzip if they exceed the minimum size threshold. // Returns compression results including the compressed payloads and compression statistics. -func (t *MessageQueueImpl) compressPayloads(payloads [][]byte) (*CompressionResult, error) { +func (t compressor) compressPayloads(payloads [][]byte) (*CompressionResult, error) { result := &CompressionResult{ Payloads: payloads, WasCompressed: false, @@ -98,7 +106,7 @@ func (t *MessageQueueImpl) compressPayloads(payloads [][]byte) (*CompressionResu } // decompressPayloads decompresses message payloads using gzip. -func (t *MessageQueueImpl) decompressPayloads(payloads [][]byte) ([][]byte, error) { +func (t compressor) decompressPayloads(payloads [][]byte) ([][]byte, error) { if len(payloads) == 0 { return payloads, nil } diff --git a/internal/msgqueue/rabbitmq/gzip_test.go b/internal/msgqueue/rabbitmq/gzip_test.go index 60e1378a80..b85ed510e6 100644 --- a/internal/msgqueue/rabbitmq/gzip_test.go +++ b/internal/msgqueue/rabbitmq/gzip_test.go @@ -25,8 +25,8 @@ func generatePayloads(count, size int) [][]byte { return payloads } -func newMQ() *MessageQueueImpl { - return &MessageQueueImpl{ +func newMQ() compressor { + return compressor{ compressionEnabled: true, compressionThreshold: 0, } @@ -52,7 +52,7 @@ func TestCompressDecompressRoundtrip(t *testing.T) { } func TestCompressPayloadsDisabled(t *testing.T) { - mq := &MessageQueueImpl{ + mq := compressor{ compressionEnabled: false, compressionThreshold: 0, } @@ -65,7 +65,7 @@ func TestCompressPayloadsDisabled(t *testing.T) { } func TestCompressPayloadsBelowThreshold(t *testing.T) { - mq := &MessageQueueImpl{ + mq := compressor{ compressionEnabled: true, compressionThreshold: 100 * 1024, } diff --git a/internal/msgqueue/rabbitmq/pubsub.go b/internal/msgqueue/rabbitmq/pubsub.go new file mode 100644 index 0000000000..f4b3ec72d0 --- /dev/null +++ b/internal/msgqueue/rabbitmq/pubsub.go @@ -0,0 +1,473 @@ +package rabbitmq + +import ( + "context" + "encoding/json" + "fmt" + "slices" + "sync" + "time" + + lru "github.com/hashicorp/golang-lru/v2" + amqp "github.com/rabbitmq/amqp091-go" + "github.com/rs/zerolog" + + "github.com/hatchet-dev/hatchet/internal/msgqueue" + "github.com/hatchet-dev/hatchet/internal/queueutils" + "github.com/hatchet-dev/hatchet/pkg/logger" + "github.com/hatchet-dev/hatchet/pkg/random" +) + +// PubSub implements msgqueue.PubSub over RabbitMQ. Tenant topics map to the +// legacy per-tenant fanout exchanges ("_v1") and scheduler partition +// topics map to the scheduler's exclusive queue on the default exchange, so +// mixed-version fleets interoperate. +// +// INVARIANT: the PubSub owns its AMQP connections and channel pools and never +// shares them with the durable MessageQueueImpl, since Pub can be called from +// within durable-write paths. +type PubSub struct { + ctx context.Context + identity string + l *zerolog.Logger + + pubChannels *channelPool + subChannels *channelPool + + // lru cache of tenant exchanges we've already declared + exchangeCache *lru.Cache[string, bool] + + compressor + + maxPayloadSize int +} + +type PubSubOpt func(*PubSubOpts) + +type PubSubOpts struct { + l *zerolog.Logger + url string + maxPubChannels int32 + maxSubChannels int32 + compressionEnabled bool + compressionThreshold int +} + +func defaultPubSubOpts() *PubSubOpts { + l := logger.NewDefaultLogger("rabbitmq-pubsub") + + return &PubSubOpts{ + l: &l, + maxPubChannels: 10, + maxSubChannels: 20, + } +} + +func WithPubSubURL(url string) PubSubOpt { + return func(opts *PubSubOpts) { + opts.url = url + } +} + +func WithPubSubLogger(l *zerolog.Logger) PubSubOpt { + return func(opts *PubSubOpts) { + opts.l = l + } +} + +func WithPubSubMaxPubChannels(maxChannels int32) PubSubOpt { + return func(opts *PubSubOpts) { + if maxChannels > 0 { + opts.maxPubChannels = maxChannels + } + } +} + +func WithPubSubMaxSubChannels(maxChannels int32) PubSubOpt { + return func(opts *PubSubOpts) { + if maxChannels > 0 { + opts.maxSubChannels = maxChannels + } + } +} + +// WithPubSubGzip enables gzip compression of payloads. Compression settings +// must match the durable MessageQueue's settings for wire compatibility, since +// both sides publish to the same tenant exchanges in mixed-version fleets. +func WithPubSubGzip(enabled bool, threshold int) PubSubOpt { + return func(opts *PubSubOpts) { + opts.compressionEnabled = enabled + + if threshold <= 0 { + threshold = 5 * 1024 // default to 5KB + } + + opts.compressionThreshold = threshold + } +} + +// NewPubSub creates a new RabbitMQ-backed PubSub with its own connections and +// channel pools. +func NewPubSub(fs ...PubSubOpt) (func() error, *PubSub, error) { + ctx, cancel := context.WithCancel(context.Background()) + + opts := defaultPubSubOpts() + + for _, f := range fs { + f(opts) + } + + newLogger := opts.l.With().Str("service", "rabbitmq-pubsub").Logger() + opts.l = &newLogger + + pubChannelPool, err := newChannelPool(ctx, opts.l, opts.url, opts.maxPubChannels) + + if err != nil { + cancel() + return nil, nil, err + } + + subChannelPool, err := newChannelPool(ctx, opts.l, opts.url, opts.maxSubChannels) + + if err != nil { + pubChannelPool.Close() + cancel() + return nil, nil, err + } + + p := &PubSub{ + ctx: ctx, + identity: identity(), + l: opts.l, + pubChannels: pubChannelPool, + subChannels: subChannelPool, + compressor: compressor{ + compressionEnabled: opts.compressionEnabled, + compressionThreshold: opts.compressionThreshold, + }, + maxPayloadSize: 16 * 1024 * 1024, // 16 MB + } + + p.exchangeCache, _ = lru.New[string, bool](2000) //nolint:errcheck // this only returns an error if the size is less than 0 + + return func() error { + cancel() + pubChannelPool.Close() + subChannelPool.Close() + return nil + }, p, nil +} + +func (p *PubSub) IsReady() bool { + return p.pubChannels.hasActiveConnection() && p.subChannels.hasActiveConnection() +} + +// Pub publishes a message to the topic. Delivery is best-effort and +// non-persistent: if no subscriber queue is bound, the message is dropped. +func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Message) error { + // compress exactly as the durable pubMessage does: in mixed-version fleets, + // old engines' mirrored copies land on the same exchanges compressed, so + // subscribers must handle both and publishers must match settings. We work + // on a shallow copy: callers dual-publish the same *Message they hand to + // the durable queue, which may not have serialized it yet. + if len(msg.Payloads) > 0 && !msg.Compressed { + compressionResult, err := p.compressPayloads(msg.Payloads) + + if err != nil { + p.l.Error().Msgf("error compressing payloads: %v", err) + return fmt.Errorf("failed to compress payloads: %w", err) + } + + if compressionResult.WasCompressed { + msgCp := *msg + msgCp.Payloads = compressionResult.Payloads + msgCp.Compressed = true + msg = &msgCp + } + } + + var pub *amqp.Channel + var acquireErr error + + for range PUB_ACQUIRE_CHANNEL_RETRIES { + poolCh, err := p.pubChannels.Acquire(ctx) + + if err != nil { + p.l.Error().Msgf("[PubSub.Pub] cannot acquire channel: %v", err) + return err + } + + pub = poolCh.Value() + + // the channel exception may be async (after a previous pub has been sent), + // so we might have acquired a closed channel from the pool + if pub.IsClosed() { + poolCh.Destroy() + acquireErr = fmt.Errorf("channel is closed") + pub = nil + continue + } + + defer poolCh.Release() + break + } + + if pub == nil { + return acquireErr + } + + body, err := json.Marshal(msg) + + if err != nil { + p.l.Error().Msgf("error marshaling pubsub message: %v", err) + return err + } + + if len(body) > p.maxPayloadSize { + if len(msg.Payloads) == 1 { + return fmt.Errorf("message size %d bytes exceeds maximum allowed size of %d bytes", len(body), p.maxPayloadSize) + } + + // split the payloads in half and publish recursively until each chunk is + // under the max size (same strategy as the durable pubMessage) + payloadsPerChunk := max(len(msg.Payloads)/2, 1) + + for chunk := range slices.Chunk(msg.Payloads, payloadsPerChunk) { + err := p.Pub(ctx, topic, &msgqueue.Message{ + ID: msg.ID, + Payloads: chunk, + TenantID: msg.TenantID, + ImmediatelyExpire: msg.ImmediatelyExpire, + Persistent: msg.Persistent, + OtelCarrier: msg.OtelCarrier, + Retries: msg.Retries, // nolint: staticcheck + Compressed: msg.Compressed, + }) + + if err != nil { + return err + } + } + + return nil + } + + exchange := "" + routingKey := topic.Name() + + if topic.IsTenantStream() { + // tenant streams ride the per-tenant fanout exchange; declare it lazily + if _, ok := p.exchangeCache.Get(topic.Name()); !ok { + if err := p.declareExchange(pub, topic.Name()); err != nil { + p.l.Error().Msgf("error declaring exchange %s: %v", topic.Name(), err) + return err + } + + p.exchangeCache.Add(topic.Name(), true) + } + + exchange = topic.Name() + routingKey = "" + } + + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + // never persistent, expire immediately without an active consumer: this is + // an at-most-once notification path + err = pub.PublishWithContext(ctx, exchange, routingKey, false, false, amqp.Publishing{ + Body: body, + Expiration: "0", + }) + + if err != nil { + p.l.Error().Msgf("error publishing pubsub message to topic %s: %v", topic.Name(), err) + return err + } + + return nil +} + +// Sub subscribes to a topic. Delivery is at-most-once: messages are acked +// before the handler runs, and handler errors are logged, never redelivered. +func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.AckHook) (func() error, error) { + ctx, cancel := context.WithCancel(p.ctx) + + p.l.Debug().Msgf("subscribing to topic: %s", topic.Name()) + + sessionCount := 0 + wg := sync.WaitGroup{} + + innerFn := func() error { + poolCh, err := p.subChannels.Acquire(ctx) + + if err != nil { + return err + } + + sub := poolCh.Value() + + if sub.IsClosed() { + poolCh.Destroy() + return fmt.Errorf("channel is closed, destroying and retrying") + } + + defer poolCh.Release() + + // the queues are exclusive and bound to the session, so we declare them + // on each new session + queueName, err := p.declareSubQueue(sub, topic) + + if err != nil { + return err + } + + deliveries, err := sub.ConsumeWithContext(ctx, queueName, p.identity, false, true, false, false, nil) + + if err != nil { + return err + } + + for rabbitMsg := range deliveries { + wg.Add(1) + + go func(rabbitMsg amqp.Delivery, session int) { + defer wg.Done() + + // always ack: at-most-once, no redelivery + if err := sub.Ack(rabbitMsg.DeliveryTag, false); err != nil { + p.l.Error().Msgf("error acknowledging message: %v", err) + return + } + + if len(rabbitMsg.Body) == 0 { + p.l.Error().Msgf("empty message body for message: %s", rabbitMsg.MessageId) + return + } + + msg := &msgqueue.Message{} + + if err := json.Unmarshal(rabbitMsg.Body, msg); err != nil { + p.l.Error().Msgf("error unmarshalling message: %v", err) + return + } + + if msg.Compressed { + decompressedPayloads, err := p.decompressPayloads(msg.Payloads) + + if err != nil { + p.l.Error().Msgf("error decompressing payloads: %v", err) + return + } + + msg.Payloads = decompressedPayloads + } + + p.l.Debug().Msgf("(session: %d) got pubsub msg", session) + + if err := handler(msg); err != nil { + p.l.Error().Msgf("error handling pubsub message %s: %v", msg.ID, err) + } + }(rabbitMsg, sessionCount) + } + + p.l.Info().Msg("pubsub deliveries channel closed") + + return nil + } + + wg.Add(1) + go func() { + defer wg.Done() + retryCount := 0 + lastRetry := time.Now() + + for { + if ctx.Err() != nil { + return + } + + sessionCount++ + + if err := innerFn(); err != nil { + if time.Since(lastRetry) > RETRY_RESET_INTERVAL { + retryCount = 0 + } + + p.l.Error().Msgf("could not run pubsub subscriber loop (retry count %d): %v", retryCount, err) + queueutils.SleepWithExponentialBackoff(10*time.Millisecond, 5*time.Second, retryCount) + lastRetry = time.Now() + retryCount++ + continue + } + } + }() + + cleanup := func() error { + p.l.Debug().Msgf("shutting down pubsub subscriber: %s", topic.Name()) + cancel() + wg.Wait() + p.l.Debug().Msgf("successfully shut down pubsub subscriber: %s", topic.Name()) + + return nil + } + + return cleanup, nil +} + +// declareSubQueue declares the consumer queue for a topic and returns its +// name. Tenant topics get a random-suffix exclusive queue bound to the tenant +// fanout exchange; scheduler topics consume the well-known partition queue on +// the default exchange. +func (p *PubSub) declareSubQueue(ch *amqp.Channel, topic msgqueue.Topic) (string, error) { + args := make(amqp.Table) + args["x-consumer-timeout"] = 300000 // 5 minutes + + name := topic.Name() + + if topic.IsTenantStream() { + if err := p.declareExchange(ch, topic.Name()); err != nil { + return "", err + } + + suffix, err := random.Generate(8) + + if err != nil { + p.l.Error().Msgf("error generating random bytes: %v", err) + return "", err + } + + name = fmt.Sprintf("%s-%s", topic.Name(), suffix) + } + + if _, err := ch.QueueDeclare(name, false, true, true, false, args); err != nil { + p.l.Error().Msgf("cannot declare queue: %q, %v", name, err) + return "", err + } + + if topic.IsTenantStream() { + p.l.Debug().Msgf("binding queue: %s to exchange: %s", name, topic.Name()) + + if err := ch.QueueBind(name, "", topic.Name(), false, nil); err != nil { + p.l.Error().Msgf("cannot bind queue: %q, %v", name, err) + return "", err + } + } + + return name, nil +} + +// declareExchange declares a tenant fanout exchange. The args are identical to +// the legacy declareTenantExchange so declares are conflict-free in +// mixed-version fleets. +func (p *PubSub) declareExchange(ch *amqp.Channel, name string) error { + return ch.ExchangeDeclare( + name, + "fanout", + true, // durable + false, // auto-deleted + false, // not internal, accepts publishings + false, // no-wait + nil, // arguments + ) +} diff --git a/internal/msgqueue/rabbitmq/pubsub_test.go b/internal/msgqueue/rabbitmq/pubsub_test.go new file mode 100644 index 0000000000..b8cbeb6d9e --- /dev/null +++ b/internal/msgqueue/rabbitmq/pubsub_test.go @@ -0,0 +1,236 @@ +//go:build integration + +package rabbitmq + +import ( + "context" + "crypto/rand" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/hatchet-dev/hatchet/internal/msgqueue" +) + +const testAMQPURL = "amqp://user:password@localhost:5672/" + +func newTestPubSub(t *testing.T, fs ...PubSubOpt) *PubSub { + t.Helper() + + opts := append([]PubSubOpt{WithPubSubURL(testAMQPURL)}, fs...) + + cleanup, ps, err := NewPubSub(opts...) + require.NoError(t, err) + require.NotNil(t, ps) + + t.Cleanup(func() { + if err := cleanup(); err != nil { + t.Errorf("error cleaning up pubsub: %v", err) + } + }) + + return ps +} + +func TestPubSubTenantFanout(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + ps := newTestPubSub(t) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + wg := &sync.WaitGroup{} + wg.Add(2) // both subscribers must receive the message + + msg, err := msgqueue.NewTenantMessage(tenantId, "task-completed", true, false, map[string]interface{}{"key": "value"}) + require.NoError(t, err) + + handler := func(received *msgqueue.Message) error { + defer wg.Done() + assert.Equal(t, msg.ID, received.ID) + return nil + } + + cleanupSub1, err := ps.Sub(topic, handler) + require.NoError(t, err) + + cleanupSub2, err := ps.Sub(topic, handler) + require.NoError(t, err) + + // give the exclusive queues time to bind + time.Sleep(1 * time.Second) + + require.NoError(t, ps.Pub(ctx, topic, msg)) + + wg.Wait() + + require.NoError(t, cleanupSub1()) + require.NoError(t, cleanupSub2()) +} + +func TestPubSubAtMostOnce(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + ps := newTestPubSub(t) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + msg, err := msgqueue.NewTenantMessage(tenantId, "task-completed", true, false, map[string]interface{}{"key": "value"}) + require.NoError(t, err) + + // publish with no subscriber: the message must be dropped + require.NoError(t, ps.Pub(ctx, topic, msg)) + + received := make(chan struct{}, 1) + + cleanupSub, err := ps.Sub(topic, func(received2 *msgqueue.Message) error { + received <- struct{}{} + return nil + }) + require.NoError(t, err) + + select { + case <-received: + t.Fatal("message published before subscription should never be delivered") + case <-time.After(3 * time.Second): + } + + require.NoError(t, cleanupSub()) +} + +func TestPubSubSchedulerTopicRoundtrip(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + ps := newTestPubSub(t) + + topic := msgqueue.SchedulerPartitionTopic(uuid.NewString()) + + wg := &sync.WaitGroup{} + wg.Add(1) + + msg, err := msgqueue.NewTenantMessage(uuid.New(), "check-tenant-queue", true, false, map[string]interface{}{"key": "value"}) + require.NoError(t, err) + + cleanupSub, err := ps.Sub(topic, func(received *msgqueue.Message) error { + defer wg.Done() + assert.Equal(t, msg.ID, received.ID) + return nil + }) + require.NoError(t, err) + + // give the exclusive queue time to be declared + time.Sleep(1 * time.Second) + + require.NoError(t, ps.Pub(ctx, topic, msg)) + + wg.Wait() + + require.NoError(t, cleanupSub()) +} + +func TestPubSubCompressedRoundtrip(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + // threshold of 1 byte forces compression + ps := newTestPubSub(t, WithPubSubGzip(true, 1)) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + payload := make([]byte, 32*1024) + _, err := rand.Read(payload) + require.NoError(t, err) + + msg, err := msgqueue.NewTenantMessage(tenantId, "task-stream-event", true, false, map[string]interface{}{"data": payload}) + require.NoError(t, err) + + originalPayload := make([]byte, len(msg.Payloads[0])) + copy(originalPayload, msg.Payloads[0]) + + wg := &sync.WaitGroup{} + wg.Add(1) + + cleanupSub, err := ps.Sub(topic, func(received *msgqueue.Message) error { + defer wg.Done() + assert.Equal(t, msg.ID, received.ID) + require.Len(t, received.Payloads, 1) + assert.Equal(t, originalPayload, received.Payloads[0], "payload should be transparently decompressed") + return nil + }) + require.NoError(t, err) + + time.Sleep(1 * time.Second) + + require.NoError(t, ps.Pub(ctx, topic, msg)) + + wg.Wait() + + require.NoError(t, cleanupSub()) +} + +func TestPubSubReconnect(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + ps := newTestPubSub(t) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + receivedCh := make(chan string, 10) + + cleanupSub, err := ps.Sub(topic, func(received *msgqueue.Message) error { + receivedCh <- received.ID + return nil + }) + require.NoError(t, err) + + time.Sleep(1 * time.Second) + + msg1, err := msgqueue.NewTenantMessage(tenantId, "task-completed", true, false, map[string]interface{}{"key": "one"}) + require.NoError(t, err) + + require.NoError(t, ps.Pub(ctx, topic, msg1)) + + select { + case id := <-receivedCh: + assert.Equal(t, "task-completed", id) + case <-ctx.Done(): + t.Fatal("timed out waiting for first message") + } + + // force-close the subscriber connection; the channel pool reconnects and the + // subscriber loop re-declares its queue and resubscribes + require.NoError(t, ps.subChannels.getConnection().Close()) + + // wait for the pool health check (1s tick) to redial and the subscriber to recover + require.Eventually(t, func() bool { + return ps.subChannels.hasActiveConnection() + }, 15*time.Second, 250*time.Millisecond) + + time.Sleep(2 * time.Second) + + msg2, err := msgqueue.NewTenantMessage(tenantId, "task-failed", true, false, map[string]interface{}{"key": "two"}) + require.NoError(t, err) + + require.NoError(t, ps.Pub(ctx, topic, msg2)) + + select { + case id := <-receivedCh: + assert.Equal(t, "task-failed", id) + case <-ctx.Done(): + t.Fatal("timed out waiting for message after reconnect") + } + + require.NoError(t, cleanupSub()) +} diff --git a/internal/msgqueue/rabbitmq/rabbitmq.go b/internal/msgqueue/rabbitmq/rabbitmq.go index 55e6848c74..4913ce593c 100644 --- a/internal/msgqueue/rabbitmq/rabbitmq.go +++ b/internal/msgqueue/rabbitmq/rabbitmq.go @@ -12,8 +12,6 @@ import ( "sync" "time" - "github.com/google/uuid" - lru "github.com/hashicorp/golang-lru/v2" "github.com/jackc/pgerrcode" "github.com/jackc/pgx/v5/pgconn" amqp "github.com/rabbitmq/amqp091-go" @@ -24,7 +22,6 @@ import ( "github.com/hatchet-dev/hatchet/internal/msgqueue" "github.com/hatchet-dev/hatchet/internal/queueutils" "github.com/hatchet-dev/hatchet/pkg/logger" - "github.com/hatchet-dev/hatchet/pkg/random" "github.com/hatchet-dev/hatchet/pkg/telemetry" ) @@ -42,19 +39,14 @@ type MessageQueueImpl struct { l *zerolog.Logger - disableTenantExchangePubs bool - - // lru cache for tenant ids - tenantIdCache *lru.Cache[uuid.UUID, bool] - pubChannels *channelPool subChannels *channelPool deadLetterBackoff time.Duration + compressor + maxPayloadSize int - compressionEnabled bool - compressionThreshold int enableMessageRejection bool maxDeathCount int } @@ -66,28 +58,26 @@ func (t *MessageQueueImpl) IsReady() bool { type MessageQueueImplOpt func(*MessageQueueImplOpts) type MessageQueueImplOpts struct { - l *zerolog.Logger - url string - qos int - disableTenantExchangePubs bool - deadLetterBackoff time.Duration - maxPubChannels int32 - maxSubChannels int32 - compressionEnabled bool - compressionThreshold int - enableMessageRejection bool - maxDeathCount int + l *zerolog.Logger + url string + qos int + deadLetterBackoff time.Duration + maxPubChannels int32 + maxSubChannels int32 + compressionEnabled bool + compressionThreshold int + enableMessageRejection bool + maxDeathCount int } func defaultMessageQueueImplOpts() *MessageQueueImplOpts { l := logger.NewDefaultLogger("rabbitmq") return &MessageQueueImplOpts{ - l: &l, - disableTenantExchangePubs: false, - deadLetterBackoff: 5 * time.Second, - enableMessageRejection: false, - maxDeathCount: 5, + l: &l, + deadLetterBackoff: 5 * time.Second, + enableMessageRejection: false, + maxDeathCount: 5, } } @@ -121,12 +111,6 @@ func WithQos(qos int) MessageQueueImplOpt { } } -func WithDisableTenantExchangePubs(disable bool) MessageQueueImplOpt { - return func(opts *MessageQueueImplOpts) { - opts.disableTenantExchangePubs = disable - } -} - func WithDeadLetterBackoff(backoff time.Duration) MessageQueueImplOpt { return func(opts *MessageQueueImplOpts) { opts.deadLetterBackoff = backoff @@ -198,24 +182,22 @@ func New(fs ...MessageQueueImplOpt) (func() error, *MessageQueueImpl, error) { } t := &MessageQueueImpl{ - ctx: ctx, - identity: identity(), - l: opts.l, - qos: opts.qos, - configFs: fs, - disableTenantExchangePubs: opts.disableTenantExchangePubs, - pubChannels: pubChannelPool, - subChannels: subChannelPool, - deadLetterBackoff: opts.deadLetterBackoff, - compressionEnabled: opts.compressionEnabled, - compressionThreshold: opts.compressionThreshold, - enableMessageRejection: opts.enableMessageRejection, - maxDeathCount: opts.maxDeathCount, - maxPayloadSize: 16 * 1024 * 1024, // 16 MB - } - - // create a new lru cache for tenant ids - t.tenantIdCache, _ = lru.New[uuid.UUID, bool](2000) //nolint:errcheck // this only returns an error if the size is less than 0 + ctx: ctx, + identity: identity(), + l: opts.l, + qos: opts.qos, + configFs: fs, + pubChannels: pubChannelPool, + subChannels: subChannelPool, + deadLetterBackoff: opts.deadLetterBackoff, + compressor: compressor{ + compressionEnabled: opts.compressionEnabled, + compressionThreshold: opts.compressionThreshold, + }, + enableMessageRejection: opts.enableMessageRejection, + maxDeathCount: opts.maxDeathCount, + maxPayloadSize: 16 * 1024 * 1024, // 16 MB + } // init the queues in a blocking fashion poolCh, err := subChannelPool.Acquire(ctx) @@ -469,28 +451,6 @@ func (t *MessageQueueImpl) pubMessage(ctx context.Context, q msgqueue.Queue, msg pubSpan.End() - if (!t.disableTenantExchangePubs || msg.ID == "task-stream-event") && msg.TenantID != uuid.Nil { - if _, ok := t.tenantIdCache.Get(msg.TenantID); !ok { - err = t.declareTenantExchange(ctx, pub, msg.TenantID) - - if err != nil { - t.l.Error().Ctx(ctx).Str("tenant_id", msg.TenantID.String()).Msgf("error registering tenant exchange: %v", err) - return err - } - } - - t.l.Debug().Ctx(ctx).Str("tenant_id", msg.TenantID.String()).Msgf("publishing tenant msg to exchange %s", msg.TenantID) - - err = pub.PublishWithContext(ctx, msgqueue.GetTenantExchangeName(msg.TenantID), "", false, false, amqp.Publishing{ - Body: body, - }) - - if err != nil { - t.l.Error().Ctx(ctx).Str("tenant_id", msg.TenantID.String()).Msgf("error publishing tenant msg: %v", err) - return err - } - } - t.l.Debug().Msgf("published msg to queue %s", q.Name()) return nil @@ -550,53 +510,6 @@ func (t *MessageQueueImpl) Subscribe( }, nil } -func (t *MessageQueueImpl) RegisterTenant(ctx context.Context, tenantId uuid.UUID) error { - poolCh, err := t.pubChannels.Acquire(ctx) - - if err != nil { - t.l.Error().Ctx(ctx).Str("tenant_id", tenantId.String()).Msgf("[RegisterTenant] cannot acquire channel: %v", err) - return err - } - - ch := poolCh.Value() - - if ch.IsClosed() { - poolCh.Destroy() - return fmt.Errorf("channel is closed") - } - - defer poolCh.Release() - - return t.declareTenantExchange(ctx, ch, tenantId) -} - -// declareTenantExchange creates the tenant fanout exchange on an already-held channel -// and records the tenant in the local cache. Callers that already hold a pub channel -// (e.g. pubMessage) must use this instead of RegisterTenant to avoid a nested Acquire. -func (t *MessageQueueImpl) declareTenantExchange(ctx context.Context, ch *amqp.Channel, tenantId uuid.UUID) error { - t.l.Debug().Ctx(ctx).Str("tenant_id", tenantId.String()).Msgf("registering tenant exchange: %s", tenantId) - - // each consumer of the fanout exchange will get notified with the tenant events - err := ch.ExchangeDeclare( - msgqueue.GetTenantExchangeName(tenantId), - "fanout", - true, // durable - false, // auto-deleted - false, // not internal, accepts publishings - false, // no-wait - nil, // arguments - ) - - if err != nil { - t.l.Error().Ctx(ctx).Str("tenant_id", tenantId.String()).Msgf("cannot declare exchange: %q, %v", tenantId, err) - return err - } - - t.tenantIdCache.Add(tenantId, true) - - return nil -} - func (t *MessageQueueImpl) initQueue(ch *amqp.Channel, q msgqueue.Queue) (string, error) { if q.IsDLQ() { return q.Name(), nil @@ -606,17 +519,6 @@ func (t *MessageQueueImpl) initQueue(ch *amqp.Channel, q msgqueue.Queue) (string args["x-consumer-timeout"] = 300000 // 5 minutes name := q.Name() - if q.FanoutExchangeKey() != "" { - suffix, err := random.Generate(8) - - if err != nil { - t.l.Error().Msgf("error generating random bytes: %v", err) - return "", err - } - - name = fmt.Sprintf("%s-%s", q.Name(), suffix) - } - if !q.IsDLQ() && q.DLQ() != nil && q.DLQ().IsAutoDLQ() { dlx1 := getTmpDLQName(q.DLQ().Name()) dlx2 := getProcDLQName(q.DLQ().Name()) @@ -662,16 +564,6 @@ func (t *MessageQueueImpl) initQueue(ch *amqp.Channel, q msgqueue.Queue) (string return "", err } - // if the queue has a subscriber key, bind it to the fanout exchange - if q.FanoutExchangeKey() != "" { - t.l.Debug().Msgf("binding queue: %s to exchange: %s", name, q.FanoutExchangeKey()) - - if err := ch.QueueBind(name, "", q.FanoutExchangeKey(), false, nil); err != nil { - t.l.Error().Msgf("cannot bind queue: %q, %v", name, err) - return "", err - } - } - return name, nil } diff --git a/internal/msgqueue/rabbitmq/rabbitmq_test.go b/internal/msgqueue/rabbitmq/rabbitmq_test.go index 4de9d71e36..9e47c3a3b1 100644 --- a/internal/msgqueue/rabbitmq/rabbitmq_test.go +++ b/internal/msgqueue/rabbitmq/rabbitmq_test.go @@ -28,7 +28,7 @@ func TestMessageQueueIntegration(t *testing.T) { defer cancel() wg := &sync.WaitGroup{} - wg.Add(2) // we wait for 2 messages here + wg.Add(1) // we wait for 1 message here url := "amqp://user:password@localhost:5672/" @@ -65,9 +65,6 @@ func TestMessageQueueIntegration(t *testing.T) { t.Fatalf("error creating task: %v", err) } - err = tq.SendMessage(ctx, staticQueue, task) - assert.NoError(t, err, "adding task to static queue should not error") - preAck := func(receivedMessage *msgqueue.Message) error { defer wg.Done() assert.Equal(t, task.ID, receivedMessage.ID, "received task ID should match sent task ID") @@ -78,27 +75,6 @@ func TestMessageQueueIntegration(t *testing.T) { cleanupQueue, err := tq.Subscribe(staticQueue, preAck, msgqueue.NoOpHook) require.NoError(t, err, "subscribing to static queue should not error") - // Test tenant registration and queue creation - err = tq.RegisterTenant(ctx, testTenantUUID) - assert.NoError(t, err, "registering tenant should not error") - - // Assuming there's a mechanism to retrieve a tenant-specific queue, e.g., by tenant ID - tenantQueue := msgqueue.TenantEventConsumerQueue(testTenantUUID) - - if err != nil { - t.Fatalf("error creating tenant-specific queue: %v", err) - } - - tqAck := func(receivedMessage *msgqueue.Message) error { - defer wg.Done() - assert.Equal(t, task.ID, receivedMessage.ID, "received tenant task ID should match sent task ID") - return nil - } - - // Test subscription to the tenant-specific queue - cleanupTenantQueue, err := tq.Subscribe(tenantQueue, tqAck, msgqueue.NoOpHook) - require.NoError(t, err, "subscribing to tenant-specific queue should not error") - // send task to queue after 1 second to give time for subscriber go func() { time.Sleep(1 * time.Second) @@ -111,9 +87,6 @@ func TestMessageQueueIntegration(t *testing.T) { if err := cleanupQueue(); err != nil { t.Fatalf("error cleaning up queue: %v", err) } - if err := cleanupTenantQueue(); err != nil { - t.Fatalf("error cleaning up queue: %v", err) - } } func TestBufferedSubMessageQueueIntegration(t *testing.T) { diff --git a/internal/msgqueue/shared_tenant_reader.go b/internal/msgqueue/shared_tenant_reader.go index d1204349c9..64c77512cf 100644 --- a/internal/msgqueue/shared_tenant_reader.go +++ b/internal/msgqueue/shared_tenant_reader.go @@ -1,7 +1,6 @@ package msgqueue import ( - "context" "sync" "time" @@ -21,13 +20,13 @@ type sharedTenantSub struct { type SharedTenantReader struct { tenants *syncx.Map[uuid.UUID, *sharedTenantSub] - mq MessageQueue + ps PubSub } -func NewSharedTenantReader(mq MessageQueue) *SharedTenantReader { +func NewSharedTenantReader(ps PubSub) *SharedTenantReader { return &SharedTenantReader{ tenants: &syncx.Map[uuid.UUID, *sharedTenantSub]{}, - mq: mq, + ps: ps, } } @@ -48,15 +47,7 @@ func (s *SharedTenantReader) Subscribe(tenantId uuid.UUID, postAck AckHook) (fun if !t.isRunning { t.isRunning = true - q := TenantEventConsumerQueue(tenantId) - - err := s.mq.RegisterTenant(context.Background(), tenantId) - - if err != nil { - return nil, err - } - - cleanupSingleSub, err := s.mq.Subscribe(q, NoOpHook, func(task *Message) error { + cleanupSingleSub, err := s.ps.Sub(TenantTopic(tenantId), func(task *Message) error { var innerErr error t.fs.Range(func(key int, f AckHook) bool { @@ -108,13 +99,13 @@ type sharedBufferedTenantSub struct { type SharedBufferedTenantReader struct { tenants *syncx.Map[uuid.UUID, *sharedBufferedTenantSub] - mq MessageQueue + ps PubSub } -func NewSharedBufferedTenantReader(mq MessageQueue) *SharedBufferedTenantReader { +func NewSharedBufferedTenantReader(ps PubSub) *SharedBufferedTenantReader { return &SharedBufferedTenantReader{ tenants: &syncx.Map[uuid.UUID, *sharedBufferedTenantSub]{}, - mq: mq, + ps: ps, } } @@ -135,15 +126,11 @@ func (s *SharedBufferedTenantReader) Subscribe(tenantId uuid.UUID, f DstFunc) (f if !t.isRunning { t.isRunning = true - q := TenantEventConsumerQueue(tenantId) - - err := s.mq.RegisterTenant(context.Background(), tenantId) - - if err != nil { - return nil, err - } - - subBuffer := NewMQSubBuffer(q, s.mq, func(tenantId uuid.UUID, msgId string, payloads [][]byte) error { + // the buffer runs in PostAck mode, which only uses the post hook, so the + // single-handler pubsub Sub maps cleanly onto the subscribe function + subBuffer := NewSubBufferFromSubscribe(func(preAck AckHook, postAck AckHook) (func() error, error) { + return s.ps.Sub(TenantTopic(tenantId), postAck) + }, func(tenantId uuid.UUID, msgId string, payloads [][]byte) error { var innerErr error t.fs.Range(func(key int, f DstFunc) bool { diff --git a/internal/services/admin/admin.go b/internal/services/admin/admin.go index 676639e152..1b3bf01439 100644 --- a/internal/services/admin/admin.go +++ b/internal/services/admin/admin.go @@ -28,6 +28,7 @@ type AdminServiceImpl struct { repov1 v1.Repository mqv1 msgqueue.MessageQueue + pubsub msgqueue.PubSub v validator.Validator analytics analytics.Analytics @@ -45,6 +46,7 @@ type AdminServiceOpt func(*AdminServiceOpts) type AdminServiceOpts struct { repov1 v1.Repository mqv1 msgqueue.MessageQueue + pubsub msgqueue.PubSub v validator.Validator analytics analytics.Analytics localScheduler *scheduler.Scheduler @@ -79,6 +81,12 @@ func WithMessageQueueV1(mq msgqueue.MessageQueue) AdminServiceOpt { } } +func WithPubSub(pubsub msgqueue.PubSub) AdminServiceOpt { + return func(opts *AdminServiceOpts) { + opts.pubsub = pubsub + } +} + func WithValidator(v validator.Validator) AdminServiceOpt { return func(opts *AdminServiceOpts) { opts.v = v @@ -148,13 +156,17 @@ func NewAdminService(fs ...AdminServiceOpt) (AdminService, error) { return nil, fmt.Errorf("task queue v1 is required. use WithMessageQueueV1") } + if opts.pubsub == nil { + return nil, fmt.Errorf("pubsub is required. use WithPubSub") + } + slots := 0 if opts.grpcTriggersEnabled { slots = opts.grpcTriggerSlots } pubBuffer := msgqueue.NewMQPubBuffer(opts.mqv1) - tw := trigger.NewTriggerWriter(opts.mqv1, opts.repov1, opts.l, pubBuffer, slots, opts.promGate) + tw := trigger.NewTriggerWriter(opts.mqv1, opts.pubsub, opts.repov1, opts.l, pubBuffer, slots, opts.promGate) var localScheduler *scheduler.Scheduler @@ -165,6 +177,7 @@ func NewAdminService(fs ...AdminServiceOpt) (AdminService, error) { return &AdminServiceImpl{ repov1: opts.repov1, mqv1: opts.mqv1, + pubsub: opts.pubsub, v: opts.v, analytics: opts.analytics, localScheduler: localScheduler, diff --git a/internal/services/admin/server.go b/internal/services/admin/server.go index fafa413876..0435ef6528 100644 --- a/internal/services/admin/server.go +++ b/internal/services/admin/server.go @@ -96,14 +96,14 @@ func (a *AdminServiceImpl) PutWorkflow(ctx context.Context, req *contracts.PutWo if err != nil { a.l.Err(err).Ctx(ctx).Msg("could not create message for notifying new queue") } else { - err = a.mqv1.SendMessage( + err = a.pubsub.Pub( notifyCtx, - msgqueue.QueueTypeFromPartitionIDAndController(tenant.SchedulerPartitionId.String, msgqueue.Scheduler), + msgqueue.SchedulerPartitionTopic(tenant.SchedulerPartitionId.String), msg, ) if err != nil { - a.l.Err(err).Ctx(ctx).Msg("could not add message to scheduler partition queue") + a.l.Err(err).Ctx(ctx).Msg("could not publish message to scheduler partition topic") } a.l.Debug().Ctx(ctx).Msgf("notified new queue for tenant %s and action %s", tenantId, action) diff --git a/internal/services/admin/v1/admin.go b/internal/services/admin/v1/admin.go index 5b44fd43de..03bb0004da 100644 --- a/internal/services/admin/v1/admin.go +++ b/internal/services/admin/v1/admin.go @@ -27,6 +27,7 @@ type AdminServiceImpl struct { repo v1.Repository mq msgqueue.MessageQueue + pubsub msgqueue.PubSub v validator.Validator analytics analytics.Analytics @@ -43,6 +44,7 @@ type AdminServiceOpt func(*AdminServiceOpts) type AdminServiceOpts struct { repo v1.Repository mq msgqueue.MessageQueue + pubsub msgqueue.PubSub v validator.Validator analytics analytics.Analytics @@ -80,6 +82,12 @@ func WithMessageQueue(mq msgqueue.MessageQueue) AdminServiceOpt { } } +func WithPubSub(pubsub msgqueue.PubSub) AdminServiceOpt { + return func(opts *AdminServiceOpts) { + opts.pubsub = pubsub + } +} + func WithValidator(v validator.Validator) AdminServiceOpt { return func(opts *AdminServiceOpts) { opts.v = v @@ -149,8 +157,12 @@ func NewAdminService(fs ...AdminServiceOpt) (AdminService, error) { return nil, fmt.Errorf("task queue is required. use WithMessageQueue") } + if opts.pubsub == nil { + return nil, fmt.Errorf("pubsub is required. use WithPubSub") + } + pubBuffer := msgqueue.NewMQPubBuffer(opts.mq) - tw := trigger.NewTriggerWriter(opts.mq, opts.repo, opts.l, pubBuffer, opts.grpcTriggerSlots, opts.promGate) + tw := trigger.NewTriggerWriter(opts.mq, opts.pubsub, opts.repo, opts.l, pubBuffer, opts.grpcTriggerSlots, opts.promGate) var localScheduler *scheduler.Scheduler @@ -161,6 +173,7 @@ func NewAdminService(fs ...AdminServiceOpt) (AdminService, error) { return &AdminServiceImpl{ repo: opts.repo, mq: opts.mq, + pubsub: opts.pubsub, v: opts.v, analytics: opts.analytics, localScheduler: localScheduler, diff --git a/internal/services/admin/v1/server.go b/internal/services/admin/v1/server.go index ebd26c136e..964c5b7cbc 100644 --- a/internal/services/admin/v1/server.go +++ b/internal/services/admin/v1/server.go @@ -760,14 +760,14 @@ func (a *AdminServiceImpl) PutWorkflow(ctx context.Context, req *contracts.Creat if err != nil { a.l.Err(err).Ctx(ctx).Msg("could not create message for notifying new queue") } else { - err = a.mq.SendMessage( + err = a.pubsub.Pub( notifyCtx, - msgqueue.QueueTypeFromPartitionIDAndController(tenant.SchedulerPartitionId.String, msgqueue.Scheduler), + msgqueue.SchedulerPartitionTopic(tenant.SchedulerPartitionId.String), msg, ) if err != nil { - a.l.Err(err).Ctx(ctx).Msg("could not add message to scheduler partition queue") + a.l.Err(err).Ctx(ctx).Msg("could not publish message to scheduler partition topic") } } } diff --git a/internal/services/controllers/olap/controller.go b/internal/services/controllers/olap/controller.go index b9fb0d809d..ab91ee0018 100644 --- a/internal/services/controllers/olap/controller.go +++ b/internal/services/controllers/olap/controller.go @@ -68,6 +68,7 @@ func (tc *OLAPControllerImpl) logRepublish(ctx context.Context, entity string, c type OLAPControllerImpl struct { mq msgqueue.MessageQueue + pubsub msgqueue.PubSub l *zerolog.Logger repo v1.Repository dv datautils.DataDecoderValidator @@ -96,6 +97,7 @@ type OLAPControllerOpt func(*OLAPControllerOpts) type OLAPControllerOpts struct { mq msgqueue.MessageQueue + pubsub msgqueue.PubSub l *zerolog.Logger repo v1.Repository dv datautils.DataDecoderValidator @@ -132,6 +134,12 @@ func WithMessageQueue(mq msgqueue.MessageQueue) OLAPControllerOpt { } } +func WithPubSub(pubsub msgqueue.PubSub) OLAPControllerOpt { + return func(opts *OLAPControllerOpts) { + opts.pubsub = pubsub + } +} + func WithLogger(l *zerolog.Logger) OLAPControllerOpt { return func(opts *OLAPControllerOpts) { opts.l = l @@ -236,6 +244,10 @@ func New(fs ...OLAPControllerOpt) (*OLAPControllerImpl, error) { return nil, fmt.Errorf("task queue is required. use WithMessageQueue") } + if opts.pubsub == nil { + return nil, fmt.Errorf("pubsub is required. use WithPubSub") + } + if opts.repo == nil { return nil, fmt.Errorf("repository is required. use WithRepository") } @@ -270,6 +282,7 @@ func New(fs ...OLAPControllerOpt) (*OLAPControllerImpl, error) { o := &OLAPControllerImpl{ mq: opts.mq, + pubsub: opts.pubsub, l: opts.l, s: s, p: opts.p, @@ -1161,6 +1174,12 @@ func (tc *OLAPControllerImpl) republishCreatedTasks(ctx context.Context, tenantI if err := tc.mq.SendMessage(ctx, msgqueue.OLAP_QUEUE, msg); err != nil { return err } + + // best-effort publish to the tenant stream, mirroring the original + // created-task publish this requeue stands in for + if err := tc.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + tc.l.Warn().Ctx(ctx).Err(err).Msg("could not publish created-task to tenant stream") + } } return nil } diff --git a/internal/services/controllers/olap/process_dag_status_updates.go b/internal/services/controllers/olap/process_dag_status_updates.go index e364eb3169..e007ac63e6 100644 --- a/internal/services/controllers/olap/process_dag_status_updates.go +++ b/internal/services/controllers/olap/process_dag_status_updates.go @@ -103,12 +103,10 @@ func (o *OLAPControllerImpl) notifyDAGsUpdated(ctx context.Context, rows []v1.Up return err } - q := msgqueue.TenantEventConsumerQueue(tenantId) - - err = o.mq.SendMessage(ctx, q, msg) - - if err != nil { - return err + // best-effort publish to the tenant stream: the dispatcher's workflow + // run subscriptions consume workflow-run-finished + if err := o.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + o.l.Warn().Ctx(ctx).Err(err).Msg("could not publish workflow-run-finished to tenant stream") } } } diff --git a/internal/services/controllers/olap/process_task_status_updates.go b/internal/services/controllers/olap/process_task_status_updates.go index 82d511cfdf..92a778bff4 100644 --- a/internal/services/controllers/olap/process_task_status_updates.go +++ b/internal/services/controllers/olap/process_task_status_updates.go @@ -103,12 +103,10 @@ func (o *OLAPControllerImpl) notifyTasksUpdated(ctx context.Context, rows []v1.U return err } - q := msgqueue.TenantEventConsumerQueue(tenantId) - - err = o.mq.SendMessage(ctx, q, msg) - - if err != nil { - return err + // best-effort publish to the tenant stream: the dispatcher's workflow + // run subscriptions consume workflow-run-finished + if err := o.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + o.l.Warn().Ctx(ctx).Err(err).Msg("could not publish workflow-run-finished to tenant stream") } } } diff --git a/internal/services/controllers/olap/signal/signal.go b/internal/services/controllers/olap/signal/signal.go index 8ee6166af3..20f2011ca6 100644 --- a/internal/services/controllers/olap/signal/signal.go +++ b/internal/services/controllers/olap/signal/signal.go @@ -20,15 +20,17 @@ import ( type OLAPSignaler struct { mq msgqueue.MessageQueue + pubsub msgqueue.PubSub repo v1.Repository pubBuffer *msgqueue.MQPubBuffer l *zerolog.Logger promGate *prometheus.Gate } -func NewOLAPSignaler(mq msgqueue.MessageQueue, repo v1.Repository, l *zerolog.Logger, pubBuffer *msgqueue.MQPubBuffer, promGate *prometheus.Gate) *OLAPSignaler { +func NewOLAPSignaler(mq msgqueue.MessageQueue, pubsub msgqueue.PubSub, repo v1.Repository, l *zerolog.Logger, pubBuffer *msgqueue.MQPubBuffer, promGate *prometheus.Gate) *OLAPSignaler { return &OLAPSignaler{ mq: mq, + pubsub: pubsub, l: l, repo: repo, pubBuffer: pubBuffer, @@ -119,6 +121,12 @@ func (s *OLAPSignaler) SignalTasksCreated(ctx context.Context, tenantId uuid.UUI s.l.Err(err).Ctx(ctx).Msg("could not add message to olap queue") continue } + + // best-effort publish to the tenant stream: the dispatcher's workflow + // event subscriptions consume created-task + if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish created-task to tenant stream") + } } eg := &errgroup.Group{} @@ -267,14 +275,14 @@ func (s *OLAPSignaler) signalTasksCreatedAndQueued(ctx context.Context, tenantId if err != nil { s.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not create message for scheduler partition queue") } else { - err = s.mq.SendMessage( + err = s.pubsub.Pub( ctx, - msgqueue.QueueTypeFromPartitionIDAndController(tenant.SchedulerPartitionId.String, msgqueue.Scheduler), + msgqueue.SchedulerPartitionTopic(tenant.SchedulerPartitionId.String), msg, ) if err != nil { - s.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not add message to scheduler partition queue") + s.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not publish message to scheduler partition topic") } } } diff --git a/internal/services/controllers/task/controller.go b/internal/services/controllers/task/controller.go index 88db238500..4d002cd8dc 100644 --- a/internal/services/controllers/task/controller.go +++ b/internal/services/controllers/task/controller.go @@ -43,6 +43,7 @@ type TasksController interface { type TasksControllerImpl struct { mq msgqueue.MessageQueue + pubsub msgqueue.PubSub pubBuffer *msgqueue.MQPubBuffer l *zerolog.Logger queueLogger *zerolog.Logger @@ -73,6 +74,7 @@ type TasksControllerOpt func(*TasksControllerOpts) type TasksControllerOpts struct { mq msgqueue.MessageQueue + pubsub msgqueue.PubSub l *zerolog.Logger repov1 v1.Repository dv datautils.DataDecoderValidator @@ -113,6 +115,12 @@ func WithMessageQueue(mq msgqueue.MessageQueue) TasksControllerOpt { } } +func WithPubSub(pubsub msgqueue.PubSub) TasksControllerOpt { + return func(opts *TasksControllerOpts) { + opts.pubsub = pubsub + } +} + func WithLogger(l *zerolog.Logger) TasksControllerOpt { return func(opts *TasksControllerOpts) { opts.l = l @@ -193,6 +201,10 @@ func New(fs ...TasksControllerOpt) (*TasksControllerImpl, error) { return nil, fmt.Errorf("task queue is required. use WithMessageQueue") } + if opts.pubsub == nil { + return nil, fmt.Errorf("pubsub is required. use WithPubSub") + } + if opts.repov1 == nil { return nil, fmt.Errorf("v2 repository is required. use WithV2Repository") } @@ -215,11 +227,12 @@ func New(fs ...TasksControllerOpt) (*TasksControllerImpl, error) { pubBuffer := msgqueue.NewMQPubBuffer(opts.mq) - signaler := signal.NewOLAPSignaler(opts.mq, opts.repov1, opts.l, pubBuffer, opts.promGate) - tw := trigger.NewTriggerWriter(opts.mq, opts.repov1, opts.l, pubBuffer, 0, opts.promGate) + signaler := signal.NewOLAPSignaler(opts.mq, opts.pubsub, opts.repov1, opts.l, pubBuffer, opts.promGate) + tw := trigger.NewTriggerWriter(opts.mq, opts.pubsub, opts.repov1, opts.l, pubBuffer, 0, opts.promGate) t := &TasksControllerImpl{ mq: opts.mq, + pubsub: opts.pubsub, pubBuffer: pubBuffer, l: opts.l, queueLogger: opts.queueLogger, @@ -991,14 +1004,14 @@ func (tc *TasksControllerImpl) notifyQueuesOnCompletion(ctx context.Context, ten if err != nil { tc.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not create message for scheduler partition queue") } else { - err = tc.mq.SendMessage( + err = tc.pubsub.Pub( ctx, - msgqueue.QueueTypeFromPartitionIDAndController(tenant.SchedulerPartitionId.String, msgqueue.Scheduler), + msgqueue.SchedulerPartitionTopic(tenant.SchedulerPartitionId.String), msg, ) if err != nil { - tc.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not add message to scheduler partition queue") + tc.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not publish message to scheduler partition topic") } } } @@ -1024,14 +1037,12 @@ func (tc *TasksControllerImpl) notifyQueuesOnCompletion(ctx context.Context, ten return } - err = tc.mq.SendMessage( - ctx, - msgqueue.TenantEventConsumerQueue(tenantId), - msg, - ) + // best-effort publish to the tenant stream: the dispatcher's workflow run + // subscriptions consume workflow-run-finished-candidate + err = tc.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg) if err != nil { - tc.l.Err(err).Ctx(ctx).Msg("could not send workflow-run-finished-candidate message") + tc.l.Err(err).Ctx(ctx).Msg("could not publish workflow-run-finished-candidate message") return } } diff --git a/internal/services/controllers/task/durable_callbacks.go b/internal/services/controllers/task/durable_callbacks.go index 71f354a927..92f54827d6 100644 --- a/internal/services/controllers/task/durable_callbacks.go +++ b/internal/services/controllers/task/durable_callbacks.go @@ -143,7 +143,7 @@ func (tc *TasksControllerImpl) notifySchedulerQueues(ctx context.Context, tenant return fmt.Errorf("failed to build check-tenant-queue message for queues %v: %w", queueNames, err) } - if err := tc.mq.SendMessage(ctx, msgqueue.QueueTypeFromPartitionIDAndController(tenant.SchedulerPartitionId.String, msgqueue.Scheduler), msg); err != nil { + if err := tc.pubsub.Pub(ctx, msgqueue.SchedulerPartitionTopic(tenant.SchedulerPartitionId.String), msg); err != nil { return fmt.Errorf("failed to notify scheduler for queues %v: %w", queueNames, err) } diff --git a/internal/services/controllers/task/trigger/trigger.go b/internal/services/controllers/task/trigger/trigger.go index 0d46c8d872..c28bec660d 100644 --- a/internal/services/controllers/task/trigger/trigger.go +++ b/internal/services/controllers/task/trigger/trigger.go @@ -32,8 +32,8 @@ var ErrNoTriggerSlots = errors.New("no trigger slots available") // NewTriggerWriter creates a new TriggerWriter with the given number of slots for concurrency control. // If the number of slots is 0, there is no limit to concurrency. -func NewTriggerWriter(mq msgqueue.MessageQueue, repo v1.Repository, l *zerolog.Logger, pubBuffer *msgqueue.MQPubBuffer, slots int, promGate *prometheus.Gate) *TriggerWriter { - s := signal.NewOLAPSignaler(mq, repo, l, pubBuffer, promGate) +func NewTriggerWriter(mq msgqueue.MessageQueue, pubsub msgqueue.PubSub, repo v1.Repository, l *zerolog.Logger, pubBuffer *msgqueue.MQPubBuffer, slots int, promGate *prometheus.Gate) *TriggerWriter { + s := signal.NewOLAPSignaler(mq, pubsub, repo, l, pubBuffer, promGate) var sem chan struct{} diff --git a/internal/services/dispatcher/dispatcher.go b/internal/services/dispatcher/dispatcher.go index 261e26a0fa..b4ebb13f2c 100644 --- a/internal/services/dispatcher/dispatcher.go +++ b/internal/services/dispatcher/dispatcher.go @@ -45,6 +45,7 @@ type DispatcherImpl struct { v validator.Validator s gocron.Scheduler mqv1 msgqueue.MessageQueue + pubsub msgqueue.PubSub analytics analytics.Analytics cache cache.Cacheable repov1 v1.Repository @@ -148,6 +149,7 @@ type DispatcherOpts struct { repov1 v1.Repository alerter hatcheterrors.Alerter mqv1 msgqueue.MessageQueue + pubsub msgqueue.PubSub analytics analytics.Analytics l *zerolog.Logger version string @@ -182,6 +184,12 @@ func WithMessageQueueV1(mqv1 msgqueue.MessageQueue) DispatcherOpt { } } +func WithPubSub(pubsub msgqueue.PubSub) DispatcherOpt { + return func(opts *DispatcherOpts) { + opts.pubsub = pubsub + } +} + func WithAlerter(a hatcheterrors.Alerter) DispatcherOpt { return func(opts *DispatcherOpts) { opts.alerter = a @@ -271,6 +279,10 @@ func New(fs ...DispatcherOpt) (*DispatcherImpl, error) { return nil, fmt.Errorf("v1 task queue is required. use WithMessageQueueV1") } + if opts.pubsub == nil { + return nil, fmt.Errorf("pubsub is required. use WithPubSub") + } + if opts.repov1 == nil { return nil, fmt.Errorf("v1 repository is required. use WithRepositoryV1") } @@ -298,6 +310,7 @@ func New(fs ...DispatcherOpt) (*DispatcherImpl, error) { return &DispatcherImpl{ mqv1: opts.mqv1, + pubsub: opts.pubsub, pubBuffer: pubBuffer, l: opts.l, dv: opts.dv, @@ -315,14 +328,14 @@ func New(fs ...DispatcherOpt) (*DispatcherImpl, error) { analytics: opts.analytics, streamEventBufferTimeout: opts.streamEventBufferTimeout, version: opts.version, - serviceV1: newDispatcherService(opts.repov1, opts.mqv1, v, opts.l, opts.dispatcherId, opts.analytics, opts.promGate), + serviceV1: newDispatcherService(opts.repov1, opts.mqv1, opts.pubsub, v, opts.l, opts.dispatcherId, opts.analytics, opts.promGate), }, nil } func (d *DispatcherImpl) Start() (func() error, error) { ctx, cancel := context.WithCancel(context.Background()) - d.sharedNonBufferedReaderv1 = msgqueue.NewSharedTenantReader(d.mqv1) - d.sharedBufferedReaderv1 = msgqueue.NewSharedBufferedTenantReader(d.mqv1) + d.sharedNonBufferedReaderv1 = msgqueue.NewSharedTenantReader(d.pubsub) + d.sharedBufferedReaderv1 = msgqueue.NewSharedBufferedTenantReader(d.pubsub) // register the dispatcher by creating a new dispatcher in the database dispatcher, err := d.repov1.Dispatcher().CreateNewDispatcher(ctx, &v1.CreateDispatcherOpts{ @@ -950,6 +963,12 @@ func (d *DispatcherImpl) handleRetries( return fmt.Errorf("could not send failed task message: %w", err) } + // best-effort publish to the tenant stream: the dispatcher's workflow + // event subscriptions consume task-failed + if err := d.pubsub.Pub(retryCtx, msgqueue.TenantTopic(tenantId), msg); err != nil { + d.l.Warn().Ctx(retryCtx).Err(err).Msg("could not publish task-failed to tenant stream") + } + return nil }) } diff --git a/internal/services/dispatcher/dispatcher_v1.go b/internal/services/dispatcher/dispatcher_v1.go index 7b6a2e6ab8..0304a0ca54 100644 --- a/internal/services/dispatcher/dispatcher_v1.go +++ b/internal/services/dispatcher/dispatcher_v1.go @@ -41,9 +41,9 @@ func (d *DispatcherServiceImpl) CancelStreamSessions() { d.streamSessions.CancelAll() } -func newDispatcherService(repo v1.Repository, mq msgqueue.MessageQueue, v validator.Validator, l *zerolog.Logger, dispatcherId uuid.UUID, a analytics.Analytics, promGate *prometheus.Gate) *DispatcherServiceImpl { +func newDispatcherService(repo v1.Repository, mq msgqueue.MessageQueue, pubsub msgqueue.PubSub, v validator.Validator, l *zerolog.Logger, dispatcherId uuid.UUID, a analytics.Analytics, promGate *prometheus.Gate) *DispatcherServiceImpl { pubBuffer := msgqueue.NewMQPubBuffer(mq) - tw := trigger.NewTriggerWriter(mq, repo, l, pubBuffer, 0, promGate) + tw := trigger.NewTriggerWriter(mq, pubsub, repo, l, pubBuffer, 0, promGate) return &DispatcherServiceImpl{ repo: repo, diff --git a/internal/services/dispatcher/server.go b/internal/services/dispatcher/server.go index b9bbf1a3b7..f6ec81d382 100644 --- a/internal/services/dispatcher/server.go +++ b/internal/services/dispatcher/server.go @@ -468,14 +468,14 @@ func (s *DispatcherImpl) Heartbeat(ctx context.Context, req *contracts.Heartbeat if err != nil { s.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not create message for notifying new worker") } else { - err = s.mqv1.SendMessage( + err = s.pubsub.Pub( notifyCtx, - msgqueue.QueueTypeFromPartitionIDAndController(tenant.SchedulerPartitionId.String, msgqueue.Scheduler), + msgqueue.SchedulerPartitionTopic(tenant.SchedulerPartitionId.String), msg, ) if err != nil { - s.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not add message to scheduler partition queue") + s.l.Err(err).Ctx(ctx).Str("scheduler_partition_id", tenant.SchedulerPartitionId.String).Msg("could not publish message to scheduler partition topic") } } }() @@ -1353,6 +1353,12 @@ func (s *DispatcherImpl) handleTaskCompleted(inputCtx context.Context, task *sql return nil, err } + // best-effort publish to the tenant stream: the dispatcher's workflow event + // subscriptions consume task-completed + if err := s.pubsub.Pub(inputCtx, msgqueue.TenantTopic(tenantId), msg); err != nil { + s.l.Warn().Ctx(inputCtx).Err(err).Msg("could not publish task-completed to tenant stream") + } + resp := &contracts.ActionEventResponse{ TenantId: tenantId.String(), WorkerId: request.WorkerId, @@ -1413,6 +1419,12 @@ func (s *DispatcherImpl) handleTaskFailed(inputCtx context.Context, task *sqlcv1 return nil, err } + // best-effort publish to the tenant stream: the dispatcher's workflow event + // subscriptions consume task-failed + if err := s.pubsub.Pub(inputCtx, msgqueue.TenantTopic(tenantId), msg); err != nil { + s.l.Warn().Ctx(inputCtx).Err(err).Msg("could not publish task-failed to tenant stream") + } + return &contracts.ActionEventResponse{ TenantId: tenantId.String(), WorkerId: request.WorkerId, diff --git a/internal/services/health/health.go b/internal/services/health/health.go index eae6320598..4d43db39d7 100644 --- a/internal/services/health/health.go +++ b/internal/services/health/health.go @@ -21,14 +21,16 @@ type Health struct { repository v1.HealthRepository queue msgqueue.MessageQueue + pubsub msgqueue.PubSub l *zerolog.Logger } -func New(repo v1.HealthRepository, queue msgqueue.MessageQueue, version string, l *zerolog.Logger) *Health { +func New(repo v1.HealthRepository, queue msgqueue.MessageQueue, pubsub msgqueue.PubSub, version string, l *zerolog.Logger) *Health { return &Health{ version: version, repository: repo, queue: queue, + pubsub: pubsub, l: l, } } @@ -44,7 +46,7 @@ func (h *Health) Start(port int) (func() error, error) { ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) defer cancel() - queueReady := h.queue.IsReady() + queueReady := h.queue.IsReady() && h.pubsub.IsReady() repositoryReady := h.repository.IsHealthy(ctx) if !queueReady || !repositoryReady { @@ -60,7 +62,7 @@ func (h *Health) Start(port int) (func() error, error) { ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) defer cancel() - queueReady := h.queue.IsReady() + queueReady := h.queue.IsReady() && h.pubsub.IsReady() repositoryReady := h.repository.IsHealthy(ctx) if h.shuttingDown || !queueReady || !repositoryReady { diff --git a/internal/services/ingestor/ingestor.go b/internal/services/ingestor/ingestor.go index 78edb2d76b..2c017dca34 100644 --- a/internal/services/ingestor/ingestor.go +++ b/internal/services/ingestor/ingestor.go @@ -36,6 +36,7 @@ type IngestorOptFunc func(*IngestorOpts) type IngestorOpts struct { mqv1 msgqueue.MessageQueue + pubsub msgqueue.PubSub repov1 v1.Repository analytics analytics.Analytics isLogIngestionEnabled bool @@ -57,6 +58,12 @@ func WithMessageQueueV1(mq msgqueue.MessageQueue) IngestorOptFunc { } } +func WithPubSub(pubsub msgqueue.PubSub) IngestorOptFunc { + return func(opts *IngestorOpts) { + opts.pubsub = pubsub + } +} + func WithRepositoryV1(r v1.Repository) IngestorOptFunc { return func(opts *IngestorOpts) { opts.repov1 = r @@ -133,6 +140,7 @@ type IngestorImpl struct { steprunTenantLookupCache *lru.Cache[string, string] mqv1 msgqueue.MessageQueue + pubsub msgqueue.PubSub v validator.Validator repov1 v1.Repository @@ -158,6 +166,10 @@ func NewIngestor(fs ...IngestorOptFunc) (Ingestor, error) { return nil, fmt.Errorf("task queue v1 is required. use WithMessageQueueV1") } + if opts.pubsub == nil { + return nil, fmt.Errorf("pubsub is required. use WithPubSub") + } + if opts.repov1 == nil { return nil, fmt.Errorf("repository v1 is required. use WithRepositoryV1") } @@ -175,7 +187,7 @@ func NewIngestor(fs ...IngestorOptFunc) (Ingestor, error) { if opts.grpcTriggersEnabled { pubBuffer = msgqueue.NewMQPubBuffer(opts.mqv1) - tw = trigger.NewTriggerWriter(opts.mqv1, opts.repov1, opts.l, pubBuffer, opts.grpcTriggerSlots, opts.promGate) + tw = trigger.NewTriggerWriter(opts.mqv1, opts.pubsub, opts.repov1, opts.l, pubBuffer, opts.grpcTriggerSlots, opts.promGate) } var localScheduler *scheduler.Scheduler @@ -187,6 +199,7 @@ func NewIngestor(fs ...IngestorOptFunc) (Ingestor, error) { return &IngestorImpl{ steprunTenantLookupCache: stepRunCache, mqv1: opts.mqv1, + pubsub: opts.pubsub, v: validator.NewDefaultValidator(), repov1: opts.repov1, analytics: opts.analytics, diff --git a/internal/services/ingestor/server_v1.go b/internal/services/ingestor/server_v1.go index 26e6e7c9e3..df0fe0f3ac 100644 --- a/internal/services/ingestor/server_v1.go +++ b/internal/services/ingestor/server_v1.go @@ -51,9 +51,8 @@ func (i *IngestorImpl) putStreamEventV1(ctx context.Context, tenant *sqlcv1.Tena return nil, err } - q := msgqueue.TenantEventConsumerQueue(tenantId) - - err = i.mqv1.SendMessage(ctx, q, msg) + // stream events are fanout-only: publish straight to the tenant stream + err = i.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg) if err != nil { return nil, err diff --git a/internal/services/scheduler/v1/scheduler.go b/internal/services/scheduler/v1/scheduler.go index 545e236118..75a3e5800c 100644 --- a/internal/services/scheduler/v1/scheduler.go +++ b/internal/services/scheduler/v1/scheduler.go @@ -33,6 +33,7 @@ type SchedulerOpt func(*SchedulerOpts) type SchedulerOpts struct { mq msgqueue.MessageQueue + pubsub msgqueue.PubSub l *zerolog.Logger repov1 repov1.Repository dv datautils.DataDecoderValidator @@ -63,6 +64,12 @@ func WithMessageQueue(mq msgqueue.MessageQueue) SchedulerOpt { } } +func WithPubSub(pubsub msgqueue.PubSub) SchedulerOpt { + return func(opts *SchedulerOpts) { + opts.pubsub = pubsub + } +} + func WithLogger(l *zerolog.Logger) SchedulerOpt { return func(opts *SchedulerOpts) { opts.l = l @@ -114,6 +121,7 @@ func WithPrometheusGate(gate *prometheus.Gate) SchedulerOpt { type Scheduler struct { mq msgqueue.MessageQueue + pubsub msgqueue.PubSub pubBuffer *msgqueue.MQPubBuffer l *zerolog.Logger repov1 repov1.Repository @@ -145,6 +153,10 @@ func New( return nil, fmt.Errorf("task queue is required. use WithMessageQueue") } + if opts.pubsub == nil { + return nil, fmt.Errorf("pubsub is required. use WithPubSub") + } + if opts.repov1 == nil { return nil, fmt.Errorf("v2 repository is required. use WithV2Repository") } @@ -171,10 +183,11 @@ func New( // TODO: replace with config or pull into a constant tasksWithNoWorkerCache := expirable.NewLRU(10000, func(string, struct{}) {}, 5*time.Minute) - signaler := signal.NewOLAPSignaler(opts.mq, opts.repov1, opts.l, pubBuffer, opts.promGate) + signaler := signal.NewOLAPSignaler(opts.mq, opts.pubsub, opts.repov1, opts.l, pubBuffer, opts.promGate) q := &Scheduler{ mq: opts.mq, + pubsub: opts.pubsub, pubBuffer: pubBuffer, l: opts.l, repov1: opts.repov1, @@ -230,9 +243,8 @@ func (s *Scheduler) Start() (func() error, error) { return nil } - cleanupQueue, err := s.mq.Subscribe( - msgqueue.QueueTypeFromPartitionIDAndController(s.p.GetSchedulerPartitionId(), msgqueue.Scheduler), - msgqueue.NoOpHook, // the only handler is to check the queue, so we acknowledge immediately with the NoOpHook + cleanupQueue, err := s.pubsub.Sub( + msgqueue.SchedulerPartitionTopic(s.p.GetSchedulerPartitionId()), postAck, ) @@ -613,6 +625,13 @@ func (s *Scheduler) scheduleStepRuns(ctx context.Context, tenantId uuid.UUID, re if err != nil { outerErr = multierror.Append(outerErr, fmt.Errorf("could not send cancelled task: %w", err)) + continue + } + + // best-effort publish to the tenant stream: the dispatcher's workflow + // event subscriptions consume task-cancelled + if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish task-cancelled to tenant stream") } } } @@ -691,6 +710,12 @@ func (s *Scheduler) internalRetry(ctx context.Context, tenantId uuid.UUID, assig s.l.Error().Ctx(ctx).Err(err).Msg("could not send failed task") continue } + + // best-effort publish to the tenant stream: the dispatcher's workflow + // event subscriptions consume task-failed + if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish task-failed to tenant stream") + } } } @@ -769,6 +794,12 @@ func (s *Scheduler) notifyAfterConcurrency(ctx context.Context, tenantId uuid.UU s.l.Error().Ctx(ctx).Err(err).Msg("could not send cancelled task") continue } + + // best-effort publish to the tenant stream: the dispatcher's workflow + // event subscriptions consume task-cancelled + if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish task-cancelled to tenant stream") + } } } @@ -926,6 +957,12 @@ func (s *Scheduler) handleDeadLetteredTaskBulkAssigned(ctx context.Context, msg // on the downstream `task-failed` processing to be idempotent. return fmt.Errorf("could not send failed task message: %w", err) } + + // best-effort publish to the tenant stream: the dispatcher's workflow + // event subscriptions consume task-failed + if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish task-failed to tenant stream") + } } return nil diff --git a/internal/services/ticker/schedule_workflow_v1_test.go b/internal/services/ticker/schedule_workflow_v1_test.go index bbab76a341..7b498d1aab 100644 --- a/internal/services/ticker/schedule_workflow_v1_test.go +++ b/internal/services/ticker/schedule_workflow_v1_test.go @@ -75,8 +75,7 @@ func (m *mockMQ) SetQOS(_ int) {} func (m *mockMQ) Subscribe(_ msgqueue.Queue, _, _ msgqueue.AckHook) (func() error, error) { return nil, nil } -func (m *mockMQ) RegisterTenant(_ context.Context, _ uuid.UUID) error { return nil } -func (m *mockMQ) IsReady() bool { return true } +func (m *mockMQ) IsReady() bool { return true } type mockRepo struct { v1.Repository diff --git a/pkg/config/database/config.go b/pkg/config/database/config.go index 4115ce59a6..45a60eeaba 100644 --- a/pkg/config/database/config.go +++ b/pkg/config/database/config.go @@ -70,6 +70,11 @@ type Layer struct { Pool *pgxpool.Pool + // DirectDatabaseURL is the resolved direct (non-pgbouncer) connection URL. + // LISTEN/NOTIFY does not survive transaction pooling, so consumers that + // need LISTEN (e.g. the postgres PubSub) must build their pools from this. + DirectDatabaseURL string + ReadReplicaPool *pgxpool.Pool // DDLPool is meant for DDL-modifying operations like DETACH PARTITION CONCURRENTLY, which diff --git a/pkg/config/loader/loader.go b/pkg/config/loader/loader.go index ed1918c8fb..4d59fdf530 100644 --- a/pkg/config/loader/loader.go +++ b/pkg/config/loader/loader.go @@ -415,10 +415,11 @@ func (c *ConfigLoader) InitDataLayer() (res *database.Layer, err error) { return cleanupV1() }, - Pool: pool, - DDLPool: ddlPool, - V1: v1, - Seed: cf.Seed, + Pool: pool, + DirectDatabaseURL: databaseUrl, + DDLPool: ddlPool, + V1: v1, + Seed: cf.Seed, }, nil } @@ -475,9 +476,13 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers } var mqv1 msgqueue.MessageQueue + var pubsubv1 msgqueue.PubSub cleanup1 := func() error { return nil } + cleanupPubSub := func() error { + return nil + } var ing ingestor.Ingestor @@ -510,7 +515,6 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers rabbitmq.WithURL(cf.MessageQueue.RabbitMQ.URL), rabbitmq.WithLogger(&l), rabbitmq.WithQos(cf.MessageQueue.RabbitMQ.Qos), - rabbitmq.WithDisableTenantExchangePubs(cf.Runtime.DisableTenantPubs), rabbitmq.WithMaxPubChannels(cf.MessageQueue.RabbitMQ.MaxPubChans), rabbitmq.WithMaxSubChannels(cf.MessageQueue.RabbitMQ.MaxSubChans), rabbitmq.WithGzipCompression( @@ -529,8 +533,91 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers } } + // The best-effort pub/sub is constructed with strictly disjoint pooled + // resources: its own AMQP connections/channel pools, or its own small + // pgx pool on the direct database URL. + pubsubKind, pubsubURL := resolvePubSubKindAndURL(cf) + + 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) + } + + pubsubv1 = msgqueue.NewGatedPubSub(pubsubv1, cf.Runtime.DisableTenantPubs) + ing, err = ingestor.NewIngestor( ingestor.WithMessageQueueV1(mqv1), + ingestor.WithPubSub(pubsubv1), ingestor.WithRepositoryV1(dc.V1), ) @@ -853,6 +940,10 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers return fmt.Errorf("error cleaning up rabbitmq: %w", err) } + if err := cleanupPubSub(); err != nil { + return fmt.Errorf("error cleaning up pubsub: %w", err) + } + if closeErr := analyticsEmitter.Close(); closeErr != nil { l.Error().Err(closeErr).Msg("error closing analytics emitter") } @@ -903,6 +994,7 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers Encryption: encryptionSvc, Layer: dc, MessageQueueV1: mqv1, + PubSubV1: pubsubv1, Services: services, PausedControllers: pausedControllers, InternalClientFactory: internalClientFactory, @@ -931,6 +1023,27 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers }, nil } +// resolvePubSubKindAndURL resolves the pub/sub kind and rabbit URL, inheriting +// from the resolved durable msgQueue settings when unset. This is the config +// backwards-compatibility invariant: a deployment configured only with today's +// durable variables (including the legacy SERVER_TASKQUEUE_* aliases) gets a +// fully working pub/sub path with zero new configuration. +func resolvePubSubKindAndURL(cf *server.ServerConfigFile) (kind string, rabbitURL string) { + kind = cf.MessageQueue.PubSub.Kind + + if kind == "" { + kind = cf.MessageQueue.Kind + } + + rabbitURL = cf.MessageQueue.PubSub.RabbitMQ.URL + + if rabbitURL == "" { + rabbitURL = cf.MessageQueue.RabbitMQ.URL + } + + return kind, rabbitURL +} + func getStrArr(v string) []string { return strings.Split(v, " ") } diff --git a/pkg/config/loader/loader_pubsub_test.go b/pkg/config/loader/loader_pubsub_test.go new file mode 100644 index 0000000000..db0ff977ef --- /dev/null +++ b/pkg/config/loader/loader_pubsub_test.go @@ -0,0 +1,88 @@ +package loader + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestPubSubSettingsInheritance checks the config backwards-compatibility +// invariant: a deployment configured only with today's durable message queue +// variables (including the legacy SERVER_TASKQUEUE_* aliases) resolves to a +// fully configured pub/sub path with zero new variables. +func TestPubSubSettingsInheritance(t *testing.T) { + cases := []struct { + name string + env map[string]string + wantKind string + wantURL string + wantMaxPub int32 + wantMaxSub int32 + }{ + { + name: "modern rabbit env only", + env: map[string]string{ + "SERVER_MSGQUEUE_KIND": "rabbitmq", + "SERVER_MSGQUEUE_RABBITMQ_URL": "amqp://user:password@rabbit:5672/", + }, + wantKind: "rabbitmq", + wantURL: "amqp://user:password@rabbit:5672/", + wantMaxPub: 10, + wantMaxSub: 20, + }, + { + name: "legacy taskqueue aliases only", + env: map[string]string{ + "SERVER_TASKQUEUE_KIND": "rabbitmq", + "SERVER_TASKQUEUE_RABBITMQ_URL": "amqp://legacy:password@rabbit:5672/", + }, + wantKind: "rabbitmq", + wantURL: "amqp://legacy:password@rabbit:5672/", + wantMaxPub: 10, + wantMaxSub: 20, + }, + { + name: "postgres durable only", + env: map[string]string{ + "SERVER_MSGQUEUE_KIND": "postgres", + }, + wantKind: "postgres", + wantURL: "", + wantMaxPub: 10, + wantMaxSub: 20, + }, + { + name: "explicit pubsub overrides", + env: map[string]string{ + "SERVER_MSGQUEUE_KIND": "postgres", + "SERVER_MSGQUEUE_PUBSUB_KIND": "rabbitmq", + "SERVER_MSGQUEUE_PUBSUB_RABBITMQ_URL": "amqp://pubsub:password@rabbit:5672/", + "SERVER_MSGQUEUE_PUBSUB_RABBITMQ_MAX_PUB_CHANS": "3", + "SERVER_MSGQUEUE_PUBSUB_RABBITMQ_MAX_SUB_CHANS": "7", + }, + wantKind: "rabbitmq", + wantURL: "amqp://pubsub:password@rabbit:5672/", + wantMaxPub: 3, + wantMaxSub: 7, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + for k, v := range tc.env { + t.Setenv(k, v) + } + + cf, err := LoadServerConfigFile() + require.NoError(t, err) + + kind, url := resolvePubSubKindAndURL(cf) + + assert.Equal(t, tc.wantKind, kind) + assert.Equal(t, tc.wantURL, url) + assert.Equal(t, tc.wantMaxPub, cf.MessageQueue.PubSub.RabbitMQ.MaxPubChans) + assert.Equal(t, tc.wantMaxSub, cf.MessageQueue.PubSub.RabbitMQ.MaxSubChans) + }) + } +} diff --git a/pkg/config/server/server.go b/pkg/config/server/server.go index f48c9cf76f..eb86119f9e 100644 --- a/pkg/config/server/server.go +++ b/pkg/config/server/server.go @@ -511,6 +511,37 @@ type MessageQueueConfigFile struct { Postgres PostgresMQConfigFile `mapstructure:"postgres" json:"postgres,omitempty"` RabbitMQ RabbitMQConfigFile `mapstructure:"rabbitmq" json:"rabbitmq,omitempty" validate:"required"` + + PubSub PubSubConfigFile `mapstructure:"pubSub" json:"pubSub,omitempty"` +} + +// PubSubConfigFile configures the best-effort pub/sub mechanism. All settings +// are optional overrides which inherit from the durable message queue settings +// when unset, so existing deployments need zero new configuration. +type PubSubConfigFile struct { + // Kind is "rabbitmq" or "postgres"; empty inherits msgQueue.kind + Kind string `mapstructure:"kind" json:"kind,omitempty" validate:"omitempty,oneof=rabbitmq postgres"` + + RabbitMQ PubSubRabbitMQConfigFile `mapstructure:"rabbitmq" json:"rabbitmq,omitempty"` + + Postgres PubSubPostgresConfigFile `mapstructure:"postgres" json:"postgres,omitempty"` +} + +type PubSubRabbitMQConfigFile struct { + // URL is the connection URL; empty inherits msgQueue.rabbitmq.url. The + // pub/sub always opens its own connections, even when the durable queue is + // also rabbitmq on the same URL. + URL string `mapstructure:"url" json:"url,omitempty"` + + MaxPubChans int32 `mapstructure:"maxPubChans" json:"maxPubChans,omitempty" default:"10"` + MaxSubChans int32 `mapstructure:"maxSubChans" json:"maxSubChans,omitempty" default:"20"` +} + +type PubSubPostgresConfigFile struct { + // The pub/sub uses a small dedicated pool built from the direct DATABASE_URL + // (never pgbouncer — LISTEN does not survive transaction pooling). + MaxConns int32 `mapstructure:"maxConns" json:"maxConns,omitempty" default:"5"` + MinConns int32 `mapstructure:"minConns" json:"minConns,omitempty" default:"1"` } type PostgresMQConfigFile struct { @@ -645,6 +676,8 @@ type ServerConfig struct { MessageQueueV1 msgqueue.MessageQueue + PubSubV1 msgqueue.PubSub + Logger *zerolog.Logger AdditionalLoggers ConfigFileAdditionalLoggers @@ -848,6 +881,14 @@ func BindAllEnv(v *viper.Viper) { // throughput options _ = v.BindEnv("msgQueue.rabbitmq.qos", "SERVER_MSGQUEUE_RABBITMQ_QOS") + + // pub/sub (all optional: inherit from the durable msgQueue settings when unset) + _ = v.BindEnv("msgQueue.pubSub.kind", "SERVER_MSGQUEUE_PUBSUB_KIND") + _ = v.BindEnv("msgQueue.pubSub.rabbitmq.url", "SERVER_MSGQUEUE_PUBSUB_RABBITMQ_URL") + _ = v.BindEnv("msgQueue.pubSub.rabbitmq.maxPubChans", "SERVER_MSGQUEUE_PUBSUB_RABBITMQ_MAX_PUB_CHANS") + _ = v.BindEnv("msgQueue.pubSub.rabbitmq.maxSubChans", "SERVER_MSGQUEUE_PUBSUB_RABBITMQ_MAX_SUB_CHANS") + _ = v.BindEnv("msgQueue.pubSub.postgres.maxConns", "SERVER_MSGQUEUE_PUBSUB_POSTGRES_MAX_CONNS") + _ = v.BindEnv("msgQueue.pubSub.postgres.minConns", "SERVER_MSGQUEUE_PUBSUB_POSTGRES_MIN_CONNS") _ = v.BindEnv("runtime.singleQueueLimit", "SERVER_SINGLE_QUEUE_LIMIT") _ = v.BindEnv("runtime.optimisticSchedulingEnabled", "SERVER_OPTIMISTIC_SCHEDULING_ENABLED") _ = v.BindEnv("runtime.optimisticSchedulingSlots", "SERVER_OPTIMISTIC_SCHEDULING_SLOTS") diff --git a/pkg/repository/mq.go b/pkg/repository/mq.go index 78b0fd9a56..861d4bd43b 100644 --- a/pkg/repository/mq.go +++ b/pkg/repository/mq.go @@ -10,6 +10,8 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/rs/zerolog" "github.com/hatchet-dev/hatchet/pkg/repository/sqlchelpers" "github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1" @@ -56,6 +58,18 @@ func newMessageQueueRepository(shared *sharedRepository) (*messageQueueRepositor } } +// NewMessageQueueRepositoryWithPool constructs a standalone MessageQueueRepository +// over a dedicated pool, with its own multiplexed listener whose LISTEN +// connection is hijacked from that pool. It is used by the best-effort PubSub, +// which must never share pooled resources with the repository layer. +func NewMessageQueueRepositoryWithPool(l *zerolog.Logger, pool *pgxpool.Pool) (MessageQueueRepository, func() error) { + return newMessageQueueRepository(&sharedRepository{ + pool: pool, + l: l, + queries: sqlcv1.New(), + }) +} + func (m *messageQueueRepository) Listen(ctx context.Context, name string, f func(ctx context.Context, notification *PubSubMessage) error) error { return m.m.listen(ctx, name, f) } diff --git a/pkg/repository/multiplexer.go b/pkg/repository/multiplexer.go index b02c237207..81a4cdc962 100644 --- a/pkg/repository/multiplexer.go +++ b/pkg/repository/multiplexer.go @@ -62,7 +62,10 @@ func (m *multiplexedListener) startListening() { if err != nil { return nil, err } - 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 }, LogError: func(innerCtx context.Context, err error) { m.l.Warn().Err(err).Msg("error in listener") From 69c3092d0c8192a033bb0a67ad0683a3c962937c Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Wed, 22 Jul 2026 10:34:27 -0700 Subject: [PATCH 2/6] address review: topic kind enum, MsgHandler rename, PubTenantMessage 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 --- internal/msgqueue/mq_pub_buffer_test.go | 2 +- internal/msgqueue/mq_sub_buffer.go | 4 +- internal/msgqueue/msgqueue.go | 7 +- internal/msgqueue/postgres/msgqueue.go | 2 +- internal/msgqueue/postgres/pubsub.go | 2 +- internal/msgqueue/pubsub.go | 77 ++++++-- internal/msgqueue/pubsub_test.go | 90 +++++++++- internal/msgqueue/rabbitmq/pubsub.go | 8 +- internal/msgqueue/rabbitmq/rabbitmq.go | 8 +- internal/msgqueue/shared_tenant_reader.go | 10 +- .../services/controllers/olap/controller.go | 8 +- .../olap/process_dag_status_updates.go | 6 +- .../olap/process_task_status_updates.go | 6 +- .../controllers/olap/signal/signal.go | 6 +- .../services/controllers/task/controller.go | 11 +- internal/services/dispatcher/dispatcher.go | 8 +- internal/services/dispatcher/server.go | 16 +- internal/services/ingestor/server_v1.go | 4 +- internal/services/scheduler/v1/scheduler.go | 45 +---- .../ticker/schedule_workflow_v1_test.go | 2 +- pkg/config/loader/loader.go | 167 +++++++++--------- 21 files changed, 291 insertions(+), 198 deletions(-) diff --git a/internal/msgqueue/mq_pub_buffer_test.go b/internal/msgqueue/mq_pub_buffer_test.go index 91d9617ffa..de3f6ad8f9 100644 --- a/internal/msgqueue/mq_pub_buffer_test.go +++ b/internal/msgqueue/mq_pub_buffer_test.go @@ -22,7 +22,7 @@ func (m *mockMessageQueue) SendMessage(ctx context.Context, q Queue, msg *Messag } return nil } -func (m *mockMessageQueue) Subscribe(_ Queue, _ AckHook, _ AckHook) (func() error, error) { +func (m *mockMessageQueue) Subscribe(_ Queue, _ MsgHandler, _ MsgHandler) (func() error, error) { return func() error { return nil }, nil } func (m *mockMessageQueue) IsReady() bool { return true } diff --git a/internal/msgqueue/mq_sub_buffer.go b/internal/msgqueue/mq_sub_buffer.go index cc7603715e..e20bfe6f4f 100644 --- a/internal/msgqueue/mq_sub_buffer.go +++ b/internal/msgqueue/mq_sub_buffer.go @@ -41,7 +41,7 @@ const ( // SubscribeFunc subscribes to a message source with pre-ack and post-ack // hooks, returning a cleanup function. -type SubscribeFunc func(preAck AckHook, postAck AckHook) (func() error, error) +type SubscribeFunc func(preAck MsgHandler, postAck MsgHandler) (func() error, error) // MQSubBuffer buffers messages coming out of the task queue, groups them by tenantId and msgId, and then flushes them // to the task handler as necessary. @@ -103,7 +103,7 @@ func defaultMQSubBufferOpts() *mqSubBufferOpts { } func NewMQSubBuffer(queue Queue, mq MessageQueue, dst DstFunc, fs ...mqSubBufferOptFunc) *MQSubBuffer { - return NewSubBufferFromSubscribe(func(preAck AckHook, postAck AckHook) (func() error, error) { + return NewSubBufferFromSubscribe(func(preAck MsgHandler, postAck MsgHandler) (func() error, error) { return mq.Subscribe(queue, preAck, postAck) }, dst, fs...) } diff --git a/internal/msgqueue/msgqueue.go b/internal/msgqueue/msgqueue.go index 375ea9b3f7..861ac46a71 100644 --- a/internal/msgqueue/msgqueue.go +++ b/internal/msgqueue/msgqueue.go @@ -151,7 +151,10 @@ func QueueTypeFromDispatcherID(d uuid.UUID) dispatcherQueue { return dispatcherQueue(d.String() + "_dispatcher_v1") } -type AckHook func(task *Message) error +// MsgHandler processes a received message. On the durable MessageQueue it is +// invoked as a pre-ack or post-ack hook; on the best-effort PubSub it is the +// sole handler with no ack semantics. +type MsgHandler func(task *Message) error type MessageQueue interface { // Clone copies the message queue with a new instance. @@ -165,7 +168,7 @@ type MessageQueue interface { // Subscribe subscribes to the task queue. It returns a cleanup function that should be called when the // subscription is no longer needed. - Subscribe(queue Queue, preAck AckHook, postAck AckHook) (func() error, error) + Subscribe(queue Queue, preAck MsgHandler, postAck MsgHandler) (func() error, error) // IsReady returns true if the task queue is ready to accept tasks. IsReady() bool diff --git a/internal/msgqueue/postgres/msgqueue.go b/internal/msgqueue/postgres/msgqueue.go index ec61b84e7a..5708379ce6 100644 --- a/internal/msgqueue/postgres/msgqueue.go +++ b/internal/msgqueue/postgres/msgqueue.go @@ -167,7 +167,7 @@ func (p *PostgresMessageQueue) addMessage(ctx context.Context, queue msgqueue.Qu return nil } -func (p *PostgresMessageQueue) Subscribe(queue msgqueue.Queue, preAck msgqueue.AckHook, postAck msgqueue.AckHook) (func() error, error) { +func (p *PostgresMessageQueue) Subscribe(queue msgqueue.Queue, preAck msgqueue.MsgHandler, postAck msgqueue.MsgHandler) (func() error, error) { err := p.upsertQueue(context.Background(), queue) if err != nil { diff --git a/internal/msgqueue/postgres/pubsub.go b/internal/msgqueue/postgres/pubsub.go index 0347539024..340ab29b9b 100644 --- a/internal/msgqueue/postgres/pubsub.go +++ b/internal/msgqueue/postgres/pubsub.go @@ -121,7 +121,7 @@ func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Me // Sub subscribes to a topic. Inline NOTIFY payloads are handled directly; // non-JSON notifications and a 1s ticker wake a poll that drains >8KB fallback // rows. Delivery is at-most-once: handler errors are logged, never redelivered. -func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.AckHook) (func() error, error) { +func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.MsgHandler) (func() error, error) { err := p.ensureQueue(context.Background(), topic) if err != nil { diff --git a/internal/msgqueue/pubsub.go b/internal/msgqueue/pubsub.go index 5267213aed..c8d3cd1c45 100644 --- a/internal/msgqueue/pubsub.go +++ b/internal/msgqueue/pubsub.go @@ -5,22 +5,32 @@ import ( "fmt" "github.com/google/uuid" + "github.com/rs/zerolog" +) + +// TopicKind identifies the kind of pub/sub destination a Topic refers to. +type TopicKind string + +const ( + // TopicKindTenantStream feeds the dispatcher's per-tenant gRPC streams. + TopicKindTenantStream TopicKind = "tenant-stream" + + // TopicKindSchedulerPartition wakes the scheduler owning a partition. + TopicKindSchedulerPartition TopicKind = "scheduler-partition" ) // Topic identifies a best-effort pub/sub destination. type Topic struct { - name string - tenantStream bool + name string + kind TopicKind } func (t Topic) Name() string { return t.name } -// IsTenantStream returns true if this topic feeds the dispatcher's per-tenant -// gRPC streams (as opposed to a scheduler partition wake-up). -func (t Topic) IsTenantStream() bool { - return t.tenantStream +func (t Topic) Kind() TopicKind { + return t.kind } // TenantTopic feeds the dispatcher's gRPC streams. The wire name matches the @@ -28,8 +38,8 @@ func (t Topic) IsTenantStream() bool { // fleets interoperate. func TenantTopic(tenantId uuid.UUID) Topic { return Topic{ - name: tenantId.String() + "_v1", - tenantStream: true, + name: tenantId.String() + "_v1", + kind: TopicKindTenantStream, } } @@ -39,6 +49,7 @@ func TenantTopic(tenantId uuid.UUID) Topic { func SchedulerPartitionTopic(partitionId string) Topic { return Topic{ name: fmt.Sprintf("%s_scheduler_v1", partitionId), + kind: TopicKindSchedulerPartition, } } @@ -57,12 +68,56 @@ type PubSub interface { // 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) + Sub(topic Topic, handler MsgHandler) (func() error, error) // IsReady returns true if the pub/sub mechanism is ready to accept messages. IsReady() bool } +// tenantStreamMsgIDs enumerates the message IDs the dispatcher's gRPC streams +// consume from tenant topics (see msgsToWorkflowEvent and +// isMatchingWorkflowRunV1 in the dispatcher). Publishing any other ID to a +// tenant topic is pure waste: the streams are the only consumers. +var tenantStreamMsgIDs = map[string]struct{}{ + MsgIDCreatedTask: {}, + MsgIDTaskCompleted: {}, + MsgIDTaskFailed: {}, + MsgIDTaskCancelled: {}, + MsgIDTaskStreamEvent: {}, + MsgIDWorkflowRunFinished: {}, + MsgIDWorkflowRunFinishedCandidate: {}, +} + +// PubTenantMessage writes a tenant-scoped message to its destinations: a +// durable send to queue via mq (when queue is non-nil), plus a publish to the +// tenant stream when the message ID is one the dispatcher's streams consume. +// +// Durable send errors are returned. Stream publish errors are best-effort +// (logged, never propagated) when the message also had a durable destination; +// when the stream is the message's only delivery path (queue == nil), the +// publish error is returned so callers can decide. +func PubTenantMessage(ctx context.Context, l *zerolog.Logger, mq MessageQueue, ps PubSub, queue Queue, msg *Message) error { + if queue != nil { + if err := mq.SendMessage(ctx, queue, msg); err != nil { + return err + } + } + + if _, ok := tenantStreamMsgIDs[msg.ID]; !ok || msg.TenantID == uuid.Nil { + return nil + } + + if err := ps.Pub(ctx, TenantTopic(msg.TenantID), msg); err != nil { + if queue == nil { + return err + } + + l.Warn().Ctx(ctx).Err(err).Str("message_id", msg.ID).Msg("could not publish message to tenant stream") + } + + return nil +} + // gatedPubSub suppresses tenant-stream publishes when disabled, except // task-stream-event (worker stream output must always flow — it has no other // delivery path). Scheduler partition wake-ups are never gated. @@ -86,14 +141,14 @@ func NewGatedPubSub(inner PubSub, disableTenantStreamPubs bool) PubSub { } func (g *gatedPubSub) Pub(ctx context.Context, topic Topic, msg *Message) error { - if g.disableTenantStreamPubs && topic.IsTenantStream() && msg.ID != MsgIDTaskStreamEvent { + if g.disableTenantStreamPubs && topic.Kind() == TopicKindTenantStream && msg.ID != MsgIDTaskStreamEvent { return nil } return g.inner.Pub(ctx, topic, msg) } -func (g *gatedPubSub) Sub(topic Topic, handler AckHook) (func() error, error) { +func (g *gatedPubSub) Sub(topic Topic, handler MsgHandler) (func() error, error) { return g.inner.Sub(topic, handler) } diff --git a/internal/msgqueue/pubsub_test.go b/internal/msgqueue/pubsub_test.go index 7e956023ab..8f75490b51 100644 --- a/internal/msgqueue/pubsub_test.go +++ b/internal/msgqueue/pubsub_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/google/uuid" + "github.com/rs/zerolog" "github.com/stretchr/testify/assert" ) @@ -13,11 +14,11 @@ func TestTopicNames(t *testing.T) { tenantTopic := TenantTopic(tenantId) assert.Equal(t, "707d0855-80ab-4e1f-a156-f1c4546cbf52_v1", tenantTopic.Name()) - assert.True(t, tenantTopic.IsTenantStream()) + assert.Equal(t, TopicKindTenantStream, tenantTopic.Kind()) schedulerTopic := SchedulerPartitionTopic("mypartition") assert.Equal(t, "mypartition_scheduler_v1", schedulerTopic.Name()) - assert.False(t, schedulerTopic.IsTenantStream()) + assert.Equal(t, TopicKindSchedulerPartition, schedulerTopic.Kind()) } type recordingPubSub struct { @@ -29,7 +30,7 @@ func (r *recordingPubSub) Pub(ctx context.Context, topic Topic, msg *Message) er return nil } -func (r *recordingPubSub) Sub(topic Topic, handler AckHook) (func() error, error) { +func (r *recordingPubSub) Sub(topic Topic, handler MsgHandler) (func() error, error) { return func() error { return nil }, nil } @@ -37,6 +38,89 @@ func (r *recordingPubSub) IsReady() bool { return true } +func TestPubTenantMessage(t *testing.T) { + tenantId := uuid.New() + l := zerolog.Nop() + + t.Run("stream-consumed id is published", func(t *testing.T) { + inner := &recordingPubSub{} + + err := PubTenantMessage(context.Background(), &l, nil, inner, nil, &Message{ID: MsgIDTaskCompleted, TenantID: tenantId}) + assert.NoError(t, err) + assert.Len(t, inner.pubs, 1) + }) + + t.Run("non-stream id is not published", func(t *testing.T) { + inner := &recordingPubSub{} + + err := PubTenantMessage(context.Background(), &l, nil, inner, nil, &Message{ID: MsgIDTaskTrigger, TenantID: tenantId}) + assert.NoError(t, err) + assert.Empty(t, inner.pubs) + }) + + t.Run("nil tenant id is not published", func(t *testing.T) { + inner := &recordingPubSub{} + + err := PubTenantMessage(context.Background(), &l, nil, inner, nil, &Message{ID: MsgIDTaskCompleted, TenantID: uuid.Nil}) + assert.NoError(t, err) + assert.Empty(t, inner.pubs) + }) + + t.Run("durable send happens before stream publish", func(t *testing.T) { + inner := &recordingPubSub{} + var sent []string + mq := &mockMessageQueue{sendMessageFn: func(ctx context.Context, q Queue, msg *Message) error { + sent = append(sent, q.Name()+"/"+msg.ID) + return nil + }} + + err := PubTenantMessage(context.Background(), &l, mq, inner, TASK_PROCESSING_QUEUE, &Message{ID: MsgIDTaskFailed, TenantID: tenantId}) + assert.NoError(t, err) + assert.Len(t, sent, 1) + assert.Len(t, inner.pubs, 1) + }) + + t.Run("durable error is returned and skips the stream publish", func(t *testing.T) { + inner := &recordingPubSub{} + mq := &mockMessageQueue{sendMessageFn: func(ctx context.Context, q Queue, msg *Message) error { + return assert.AnError + }} + + err := PubTenantMessage(context.Background(), &l, mq, inner, TASK_PROCESSING_QUEUE, &Message{ID: MsgIDTaskFailed, TenantID: tenantId}) + assert.Error(t, err) + assert.Empty(t, inner.pubs) + }) + + t.Run("stream publish error is swallowed when durably sent", func(t *testing.T) { + inner := &erroringPubSub{} + mq := &mockMessageQueue{} + + err := PubTenantMessage(context.Background(), &l, mq, inner, TASK_PROCESSING_QUEUE, &Message{ID: MsgIDTaskFailed, TenantID: tenantId}) + assert.NoError(t, err) + }) + + t.Run("stream publish error is returned when it is the only delivery path", func(t *testing.T) { + inner := &erroringPubSub{} + + err := PubTenantMessage(context.Background(), &l, nil, inner, nil, &Message{ID: MsgIDTaskStreamEvent, TenantID: tenantId}) + assert.Error(t, err) + }) +} + +type erroringPubSub struct{} + +func (e *erroringPubSub) Pub(ctx context.Context, topic Topic, msg *Message) error { + return assert.AnError +} + +func (e *erroringPubSub) Sub(topic Topic, handler MsgHandler) (func() error, error) { + return func() error { return nil }, nil +} + +func (e *erroringPubSub) IsReady() bool { + return true +} + func TestGatedPubSub(t *testing.T) { tenantId := uuid.New() diff --git a/internal/msgqueue/rabbitmq/pubsub.go b/internal/msgqueue/rabbitmq/pubsub.go index f4b3ec72d0..e42fc6acaa 100644 --- a/internal/msgqueue/rabbitmq/pubsub.go +++ b/internal/msgqueue/rabbitmq/pubsub.go @@ -255,7 +255,7 @@ func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Me exchange := "" routingKey := topic.Name() - if topic.IsTenantStream() { + if topic.Kind() == msgqueue.TopicKindTenantStream { // tenant streams ride the per-tenant fanout exchange; declare it lazily if _, ok := p.exchangeCache.Get(topic.Name()); !ok { if err := p.declareExchange(pub, topic.Name()); err != nil { @@ -290,7 +290,7 @@ func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Me // Sub subscribes to a topic. Delivery is at-most-once: messages are acked // before the handler runs, and handler errors are logged, never redelivered. -func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.AckHook) (func() error, error) { +func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.MsgHandler) (func() error, error) { ctx, cancel := context.WithCancel(p.ctx) p.l.Debug().Msgf("subscribing to topic: %s", topic.Name()) @@ -425,7 +425,7 @@ func (p *PubSub) declareSubQueue(ch *amqp.Channel, topic msgqueue.Topic) (string name := topic.Name() - if topic.IsTenantStream() { + if topic.Kind() == msgqueue.TopicKindTenantStream { if err := p.declareExchange(ch, topic.Name()); err != nil { return "", err } @@ -445,7 +445,7 @@ func (p *PubSub) declareSubQueue(ch *amqp.Channel, topic msgqueue.Topic) (string return "", err } - if topic.IsTenantStream() { + if topic.Kind() == msgqueue.TopicKindTenantStream { p.l.Debug().Msgf("binding queue: %s to exchange: %s", name, topic.Name()) if err := ch.QueueBind(name, "", topic.Name(), false, nil); err != nil { diff --git a/internal/msgqueue/rabbitmq/rabbitmq.go b/internal/msgqueue/rabbitmq/rabbitmq.go index 4913ce593c..338abec613 100644 --- a/internal/msgqueue/rabbitmq/rabbitmq.go +++ b/internal/msgqueue/rabbitmq/rabbitmq.go @@ -459,8 +459,8 @@ func (t *MessageQueueImpl) pubMessage(ctx context.Context, q msgqueue.Queue, msg // Subscribe subscribes to the msg queue. func (t *MessageQueueImpl) Subscribe( q msgqueue.Queue, - preAck msgqueue.AckHook, - postAck msgqueue.AckHook, + preAck msgqueue.MsgHandler, + postAck msgqueue.MsgHandler, ) (func() error, error) { ctx, cancel := context.WithCancel(context.Background()) @@ -571,8 +571,8 @@ func (t *MessageQueueImpl) subscribe( ctx context.Context, subId string, q msgqueue.Queue, - preAck msgqueue.AckHook, - postAck msgqueue.AckHook, + preAck msgqueue.MsgHandler, + postAck msgqueue.MsgHandler, ) (func() error, error) { sessionCount := 0 diff --git a/internal/msgqueue/shared_tenant_reader.go b/internal/msgqueue/shared_tenant_reader.go index 64c77512cf..83edf370d6 100644 --- a/internal/msgqueue/shared_tenant_reader.go +++ b/internal/msgqueue/shared_tenant_reader.go @@ -11,7 +11,7 @@ import ( ) type sharedTenantSub struct { - fs *syncx.Map[int, AckHook] + fs *syncx.Map[int, MsgHandler] counter int isRunning bool mu sync.Mutex @@ -30,9 +30,9 @@ func NewSharedTenantReader(ps PubSub) *SharedTenantReader { } } -func (s *SharedTenantReader) Subscribe(tenantId uuid.UUID, postAck AckHook) (func() error, error) { +func (s *SharedTenantReader) Subscribe(tenantId uuid.UUID, postAck MsgHandler) (func() error, error) { t, _ := s.tenants.LoadOrStore(tenantId, &sharedTenantSub{ - fs: &syncx.Map[int, AckHook]{}, + fs: &syncx.Map[int, MsgHandler]{}, }) t.mu.Lock() @@ -50,7 +50,7 @@ func (s *SharedTenantReader) Subscribe(tenantId uuid.UUID, postAck AckHook) (fun cleanupSingleSub, err := s.ps.Sub(TenantTopic(tenantId), func(task *Message) error { var innerErr error - t.fs.Range(func(key int, f AckHook) bool { + t.fs.Range(func(key int, f MsgHandler) bool { if err := f(task); err != nil { innerErr = multierror.Append(innerErr, err) } @@ -128,7 +128,7 @@ func (s *SharedBufferedTenantReader) Subscribe(tenantId uuid.UUID, f DstFunc) (f // the buffer runs in PostAck mode, which only uses the post hook, so the // single-handler pubsub Sub maps cleanly onto the subscribe function - subBuffer := NewSubBufferFromSubscribe(func(preAck AckHook, postAck AckHook) (func() error, error) { + subBuffer := NewSubBufferFromSubscribe(func(preAck MsgHandler, postAck MsgHandler) (func() error, error) { return s.ps.Sub(TenantTopic(tenantId), postAck) }, func(tenantId uuid.UUID, msgId string, payloads [][]byte) error { var innerErr error diff --git a/internal/services/controllers/olap/controller.go b/internal/services/controllers/olap/controller.go index ab91ee0018..8772c803fb 100644 --- a/internal/services/controllers/olap/controller.go +++ b/internal/services/controllers/olap/controller.go @@ -1171,15 +1171,9 @@ func (tc *OLAPControllerImpl) republishCreatedTasks(ctx context.Context, tenantI if err != nil { return err } - if err := tc.mq.SendMessage(ctx, msgqueue.OLAP_QUEUE, msg); err != nil { + if err := msgqueue.PubTenantMessage(ctx, tc.l, tc.mq, tc.pubsub, msgqueue.OLAP_QUEUE, msg); err != nil { return err } - - // best-effort publish to the tenant stream, mirroring the original - // created-task publish this requeue stands in for - if err := tc.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { - tc.l.Warn().Ctx(ctx).Err(err).Msg("could not publish created-task to tenant stream") - } } return nil } diff --git a/internal/services/controllers/olap/process_dag_status_updates.go b/internal/services/controllers/olap/process_dag_status_updates.go index e007ac63e6..e69fcfac55 100644 --- a/internal/services/controllers/olap/process_dag_status_updates.go +++ b/internal/services/controllers/olap/process_dag_status_updates.go @@ -103,9 +103,9 @@ func (o *OLAPControllerImpl) notifyDAGsUpdated(ctx context.Context, rows []v1.Up return err } - // best-effort publish to the tenant stream: the dispatcher's workflow - // run subscriptions consume workflow-run-finished - if err := o.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + // fanout-only: the dispatcher's workflow run subscriptions consume + // workflow-run-finished off the tenant stream + if err := msgqueue.PubTenantMessage(ctx, o.l, nil, o.pubsub, nil, msg); err != nil { o.l.Warn().Ctx(ctx).Err(err).Msg("could not publish workflow-run-finished to tenant stream") } } diff --git a/internal/services/controllers/olap/process_task_status_updates.go b/internal/services/controllers/olap/process_task_status_updates.go index 92a778bff4..6fe082b32c 100644 --- a/internal/services/controllers/olap/process_task_status_updates.go +++ b/internal/services/controllers/olap/process_task_status_updates.go @@ -103,9 +103,9 @@ func (o *OLAPControllerImpl) notifyTasksUpdated(ctx context.Context, rows []v1.U return err } - // best-effort publish to the tenant stream: the dispatcher's workflow - // run subscriptions consume workflow-run-finished - if err := o.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + // fanout-only: the dispatcher's workflow run subscriptions consume + // workflow-run-finished off the tenant stream + if err := msgqueue.PubTenantMessage(ctx, o.l, nil, o.pubsub, nil, msg); err != nil { o.l.Warn().Ctx(ctx).Err(err).Msg("could not publish workflow-run-finished to tenant stream") } } diff --git a/internal/services/controllers/olap/signal/signal.go b/internal/services/controllers/olap/signal/signal.go index 20f2011ca6..00ca076e13 100644 --- a/internal/services/controllers/olap/signal/signal.go +++ b/internal/services/controllers/olap/signal/signal.go @@ -122,9 +122,9 @@ func (s *OLAPSignaler) SignalTasksCreated(ctx context.Context, tenantId uuid.UUI continue } - // best-effort publish to the tenant stream: the dispatcher's workflow - // event subscriptions consume created-task - if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { + // the durable write went through the pub buffer above; this publishes the + // stream copy the dispatcher's workflow event subscriptions consume + if err := msgqueue.PubTenantMessage(ctx, s.l, nil, s.pubsub, nil, msg); err != nil { s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish created-task to tenant stream") } } diff --git a/internal/services/controllers/task/controller.go b/internal/services/controllers/task/controller.go index 4d002cd8dc..74fcbcef54 100644 --- a/internal/services/controllers/task/controller.go +++ b/internal/services/controllers/task/controller.go @@ -823,8 +823,11 @@ func (tc *TasksControllerImpl) handleCancelTasks(ctx context.Context, tenantId u return fmt.Errorf("could not create message for task cancellation: %w", err) } - return tc.mq.SendMessage( + return msgqueue.PubTenantMessage( ctx, + tc.l, + tc.mq, + tc.pubsub, msgqueue.TASK_PROCESSING_QUEUE, msg, ) @@ -1037,9 +1040,9 @@ func (tc *TasksControllerImpl) notifyQueuesOnCompletion(ctx context.Context, ten return } - // best-effort publish to the tenant stream: the dispatcher's workflow run - // subscriptions consume workflow-run-finished-candidate - err = tc.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg) + // fanout-only: the dispatcher's workflow run subscriptions consume + // workflow-run-finished-candidate off the tenant stream + err = msgqueue.PubTenantMessage(ctx, tc.l, nil, tc.pubsub, nil, msg) if err != nil { tc.l.Err(err).Ctx(ctx).Msg("could not publish workflow-run-finished-candidate message") diff --git a/internal/services/dispatcher/dispatcher.go b/internal/services/dispatcher/dispatcher.go index b4ebb13f2c..fb402a9d25 100644 --- a/internal/services/dispatcher/dispatcher.go +++ b/internal/services/dispatcher/dispatcher.go @@ -957,18 +957,12 @@ func (d *DispatcherImpl) handleRetries( queueutils.SleepWithExponentialBackoff(100*time.Millisecond, 5*time.Second, int(task.InternalRetryCount)) - err = d.mqv1.SendMessage(retryCtx, msgqueue.TASK_PROCESSING_QUEUE, msg) + err = msgqueue.PubTenantMessage(retryCtx, d.l, d.mqv1, d.pubsub, msgqueue.TASK_PROCESSING_QUEUE, msg) if err != nil { return fmt.Errorf("could not send failed task message: %w", err) } - // best-effort publish to the tenant stream: the dispatcher's workflow - // event subscriptions consume task-failed - if err := d.pubsub.Pub(retryCtx, msgqueue.TenantTopic(tenantId), msg); err != nil { - d.l.Warn().Ctx(retryCtx).Err(err).Msg("could not publish task-failed to tenant stream") - } - return nil }) } diff --git a/internal/services/dispatcher/server.go b/internal/services/dispatcher/server.go index f6ec81d382..386a327c52 100644 --- a/internal/services/dispatcher/server.go +++ b/internal/services/dispatcher/server.go @@ -1347,18 +1347,12 @@ func (s *DispatcherImpl) handleTaskCompleted(inputCtx context.Context, task *sql return nil, err } - err = s.mqv1.SendMessage(inputCtx, msgqueue.TASK_PROCESSING_QUEUE, msg) + err = msgqueue.PubTenantMessage(inputCtx, s.l, s.mqv1, s.pubsub, msgqueue.TASK_PROCESSING_QUEUE, msg) if err != nil { return nil, err } - // best-effort publish to the tenant stream: the dispatcher's workflow event - // subscriptions consume task-completed - if err := s.pubsub.Pub(inputCtx, msgqueue.TenantTopic(tenantId), msg); err != nil { - s.l.Warn().Ctx(inputCtx).Err(err).Msg("could not publish task-completed to tenant stream") - } - resp := &contracts.ActionEventResponse{ TenantId: tenantId.String(), WorkerId: request.WorkerId, @@ -1413,18 +1407,12 @@ func (s *DispatcherImpl) handleTaskFailed(inputCtx context.Context, task *sqlcv1 return nil, err } - err = s.mqv1.SendMessage(inputCtx, msgqueue.TASK_PROCESSING_QUEUE, msg) + err = msgqueue.PubTenantMessage(inputCtx, s.l, s.mqv1, s.pubsub, msgqueue.TASK_PROCESSING_QUEUE, msg) if err != nil { return nil, err } - // best-effort publish to the tenant stream: the dispatcher's workflow event - // subscriptions consume task-failed - if err := s.pubsub.Pub(inputCtx, msgqueue.TenantTopic(tenantId), msg); err != nil { - s.l.Warn().Ctx(inputCtx).Err(err).Msg("could not publish task-failed to tenant stream") - } - return &contracts.ActionEventResponse{ TenantId: tenantId.String(), WorkerId: request.WorkerId, diff --git a/internal/services/ingestor/server_v1.go b/internal/services/ingestor/server_v1.go index df0fe0f3ac..977c5b20d5 100644 --- a/internal/services/ingestor/server_v1.go +++ b/internal/services/ingestor/server_v1.go @@ -51,8 +51,8 @@ func (i *IngestorImpl) putStreamEventV1(ctx context.Context, tenant *sqlcv1.Tena return nil, err } - // stream events are fanout-only: publish straight to the tenant stream - err = i.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg) + // stream events are fanout-only: the tenant stream is their sole delivery path + err = msgqueue.PubTenantMessage(ctx, i.l, nil, i.pubsub, nil, msg) if err != nil { return nil, err diff --git a/internal/services/scheduler/v1/scheduler.go b/internal/services/scheduler/v1/scheduler.go index 75a3e5800c..8ddd7ef6de 100644 --- a/internal/services/scheduler/v1/scheduler.go +++ b/internal/services/scheduler/v1/scheduler.go @@ -617,21 +617,10 @@ func (s *Scheduler) scheduleStepRuns(ctx context.Context, tenantId uuid.UUID, re continue } - err = s.mq.SendMessage( - ctx, - msgqueue.TASK_PROCESSING_QUEUE, - msg, - ) + err = msgqueue.PubTenantMessage(ctx, s.l, s.mq, s.pubsub, msgqueue.TASK_PROCESSING_QUEUE, msg) if err != nil { outerErr = multierror.Append(outerErr, fmt.Errorf("could not send cancelled task: %w", err)) - continue - } - - // best-effort publish to the tenant stream: the dispatcher's workflow - // event subscriptions consume task-cancelled - if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { - s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish task-cancelled to tenant stream") } } } @@ -700,22 +689,12 @@ func (s *Scheduler) internalRetry(ctx context.Context, tenantId uuid.UUID, assig continue } - err = s.mq.SendMessage( - ctx, - msgqueue.TASK_PROCESSING_QUEUE, - msg, - ) + err = msgqueue.PubTenantMessage(ctx, s.l, s.mq, s.pubsub, msgqueue.TASK_PROCESSING_QUEUE, msg) if err != nil { s.l.Error().Ctx(ctx).Err(err).Msg("could not send failed task") continue } - - // best-effort publish to the tenant stream: the dispatcher's workflow - // event subscriptions consume task-failed - if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { - s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish task-failed to tenant stream") - } } } @@ -784,22 +763,12 @@ func (s *Scheduler) notifyAfterConcurrency(ctx context.Context, tenantId uuid.UU continue } - err = s.mq.SendMessage( - ctx, - msgqueue.TASK_PROCESSING_QUEUE, - msg, - ) + err = msgqueue.PubTenantMessage(ctx, s.l, s.mq, s.pubsub, msgqueue.TASK_PROCESSING_QUEUE, msg) if err != nil { s.l.Error().Ctx(ctx).Err(err).Msg("could not send cancelled task") continue } - - // best-effort publish to the tenant stream: the dispatcher's workflow - // event subscriptions consume task-cancelled - if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { - s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish task-cancelled to tenant stream") - } } } @@ -949,7 +918,7 @@ func (s *Scheduler) handleDeadLetteredTaskBulkAssigned(ctx context.Context, msg return fmt.Errorf("could not create failed task message: %w", err) } - err = s.mq.SendMessage(ctx, msgqueue.TASK_PROCESSING_QUEUE, msg) + err = msgqueue.PubTenantMessage(ctx, s.l, s.mq, s.pubsub, msgqueue.TASK_PROCESSING_QUEUE, msg) if err != nil { // NOTE: failure to send on the MQ is likely not transient; ideally we could only retry individual @@ -957,12 +926,6 @@ func (s *Scheduler) handleDeadLetteredTaskBulkAssigned(ctx context.Context, msg // on the downstream `task-failed` processing to be idempotent. return fmt.Errorf("could not send failed task message: %w", err) } - - // best-effort publish to the tenant stream: the dispatcher's workflow - // event subscriptions consume task-failed - if err := s.pubsub.Pub(ctx, msgqueue.TenantTopic(tenantId), msg); err != nil { - s.l.Warn().Ctx(ctx).Err(err).Msg("could not publish task-failed to tenant stream") - } } return nil diff --git a/internal/services/ticker/schedule_workflow_v1_test.go b/internal/services/ticker/schedule_workflow_v1_test.go index 7b498d1aab..422fa94c1b 100644 --- a/internal/services/ticker/schedule_workflow_v1_test.go +++ b/internal/services/ticker/schedule_workflow_v1_test.go @@ -72,7 +72,7 @@ func (m *mockMQ) SendMessage(_ context.Context, _ msgqueue.Queue, _ *msgqueue.Me } func (m *mockMQ) Clone() (func() error, msgqueue.MessageQueue, error) { return nil, nil, nil } func (m *mockMQ) SetQOS(_ int) {} -func (m *mockMQ) Subscribe(_ msgqueue.Queue, _, _ msgqueue.AckHook) (func() error, error) { +func (m *mockMQ) Subscribe(_ msgqueue.Queue, _, _ msgqueue.MsgHandler) (func() error, error) { return nil, nil } func (m *mockMQ) IsReady() bool { return true } diff --git a/pkg/config/loader/loader.go b/pkg/config/loader/loader.go index 4d59fdf530..62180aef4d 100644 --- a/pkg/config/loader/loader.go +++ b/pkg/config/loader/loader.go @@ -533,88 +533,12 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers } } - // The best-effort pub/sub is constructed with strictly disjoint pooled - // resources: its own AMQP connections/channel pools, or its own small - // pgx pool on the direct database URL. - pubsubKind, pubsubURL := resolvePubSubKindAndURL(cf) + cleanupPubSub, pubsubv1, err = createPubSubV1(dc, cf, &l) - 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) + if err != nil { + return nil, nil, err } - pubsubv1 = msgqueue.NewGatedPubSub(pubsubv1, cf.Runtime.DisableTenantPubs) - ing, err = ingestor.NewIngestor( ingestor.WithMessageQueueV1(mqv1), ingestor.WithPubSub(pubsubv1), @@ -1023,6 +947,91 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers }, nil } +// createPubSubV1 constructs the best-effort pub/sub with strictly disjoint +// pooled resources: its own AMQP connections/channel pools, or its own small +// pgx pool on the direct database URL. The result is wrapped in the +// DisableTenantPubs gate. +func createPubSubV1(dc *database.Layer, cf *server.ServerConfigFile, l *zerolog.Logger) (cleanup func() error, ps msgqueue.PubSub, err error) { + pubsubKind, pubsubURL := resolvePubSubKindAndURL(cf) + + 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, pgps, err := pgmq.NewPubSub( + pubsubRepo, + pgmq.WithPubSubLogger(l), + ) + + if err != nil { + pubsubPool.Close() + return nil, nil, fmt.Errorf("could not init postgres pubsub: %w", err) + } + + ps = pgps + cleanup = 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, rmqps, 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) + } + + ps = rmqps + cleanup = cleanupRmq + default: + return nil, nil, fmt.Errorf("invalid pubsub kind %q, must be 'rabbitmq' or 'postgres'", pubsubKind) + } + + return cleanup, msgqueue.NewGatedPubSub(ps, cf.Runtime.DisableTenantPubs), nil +} + // resolvePubSubKindAndURL resolves the pub/sub kind and rabbit URL, inheriting // from the resolved durable msgQueue settings when unset. This is the config // backwards-compatibility invariant: a deployment configured only with today's From a825a7b664eeb202b0e22cd21bef1cf2ce34429c Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Mon, 27 Jul 2026 14:11:05 -0400 Subject: [PATCH 3/6] address review: drop pubsub from health probes, decompress cross-backend 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). --- cmd/hatchet-engine/engine/run.go | 4 +- internal/msgqueue/gzip.go | 42 +++++++++++++++++ internal/msgqueue/postgres/pubsub.go | 14 ++++++ internal/msgqueue/postgres/pubsub_test.go | 55 +++++++++++++++++++++++ internal/msgqueue/rabbitmq/gzip.go | 31 ------------- internal/msgqueue/rabbitmq/gzip_test.go | 4 +- internal/msgqueue/rabbitmq/pubsub.go | 2 +- internal/msgqueue/rabbitmq/rabbitmq.go | 13 ++++-- internal/services/health/health.go | 12 ++--- 9 files changed, 133 insertions(+), 44 deletions(-) create mode 100644 internal/msgqueue/gzip.go diff --git a/cmd/hatchet-engine/engine/run.go b/cmd/hatchet-engine/engine/run.go index 52bb95c4f4..396bfdad36 100644 --- a/cmd/hatchet-engine/engine/run.go +++ b/cmd/hatchet-engine/engine/run.go @@ -160,7 +160,7 @@ func runV0Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. var healthCleanup func() error healthProbes := sc.HasService("health") if healthProbes { - h = health.New(sc.V1.Health(), sc.MessageQueueV1, sc.PubSubV1, sc.Version, l) + h = health.New(sc.V1.Health(), sc.MessageQueueV1, sc.Version, l) cleanupHealth, err := h.Start(sc.Runtime.HealthcheckPort) if err != nil { return fmt.Errorf("could not start health: %w", err) @@ -572,7 +572,7 @@ func runV1Config(ctx context.Context, sc *server.ServerConfig, cleanup *cleanup. var healthCleanup func() error if healthProbes { - h = health.New(sc.V1.Health(), sc.MessageQueueV1, sc.PubSubV1, sc.Version, l) + h = health.New(sc.V1.Health(), sc.MessageQueueV1, sc.Version, l) clean, err := h.Start(sc.Runtime.HealthcheckPort) diff --git a/internal/msgqueue/gzip.go b/internal/msgqueue/gzip.go new file mode 100644 index 0000000000..a50309ef25 --- /dev/null +++ b/internal/msgqueue/gzip.go @@ -0,0 +1,42 @@ +package msgqueue + +import ( + "bytes" + "compress/gzip" + "fmt" + "io" +) + +// DecompressPayloads decompresses gzip-compressed message payloads. It lives +// in the msgqueue package (rather than a single implementation) because a +// compressed message can cross backends: a producer whose durable queue is +// RabbitMQ with compression enabled hands the same *Message to whichever +// pub/sub is configured, so every subscriber must be able to decompress. +func DecompressPayloads(payloads [][]byte) ([][]byte, error) { + if len(payloads) == 0 { + return payloads, nil + } + + decompressed := make([][]byte, len(payloads)) + + for i, payload := range payloads { + reader, err := gzip.NewReader(bytes.NewReader(payload)) + if err != nil { + return nil, fmt.Errorf("failed to create gzip reader for payload %d: %w", i, err) + } + + decompressedData, err := io.ReadAll(reader) + if err != nil { + reader.Close() + return nil, fmt.Errorf("failed to read from gzip reader for payload %d: %w", i, err) + } + + if err := reader.Close(); err != nil { + return nil, fmt.Errorf("failed to close gzip reader for payload %d: %w", i, err) + } + + decompressed[i] = decompressedData + } + + return decompressed, nil +} diff --git a/internal/msgqueue/postgres/pubsub.go b/internal/msgqueue/postgres/pubsub.go index 340ab29b9b..d4c171a3f1 100644 --- a/internal/msgqueue/postgres/pubsub.go +++ b/internal/msgqueue/postgres/pubsub.go @@ -157,6 +157,20 @@ func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.MsgHandler) (func() } } + // compressed messages can cross backends: a producer whose durable queue + // is rabbitmq with compression enabled publishes here with gzip payloads + if task.Compressed { + decompressedPayloads, err := msgqueue.DecompressPayloads(task.Payloads) + + if err != nil { + p.l.Error().Err(err).Msgf("error decompressing payloads for pubsub message %s", task.ID) + return + } + + task.Payloads = decompressedPayloads + task.Compressed = false + } + if err := handler(task); err != nil { p.l.Error().Err(err).Msgf("error handling pubsub message %s", task.ID) } diff --git a/internal/msgqueue/postgres/pubsub_test.go b/internal/msgqueue/postgres/pubsub_test.go index ff3fb89858..6979a53fe9 100644 --- a/internal/msgqueue/postgres/pubsub_test.go +++ b/internal/msgqueue/postgres/pubsub_test.go @@ -3,6 +3,8 @@ package postgres import ( + "bytes" + "compress/gzip" "context" "crypto/rand" "encoding/json" @@ -188,6 +190,59 @@ func TestPostgresPubSubTwoSubscribers(t *testing.T) { require.NoError(t, cleanupSub2()) } +// TestPostgresPubSubCompressedPayload simulates the cross-backend case: a +// producer whose durable queue is RabbitMQ with compression enabled publishes +// a message that already carries gzip payloads (Compressed=true). The postgres +// subscriber must transparently decompress before invoking the handler. +func TestPostgresPubSubCompressedPayload(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + ps, _ := newTestPubSub(t) + + tenantId := uuid.New() + topic := msgqueue.TenantTopic(tenantId) + + original := []byte(`{"key":"value"}`) + + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + _, err := w.Write(original) + require.NoError(t, err) + require.NoError(t, w.Close()) + + msg := &msgqueue.Message{ + ID: "task-completed", + TenantID: tenantId, + Payloads: [][]byte{buf.Bytes()}, + ImmediatelyExpire: true, + Compressed: true, + } + + received := make(chan *msgqueue.Message, 1) + + cleanupSub, err := ps.Sub(topic, func(m *msgqueue.Message) error { + received <- m + return nil + }) + require.NoError(t, err) + + time.Sleep(1 * time.Second) + + require.NoError(t, ps.Pub(ctx, topic, msg)) + + select { + case m := <-received: + require.Len(t, m.Payloads, 1) + assert.Equal(t, original, m.Payloads[0], "payload should be transparently decompressed") + assert.False(t, m.Compressed) + case <-ctx.Done(): + t.Fatal("timed out waiting for compressed message delivery") + } + + require.NoError(t, cleanupSub()) +} + // TestPostgresPubSubDisjointPools asserts the core invariant: a full pub/sub // cycle never acquires from any pool other than the PubSub's own dedicated pool. func TestPostgresPubSubDisjointPools(t *testing.T) { diff --git a/internal/msgqueue/rabbitmq/gzip.go b/internal/msgqueue/rabbitmq/gzip.go index 3ae77b14e5..5faeaa09f9 100644 --- a/internal/msgqueue/rabbitmq/gzip.go +++ b/internal/msgqueue/rabbitmq/gzip.go @@ -4,7 +4,6 @@ import ( "bytes" "compress/gzip" "fmt" - "io" "sync" ) @@ -104,33 +103,3 @@ func (t compressor) compressPayloads(payloads [][]byte) (*CompressionResult, err return result, nil } - -// decompressPayloads decompresses message payloads using gzip. -func (t compressor) decompressPayloads(payloads [][]byte) ([][]byte, error) { - if len(payloads) == 0 { - return payloads, nil - } - - decompressed := make([][]byte, len(payloads)) - - for i, payload := range payloads { - reader, err := gzip.NewReader(bytes.NewReader(payload)) - if err != nil { - return nil, fmt.Errorf("failed to create gzip reader for payload %d: %w", i, err) - } - - decompressedData, err := io.ReadAll(reader) - if err != nil { - reader.Close() - return nil, fmt.Errorf("failed to read from gzip reader for payload %d: %w", i, err) - } - - if err := reader.Close(); err != nil { - return nil, fmt.Errorf("failed to close gzip reader for payload %d: %w", i, err) - } - - decompressed[i] = decompressedData - } - - return decompressed, nil -} diff --git a/internal/msgqueue/rabbitmq/gzip_test.go b/internal/msgqueue/rabbitmq/gzip_test.go index b85ed510e6..51b213ead5 100644 --- a/internal/msgqueue/rabbitmq/gzip_test.go +++ b/internal/msgqueue/rabbitmq/gzip_test.go @@ -6,6 +6,8 @@ import ( "testing" "github.com/stretchr/testify/assert" + + "github.com/hatchet-dev/hatchet/internal/msgqueue" ) func generatePayloads(count, size int) [][]byte { @@ -41,7 +43,7 @@ func TestCompressDecompressRoundtrip(t *testing.T) { assert.True(t, result.WasCompressed, "expected WasCompressed to be true") assert.Equal(t, len(payloads), len(result.Payloads), "expected %d payloads, got %d", len(payloads), len(result.Payloads)) - decompressed, err := mq.decompressPayloads(result.Payloads) + decompressed, err := msgqueue.DecompressPayloads(result.Payloads) assert.NoError(t, err) diff --git a/internal/msgqueue/rabbitmq/pubsub.go b/internal/msgqueue/rabbitmq/pubsub.go index e42fc6acaa..cb22473c05 100644 --- a/internal/msgqueue/rabbitmq/pubsub.go +++ b/internal/msgqueue/rabbitmq/pubsub.go @@ -353,7 +353,7 @@ func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.MsgHandler) (func() } if msg.Compressed { - decompressedPayloads, err := p.decompressPayloads(msg.Payloads) + decompressedPayloads, err := msgqueue.DecompressPayloads(msg.Payloads) if err != nil { p.l.Error().Msgf("error decompressing payloads: %v", err) diff --git a/internal/msgqueue/rabbitmq/rabbitmq.go b/internal/msgqueue/rabbitmq/rabbitmq.go index 338abec613..36def2c374 100644 --- a/internal/msgqueue/rabbitmq/rabbitmq.go +++ b/internal/msgqueue/rabbitmq/rabbitmq.go @@ -282,7 +282,10 @@ func (t *MessageQueueImpl) pubMessage(ctx context.Context, q msgqueue.Queue, msg var compressionResult *CompressionResult - // don't re-compress if the message was already compressed + // don't re-compress if the message was already compressed. We work on a + // shallow copy rather than mutating in place: callers may hand the same + // *Message to a pub/sub on a different backend afterwards (PubTenantMessage), + // and compression is a rabbitmq wire concern if len(msg.Payloads) > 0 && !msg.Compressed { var err error compressionResult, err = t.compressPayloads(msg.Payloads) @@ -292,8 +295,10 @@ func (t *MessageQueueImpl) pubMessage(ctx context.Context, q msgqueue.Queue, msg } if compressionResult.WasCompressed { - msg.Payloads = compressionResult.Payloads - msg.Compressed = true + msgCp := *msg + msgCp.Payloads = compressionResult.Payloads + msgCp.Compressed = true + msg = &msgCp t.l.Debug().Msgf("compressed payloads for message %s: original=%d bytes, compressed=%d bytes, ratio=%.2f%%", msg.ID, compressionResult.OriginalSize, compressionResult.CompressedSize, compressionResult.CompressionRatio*100) @@ -679,7 +684,7 @@ func (t *MessageQueueImpl) subscribe( } if msg.Compressed { - decompressedPayloads, err := t.decompressPayloads(msg.Payloads) + decompressedPayloads, err := msgqueue.DecompressPayloads(msg.Payloads) if err != nil { t.l.Error().Msgf("error decompressing payloads: %v", err) // reject this message diff --git a/internal/services/health/health.go b/internal/services/health/health.go index 4d43db39d7..35a6b41823 100644 --- a/internal/services/health/health.go +++ b/internal/services/health/health.go @@ -21,16 +21,18 @@ type Health struct { repository v1.HealthRepository queue msgqueue.MessageQueue - pubsub msgqueue.PubSub l *zerolog.Logger } -func New(repo v1.HealthRepository, queue msgqueue.MessageQueue, pubsub msgqueue.PubSub, version string, l *zerolog.Logger) *Health { +// New creates a health service over the durable message queue and repository. +// The best-effort PubSub is deliberately NOT part of the probes: 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. +func New(repo v1.HealthRepository, queue msgqueue.MessageQueue, version string, l *zerolog.Logger) *Health { return &Health{ version: version, repository: repo, queue: queue, - pubsub: pubsub, l: l, } } @@ -46,7 +48,7 @@ func (h *Health) Start(port int) (func() error, error) { ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) defer cancel() - queueReady := h.queue.IsReady() && h.pubsub.IsReady() + queueReady := h.queue.IsReady() repositoryReady := h.repository.IsHealthy(ctx) if !queueReady || !repositoryReady { @@ -62,7 +64,7 @@ func (h *Health) Start(port int) (func() error, error) { ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) defer cancel() - queueReady := h.queue.IsReady() && h.pubsub.IsReady() + queueReady := h.queue.IsReady() repositoryReady := h.repository.IsHealthy(ctx) if h.shuttingDown || !queueReady || !repositoryReady { From d199b356df1922b37902b90255f5a218e1b0ce29 Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Mon, 27 Jul 2026 15:52:32 -0400 Subject: [PATCH 4/6] address review: move Compressor to msgqueue, Message.Clone, real postgres IsReady, constants, comment cleanups MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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. --- internal/msgqueue/gzip.go | 102 +++++++++++++++++ internal/msgqueue/{rabbitmq => }/gzip_test.go | 50 ++++----- internal/msgqueue/msg.go | 7 ++ internal/msgqueue/postgres/pubsub.go | 45 ++++++-- internal/msgqueue/rabbitmq/gzip.go | 105 ------------------ internal/msgqueue/rabbitmq/pubsub.go | 39 +++---- internal/msgqueue/rabbitmq/rabbitmq.go | 48 ++++---- .../olap/process_dag_status_updates.go | 2 - .../olap/process_task_status_updates.go | 2 - internal/services/health/health.go | 4 - pkg/config/loader/loader.go | 1 + 11 files changed, 207 insertions(+), 198 deletions(-) rename internal/msgqueue/{rabbitmq => }/gzip_test.go (73%) delete mode 100644 internal/msgqueue/rabbitmq/gzip.go diff --git a/internal/msgqueue/gzip.go b/internal/msgqueue/gzip.go index a50309ef25..d13bdcb2cd 100644 --- a/internal/msgqueue/gzip.go +++ b/internal/msgqueue/gzip.go @@ -5,8 +5,110 @@ import ( "compress/gzip" "fmt" "io" + "sync" ) +// DefaultCompressionThreshold is the payload size above which messages are +// gzip-compressed when compression is enabled. +const DefaultCompressionThreshold = 5 * 1024 // 5KB + +// Compressor holds gzip compression settings for a message queue +// implementation. Compression settings must agree between the durable queue +// and the pub/sub, since both publish to the same tenant topics. +type Compressor struct { + Enabled bool + Threshold int +} + +type CompressionResult struct { + Payloads [][]byte + WasCompressed bool + OriginalSize int + CompressedSize int + + // CompressionRatio is the ratio of compressed size to original size (compressed / original) + CompressionRatio float64 +} + +// gzipWriterPool reuses gzip.Writer instances to avoid repeated allocations. +// No explicit size cap is needed: sync.Pool is self-limiting because the Go +// runtime evicts pooled objects during GC, so the pool cannot grow unbounded. +// In practice the pool size is also bounded by the number of goroutines +// concurrently compressing, which is small for a publish path. +var gzipWriterPool = sync.Pool{ + New: func() any { + return gzip.NewWriter(nil) + }, +} + +func getPayloadSize(payloads [][]byte) int { + totalSize := 0 + for _, payload := range payloads { + totalSize += len(payload) + } + return totalSize +} + +// CompressPayloads compresses message payloads using gzip if they exceed the +// minimum size threshold. Returns compression results including the compressed +// payloads and compression statistics. +func (t Compressor) CompressPayloads(payloads [][]byte) (*CompressionResult, error) { + result := &CompressionResult{ + Payloads: payloads, + WasCompressed: false, + } + + if !t.Enabled || len(payloads) == 0 { + return result, nil + } + + // Calculate total size to determine if compression is worthwhile + totalSize := getPayloadSize(payloads) + result.OriginalSize = totalSize + + // Only compress if total size exceeds threshold + if totalSize < t.Threshold { + result.CompressedSize = totalSize + result.CompressionRatio = 1.0 + return result, nil + } + + compressed := make([][]byte, len(payloads)) + compressedSize := 0 + + for i, payload := range payloads { + var buf bytes.Buffer + + w := gzipWriterPool.Get().(*gzip.Writer) + w.Reset(&buf) + + if _, err := w.Write(payload); err != nil { + w.Close() + return nil, fmt.Errorf("failed to write to gzip writer: %w", err) + } + + if err := w.Close(); err != nil { + return nil, fmt.Errorf("failed to close gzip writer: %w", err) + } + + gzipWriterPool.Put(w) + + compressed[i] = buf.Bytes() + compressedSize += len(compressed[i]) + } + + result.Payloads = compressed + result.WasCompressed = true + result.CompressedSize = compressedSize + + // Calculate compression ratio (compressed / original) + if totalSize > 0 { + result.CompressionRatio = float64(compressedSize) / float64(totalSize) + } + + return result, nil +} + // DecompressPayloads decompresses gzip-compressed message payloads. It lives // in the msgqueue package (rather than a single implementation) because a // compressed message can cross backends: a producer whose durable queue is diff --git a/internal/msgqueue/rabbitmq/gzip_test.go b/internal/msgqueue/gzip_test.go similarity index 73% rename from internal/msgqueue/rabbitmq/gzip_test.go rename to internal/msgqueue/gzip_test.go index 51b213ead5..85117f2b9c 100644 --- a/internal/msgqueue/rabbitmq/gzip_test.go +++ b/internal/msgqueue/gzip_test.go @@ -1,4 +1,4 @@ -package rabbitmq +package msgqueue import ( "bytes" @@ -6,8 +6,6 @@ import ( "testing" "github.com/stretchr/testify/assert" - - "github.com/hatchet-dev/hatchet/internal/msgqueue" ) func generatePayloads(count, size int) [][]byte { @@ -27,23 +25,23 @@ func generatePayloads(count, size int) [][]byte { return payloads } -func newMQ() compressor { - return compressor{ - compressionEnabled: true, - compressionThreshold: 0, +func newCompressor() Compressor { + return Compressor{ + Enabled: true, + Threshold: 0, } } func TestCompressDecompressRoundtrip(t *testing.T) { - mq := newMQ() + mq := newCompressor() payloads := generatePayloads(5, 10*1024) - result, err := mq.compressPayloads(payloads) + result, err := mq.CompressPayloads(payloads) assert.NoError(t, err) assert.True(t, result.WasCompressed, "expected WasCompressed to be true") assert.Equal(t, len(payloads), len(result.Payloads), "expected %d payloads, got %d", len(payloads), len(result.Payloads)) - decompressed, err := msgqueue.DecompressPayloads(result.Payloads) + decompressed, err := DecompressPayloads(result.Payloads) assert.NoError(t, err) @@ -54,69 +52,69 @@ func TestCompressDecompressRoundtrip(t *testing.T) { } func TestCompressPayloadsDisabled(t *testing.T) { - mq := compressor{ - compressionEnabled: false, - compressionThreshold: 0, + mq := Compressor{ + Enabled: false, + Threshold: 0, } payloads := generatePayloads(3, 1024) - result, err := mq.compressPayloads(payloads) + result, err := mq.CompressPayloads(payloads) assert.NoError(t, err) assert.False(t, result.WasCompressed, "expected WasCompressed to be false when compression is disabled") } func TestCompressPayloadsBelowThreshold(t *testing.T) { - mq := compressor{ - compressionEnabled: true, - compressionThreshold: 100 * 1024, + mq := Compressor{ + Enabled: true, + Threshold: 100 * 1024, } payloads := generatePayloads(1, 1024) - result, err := mq.compressPayloads(payloads) + result, err := mq.CompressPayloads(payloads) assert.NoError(t, err) assert.False(t, result.WasCompressed, "expected WasCompressed to be false when below threshold") } func BenchmarkCompressPayloads_1x10KiB(b *testing.B) { - mq := newMQ() + mq := newCompressor() payloads := generatePayloads(1, 10*1024) b.ResetTimer() b.ReportAllocs() for i := 0; i < b.N; i++ { - _, _ = mq.compressPayloads(payloads) + _, _ = mq.CompressPayloads(payloads) } } func BenchmarkCompressPayloads_10x10KiB(b *testing.B) { - mq := newMQ() + mq := newCompressor() payloads := generatePayloads(10, 10*1024) b.ResetTimer() b.ReportAllocs() for i := 0; i < b.N; i++ { - _, _ = mq.compressPayloads(payloads) + _, _ = mq.CompressPayloads(payloads) } } func BenchmarkCompressPayloads_10x100KiB(b *testing.B) { - mq := newMQ() + mq := newCompressor() payloads := generatePayloads(10, 100*1024) b.ResetTimer() b.ReportAllocs() for i := 0; i < b.N; i++ { - _, _ = mq.compressPayloads(payloads) + _, _ = mq.CompressPayloads(payloads) } } func BenchmarkCompressPayloads_Concurrent(b *testing.B) { - mq := newMQ() + mq := newCompressor() payloads := generatePayloads(5, 10*1024) b.ResetTimer() b.ReportAllocs() b.RunParallel(func(pb *testing.PB) { for pb.Next() { - _, _ = mq.compressPayloads(payloads) + _, _ = mq.CompressPayloads(payloads) } }) } diff --git a/internal/msgqueue/msg.go b/internal/msgqueue/msg.go index 23c802a3e1..84b8313db7 100644 --- a/internal/msgqueue/msg.go +++ b/internal/msgqueue/msg.go @@ -68,6 +68,13 @@ func DecodeAndValidateSingleton(dv datautils.DataDecoderValidator, payloads [][] return dv.DecodeAndValidate(payloads[0], target) } +// Clone returns a shallow copy of the message. Payloads and OtelCarrier are +// shared with the original: callers must replace them, never mutate in place. +func (t *Message) Clone() *Message { + c := *t + return &c +} + func (t *Message) Serialize() ([]byte, error) { return json.Marshal(t) } diff --git a/internal/msgqueue/postgres/pubsub.go b/internal/msgqueue/postgres/pubsub.go index d4c171a3f1..9dd36e3e36 100644 --- a/internal/msgqueue/postgres/pubsub.go +++ b/internal/msgqueue/postgres/pubsub.go @@ -5,6 +5,7 @@ import ( "encoding/json" "time" + "github.com/jackc/pgx/v5/pgxpool" "github.com/rs/zerolog" "golang.org/x/sync/errgroup" @@ -18,19 +19,16 @@ import ( // fallbackPollBatchSize is the max number of >8KB fallback rows drained per poll. const fallbackPollBatchSize = 100 -// PubSub implements msgqueue.PubSub over Postgres LISTEN/NOTIFY. Topic names -// match the legacy non-durable queue / NOTIFY names, and all traffic -// multiplexes over the single "hatchet_listener" channel, so mixed-version -// fleets interoperate. -// -// INVARIANT: the repo passed to NewPubSub must be built on a dedicated pool -// from the direct (non-pgbouncer) database URL — never the shared repository -// pool. Pub can be called from within durable-write paths, and LISTEN does not -// survive transaction pooling. +// PubSub implements msgqueue.PubSub in Postgres, using LISTEN/NOTIFY for +// messages less than 8KB and the "MessageQueueItem" table for messages +// greater than 8KB. type PubSub struct { repo v1.MessageQueueRepository l *zerolog.Logger + // pool is the dedicated pool backing repo; used only for readiness checks + pool *pgxpool.Pool + // ttlCache dedupes queue-row upserts, which are needed so >8KB payloads can // fall back to durable rows ttlCache *cache.TTLCache[string, bool] @@ -39,7 +37,8 @@ type PubSub struct { type PubSubOpt func(*PubSubOpts) type PubSubOpts struct { - l *zerolog.Logger + l *zerolog.Logger + pool *pgxpool.Pool } func defaultPubSubOpts() *PubSubOpts { @@ -56,8 +55,21 @@ func WithPubSubLogger(l *zerolog.Logger) PubSubOpt { } } +// WithPubSubPool provides the pool backing the repository so IsReady can ping +// the database; without it, IsReady always reports true. +func WithPubSubPool(pool *pgxpool.Pool) PubSubOpt { + return func(opts *PubSubOpts) { + opts.pool = pool + } +} + // NewPubSub creates a new Postgres-backed PubSub over the given message queue // repository. +// +// The repo must be built on a dedicated pool from the direct (non-pgbouncer) +// database URL — never the shared repository pool or the pgbouncer URL. Pub +// can be called from within durable-write paths (so sharing pools is a +// deadlock risk), and LISTEN does not survive transaction pooling. func NewPubSub(repo v1.MessageQueueRepository, fs ...PubSubOpt) (func() error, *PubSub, error) { opts := defaultPubSubOpts() @@ -70,6 +82,7 @@ func NewPubSub(repo v1.MessageQueueRepository, fs ...PubSubOpt) (func() error, * p := &PubSub{ repo: repo, l: opts.l, + pool: opts.pool, ttlCache: c, } @@ -79,8 +92,18 @@ func NewPubSub(repo v1.MessageQueueRepository, fs ...PubSubOpt) (func() error, * }, p, nil } +// IsReady reports whether the database backing the pub/sub is reachable. Note +// that this is informational: it is deliberately not wired into the health +// probes, since a degraded pub/sub shouldn't fail liveness or readiness. func (p *PubSub) IsReady() bool { - return true + if p.pool == nil { + return true + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + return p.pool.Ping(ctx) == nil } // Pub publishes a message to the topic via NOTIFY. Payloads whose wrapped diff --git a/internal/msgqueue/rabbitmq/gzip.go b/internal/msgqueue/rabbitmq/gzip.go deleted file mode 100644 index 5faeaa09f9..0000000000 --- a/internal/msgqueue/rabbitmq/gzip.go +++ /dev/null @@ -1,105 +0,0 @@ -package rabbitmq - -import ( - "bytes" - "compress/gzip" - "fmt" - - "sync" -) - -// compressor holds gzip compression settings shared by the durable -// MessageQueueImpl and the best-effort PubSub (compression must agree on both -// sides for wire compatibility). -type compressor struct { - compressionEnabled bool - compressionThreshold int -} - -type CompressionResult struct { - Payloads [][]byte - WasCompressed bool - OriginalSize int - CompressedSize int - - // CompressionRatio is the ratio of compressed size to original size (compressed / original) - CompressionRatio float64 -} - -// gzipWriterPool reuses gzip.Writer instances to avoid repeated allocations. -// No explicit size cap is needed: sync.Pool is self-limiting because the Go -// runtime evicts pooled objects during GC, so the pool cannot grow unbounded. -// In practice the pool size is also bounded by the number of goroutines -// concurrently compressing, which is small for a RabbitMQ publish path. -var gzipWriterPool = sync.Pool{ - New: func() any { - return gzip.NewWriter(nil) - }, -} - -func getPayloadSize(payloads [][]byte) int { - totalSize := 0 - for _, payload := range payloads { - totalSize += len(payload) - } - return totalSize -} - -// compressPayloads compresses message payloads using gzip if they exceed the minimum size threshold. -// Returns compression results including the compressed payloads and compression statistics. -func (t compressor) compressPayloads(payloads [][]byte) (*CompressionResult, error) { - result := &CompressionResult{ - Payloads: payloads, - WasCompressed: false, - } - - if !t.compressionEnabled || len(payloads) == 0 { - return result, nil - } - - // Calculate total size to determine if compression is worthwhile - totalSize := getPayloadSize(payloads) - result.OriginalSize = totalSize - - // Only compress if total size exceeds threshold - if totalSize < t.compressionThreshold { - result.CompressedSize = totalSize - result.CompressionRatio = 1.0 - return result, nil - } - - compressed := make([][]byte, len(payloads)) - compressedSize := 0 - - for i, payload := range payloads { - var buf bytes.Buffer - - w := gzipWriterPool.Get().(*gzip.Writer) - w.Reset(&buf) - - if _, err := w.Write(payload); err != nil { - w.Close() - return nil, fmt.Errorf("failed to write to gzip writer: %w", err) - } - - if err := w.Close(); err != nil { - return nil, fmt.Errorf("failed to close gzip writer: %w", err) - } - - gzipWriterPool.Put(w) - - compressed[i] = buf.Bytes() - compressedSize += len(compressed[i]) - } - - result.Payloads = compressed - result.WasCompressed = true - result.CompressedSize = compressedSize - - // Calculate compression ratio (compressed / original) - if totalSize > 0 { - result.CompressionRatio = float64(compressedSize) / float64(totalSize) - } - - return result, nil -} diff --git a/internal/msgqueue/rabbitmq/pubsub.go b/internal/msgqueue/rabbitmq/pubsub.go index cb22473c05..d499064a5b 100644 --- a/internal/msgqueue/rabbitmq/pubsub.go +++ b/internal/msgqueue/rabbitmq/pubsub.go @@ -37,9 +37,7 @@ type PubSub struct { // lru cache of tenant exchanges we've already declared exchangeCache *lru.Cache[string, bool] - compressor - - maxPayloadSize int + compressor msgqueue.Compressor } type PubSubOpt func(*PubSubOpts) @@ -99,7 +97,7 @@ func WithPubSubGzip(enabled bool, threshold int) PubSubOpt { opts.compressionEnabled = enabled if threshold <= 0 { - threshold = 5 * 1024 // default to 5KB + threshold = msgqueue.DefaultCompressionThreshold } opts.compressionThreshold = threshold @@ -141,11 +139,10 @@ func NewPubSub(fs ...PubSubOpt) (func() error, *PubSub, error) { l: opts.l, pubChannels: pubChannelPool, subChannels: subChannelPool, - compressor: compressor{ - compressionEnabled: opts.compressionEnabled, - compressionThreshold: opts.compressionThreshold, + compressor: msgqueue.Compressor{ + Enabled: opts.compressionEnabled, + Threshold: opts.compressionThreshold, }, - maxPayloadSize: 16 * 1024 * 1024, // 16 MB } p.exchangeCache, _ = lru.New[string, bool](2000) //nolint:errcheck // this only returns an error if the size is less than 0 @@ -171,7 +168,7 @@ func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Me // on a shallow copy: callers dual-publish the same *Message they hand to // the durable queue, which may not have serialized it yet. if len(msg.Payloads) > 0 && !msg.Compressed { - compressionResult, err := p.compressPayloads(msg.Payloads) + compressionResult, err := p.compressor.CompressPayloads(msg.Payloads) if err != nil { p.l.Error().Msgf("error compressing payloads: %v", err) @@ -179,10 +176,10 @@ func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Me } if compressionResult.WasCompressed { - msgCp := *msg + msgCp := msg.Clone() msgCp.Payloads = compressionResult.Payloads msgCp.Compressed = true - msg = &msgCp + msg = msgCp } } @@ -223,9 +220,9 @@ func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Me return err } - if len(body) > p.maxPayloadSize { + if len(body) > maxPayloadSize { if len(msg.Payloads) == 1 { - return fmt.Errorf("message size %d bytes exceeds maximum allowed size of %d bytes", len(body), p.maxPayloadSize) + return fmt.Errorf("message size %d bytes exceeds maximum allowed size of %d bytes", len(body), maxPayloadSize) } // split the payloads in half and publish recursively until each chunk is @@ -233,16 +230,10 @@ func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Me payloadsPerChunk := max(len(msg.Payloads)/2, 1) for chunk := range slices.Chunk(msg.Payloads, payloadsPerChunk) { - err := p.Pub(ctx, topic, &msgqueue.Message{ - ID: msg.ID, - Payloads: chunk, - TenantID: msg.TenantID, - ImmediatelyExpire: msg.ImmediatelyExpire, - Persistent: msg.Persistent, - OtelCarrier: msg.OtelCarrier, - Retries: msg.Retries, // nolint: staticcheck - Compressed: msg.Compressed, - }) + msgCp := msg.Clone() + msgCp.Payloads = chunk + + err := p.Pub(ctx, topic, msgCp) if err != nil { return err @@ -421,7 +412,7 @@ func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.MsgHandler) (func() // the default exchange. func (p *PubSub) declareSubQueue(ch *amqp.Channel, topic msgqueue.Topic) (string, error) { args := make(amqp.Table) - args["x-consumer-timeout"] = 300000 // 5 minutes + args["x-consumer-timeout"] = consumerTimeout.Milliseconds() name := topic.Name() diff --git a/internal/msgqueue/rabbitmq/rabbitmq.go b/internal/msgqueue/rabbitmq/rabbitmq.go index 36def2c374..2d822388ab 100644 --- a/internal/msgqueue/rabbitmq/rabbitmq.go +++ b/internal/msgqueue/rabbitmq/rabbitmq.go @@ -29,6 +29,14 @@ const MAX_RETRY_COUNT = 15 const RETRY_INTERVAL = 2 * time.Second const RETRY_RESET_INTERVAL = 30 * time.Second +// maxPayloadSize is the maximum size of a single published AMQP message body; +// larger messages are recursively split into payload chunks. +const maxPayloadSize = 16 * 1024 * 1024 // 16 MB + +// consumerTimeout is the x-consumer-timeout applied to declared queues: the +// broker requeues a delivery if it isn't acked within this window. +const consumerTimeout = 5 * time.Minute + // MessageQueueImpl implements MessageQueue interface using AMQP. type MessageQueueImpl struct { ctx context.Context @@ -44,9 +52,8 @@ type MessageQueueImpl struct { deadLetterBackoff time.Duration - compressor + compressor msgqueue.Compressor - maxPayloadSize int enableMessageRejection bool maxDeathCount int } @@ -122,7 +129,7 @@ func WithGzipCompression(enabled bool, threshold int) MessageQueueImplOpt { opts.compressionEnabled = enabled if threshold <= 0 { - threshold = 5 * 1024 // default to 5KB + threshold = msgqueue.DefaultCompressionThreshold } opts.compressionThreshold = threshold @@ -190,13 +197,12 @@ func New(fs ...MessageQueueImplOpt) (func() error, *MessageQueueImpl, error) { pubChannels: pubChannelPool, subChannels: subChannelPool, deadLetterBackoff: opts.deadLetterBackoff, - compressor: compressor{ - compressionEnabled: opts.compressionEnabled, - compressionThreshold: opts.compressionThreshold, + compressor: msgqueue.Compressor{ + Enabled: opts.compressionEnabled, + Threshold: opts.compressionThreshold, }, enableMessageRejection: opts.enableMessageRejection, maxDeathCount: opts.maxDeathCount, - maxPayloadSize: 16 * 1024 * 1024, // 16 MB } // init the queues in a blocking fashion @@ -280,7 +286,7 @@ func (t *MessageQueueImpl) pubMessage(ctx context.Context, q msgqueue.Queue, msg msg.SetOtelCarrier(otelCarrier) - var compressionResult *CompressionResult + var compressionResult *msgqueue.CompressionResult // don't re-compress if the message was already compressed. We work on a // shallow copy rather than mutating in place: callers may hand the same @@ -288,17 +294,17 @@ func (t *MessageQueueImpl) pubMessage(ctx context.Context, q msgqueue.Queue, msg // and compression is a rabbitmq wire concern if len(msg.Payloads) > 0 && !msg.Compressed { var err error - compressionResult, err = t.compressPayloads(msg.Payloads) + compressionResult, err = t.compressor.CompressPayloads(msg.Payloads) if err != nil { t.l.Error().Msgf("error compressing payloads: %v", err) return fmt.Errorf("failed to compress payloads: %w", err) } if compressionResult.WasCompressed { - msgCp := *msg + msgCp := msg.Clone() msgCp.Payloads = compressionResult.Payloads msgCp.Compressed = true - msg = &msgCp + msg = msgCp t.l.Debug().Msgf("compressed payloads for message %s: original=%d bytes, compressed=%d bytes, ratio=%.2f%%", msg.ID, compressionResult.OriginalSize, compressionResult.CompressedSize, compressionResult.CompressionRatio*100) @@ -353,9 +359,9 @@ func (t *MessageQueueImpl) pubMessage(ctx context.Context, q msgqueue.Queue, msg bodySize := len(body) - if bodySize > t.maxPayloadSize { + if bodySize > maxPayloadSize { if len(msg.Payloads) == 1 { - err := fmt.Errorf("message size %d bytes exceeds maximum allowed size of %d bytes", bodySize, t.maxPayloadSize) + err := fmt.Errorf("message size %d bytes exceeds maximum allowed size of %d bytes", bodySize, maxPayloadSize) span.RecordError(err) span.SetStatus(codes.Error, "message size exceeds maximum allowed size") return err @@ -378,16 +384,10 @@ func (t *MessageQueueImpl) pubMessage(ctx context.Context, q msgqueue.Queue, msg // recursively call pubMessage with the chunked payloads // if the payload chunks are still too large, this will continue to split them // until they are under the max size. - err := t.pubMessage(ctx, q, &msgqueue.Message{ - ID: msg.ID, - Payloads: chunk, - TenantID: msg.TenantID, - ImmediatelyExpire: msg.ImmediatelyExpire, - Persistent: msg.Persistent, - OtelCarrier: msg.OtelCarrier, - Retries: msg.Retries, - Compressed: msg.Compressed, - }) + msgCp := msg.Clone() + msgCp.Payloads = chunk + + err := t.pubMessage(ctx, q, msgCp) if err != nil { return err @@ -521,7 +521,7 @@ func (t *MessageQueueImpl) initQueue(ch *amqp.Channel, q msgqueue.Queue) (string } args := make(amqp.Table) - args["x-consumer-timeout"] = 300000 // 5 minutes + args["x-consumer-timeout"] = consumerTimeout.Milliseconds() name := q.Name() if !q.IsDLQ() && q.DLQ() != nil && q.DLQ().IsAutoDLQ() { diff --git a/internal/services/controllers/olap/process_dag_status_updates.go b/internal/services/controllers/olap/process_dag_status_updates.go index e69fcfac55..e80b74d847 100644 --- a/internal/services/controllers/olap/process_dag_status_updates.go +++ b/internal/services/controllers/olap/process_dag_status_updates.go @@ -103,8 +103,6 @@ func (o *OLAPControllerImpl) notifyDAGsUpdated(ctx context.Context, rows []v1.Up return err } - // fanout-only: the dispatcher's workflow run subscriptions consume - // workflow-run-finished off the tenant stream if err := msgqueue.PubTenantMessage(ctx, o.l, nil, o.pubsub, nil, msg); err != nil { o.l.Warn().Ctx(ctx).Err(err).Msg("could not publish workflow-run-finished to tenant stream") } diff --git a/internal/services/controllers/olap/process_task_status_updates.go b/internal/services/controllers/olap/process_task_status_updates.go index 6fe082b32c..e90d5465a9 100644 --- a/internal/services/controllers/olap/process_task_status_updates.go +++ b/internal/services/controllers/olap/process_task_status_updates.go @@ -103,8 +103,6 @@ func (o *OLAPControllerImpl) notifyTasksUpdated(ctx context.Context, rows []v1.U return err } - // fanout-only: the dispatcher's workflow run subscriptions consume - // workflow-run-finished off the tenant stream if err := msgqueue.PubTenantMessage(ctx, o.l, nil, o.pubsub, nil, msg); err != nil { o.l.Warn().Ctx(ctx).Err(err).Msg("could not publish workflow-run-finished to tenant stream") } diff --git a/internal/services/health/health.go b/internal/services/health/health.go index 35a6b41823..eae6320598 100644 --- a/internal/services/health/health.go +++ b/internal/services/health/health.go @@ -24,10 +24,6 @@ type Health struct { l *zerolog.Logger } -// New creates a health service over the durable message queue and repository. -// The best-effort PubSub is deliberately NOT part of the probes: 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. func New(repo v1.HealthRepository, queue msgqueue.MessageQueue, version string, l *zerolog.Logger) *Health { return &Health{ version: version, diff --git a/pkg/config/loader/loader.go b/pkg/config/loader/loader.go index e8b92fc234..bf0f49dc8a 100644 --- a/pkg/config/loader/loader.go +++ b/pkg/config/loader/loader.go @@ -982,6 +982,7 @@ func createPubSubV1(dc *database.Layer, cf *server.ServerConfigFile, l *zerolog. cleanupPg, pgps, err := pgmq.NewPubSub( pubsubRepo, pgmq.WithPubSubLogger(l), + pgmq.WithPubSubPool(pubsubPool), ) if err != nil { From ebafd2fc211f15f994856e87550e7d5e7e7892dd Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Mon, 27 Jul 2026 19:56:46 -0400 Subject: [PATCH 5/6] fix TestMsgIdBufferMemoryLeak failing after new tests joined the package MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- internal/msgqueue/mq_sub_buffer_test.go | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/internal/msgqueue/mq_sub_buffer_test.go b/internal/msgqueue/mq_sub_buffer_test.go index f3ef8c517b..776ce2bdc2 100644 --- a/internal/msgqueue/mq_sub_buffer_test.go +++ b/internal/msgqueue/mq_sub_buffer_test.go @@ -116,7 +116,12 @@ func TestMsgIdBufferMemoryLeak(t *testing.T) { // Create a buffer buf := newMsgIDBuffer(ctx, testTenantID, "test-msg", dst, 10*time.Millisecond, testBufSize, 10, false) - // Force GC and get baseline + // Force GC twice and get baseline: earlier tests in this package leave heap + // state behind (notably sync.Pool contents, which survive one GC in the + // victim cache and are only freed by a second). A single GC would count + // that memory in the baseline and release it before the after-measurement, + // systematically biasing the growth negative. + runtime.GC() runtime.GC() time.Sleep(50 * time.Millisecond) var baselineMemStats runtime.MemStats @@ -172,7 +177,9 @@ func TestMsgIdBufferMemoryLeak(t *testing.T) { // Verify memory didn't grow excessively // With 1000 flushes, if we were creating goroutines+timers for each, // we'd see significant memory growth (multiple MB) - memGrowthMB := float64(afterMemStats.Alloc-baselineMemStats.Alloc) / 1024 / 1024 + // subtract as floats: Alloc is a uint64, and the heap can legitimately be + // smaller than the baseline after the final GC, which would underflow + memGrowthMB := (float64(afterMemStats.Alloc) - float64(baselineMemStats.Alloc)) / 1024 / 1024 if memGrowthMB > 5 { t.Errorf("Excessive memory growth: %.2f MB (expected <5MB)", memGrowthMB) } From 8fdbad1d25d0da20ecd2d9b62c3f338a690e5231 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20Kupczy=C5=84ski?= Date: Tue, 28 Jul 2026 11:20:42 +0200 Subject: [PATCH 6/6] address review: reset Compressed after rabbitmq pubsub decompress, document 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. --- internal/msgqueue/postgres/pubsub.go | 5 +++++ internal/msgqueue/pubsub.go | 5 +++++ internal/msgqueue/rabbitmq/pubsub.go | 1 + 3 files changed, 11 insertions(+) diff --git a/internal/msgqueue/postgres/pubsub.go b/internal/msgqueue/postgres/pubsub.go index 9dd36e3e36..c7e16de689 100644 --- a/internal/msgqueue/postgres/pubsub.go +++ b/internal/msgqueue/postgres/pubsub.go @@ -144,6 +144,11 @@ func (p *PubSub) Pub(ctx context.Context, topic msgqueue.Topic, msg *msgqueue.Me // Sub subscribes to a topic. Inline NOTIFY payloads are handled directly; // non-JSON notifications and a 1s ticker wake a poll that drains >8KB fallback // rows. Delivery is at-most-once: handler errors are logged, never redelivered. +// +// Fanout divergence: inline NOTIFY payloads reach every subscriber of the +// topic, but each >8KB fallback row is read and acked by exactly one +// subscriber, unlike the rabbitmq fanout, which delivers every message to +// every subscriber. func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.MsgHandler) (func() error, error) { err := p.ensureQueue(context.Background(), topic) diff --git a/internal/msgqueue/pubsub.go b/internal/msgqueue/pubsub.go index c8d3cd1c45..d573d319ff 100644 --- a/internal/msgqueue/pubsub.go +++ b/internal/msgqueue/pubsub.go @@ -68,6 +68,11 @@ type PubSub interface { // semantics: handler errors are logged, never redelivered. It returns a // cleanup function that should be called when the subscription is no // longer needed. + // + // Fanout to multiple subscribers of the same topic is backend-dependent: + // rabbitmq delivers every message to every subscriber, while the postgres + // implementation delivers messages over 8KB to only one subscriber (see + // postgres.PubSub.Sub). Sub(topic Topic, handler MsgHandler) (func() error, error) // IsReady returns true if the pub/sub mechanism is ready to accept messages. diff --git a/internal/msgqueue/rabbitmq/pubsub.go b/internal/msgqueue/rabbitmq/pubsub.go index d499064a5b..b1b50512f1 100644 --- a/internal/msgqueue/rabbitmq/pubsub.go +++ b/internal/msgqueue/rabbitmq/pubsub.go @@ -352,6 +352,7 @@ func (p *PubSub) Sub(topic msgqueue.Topic, handler msgqueue.MsgHandler) (func() } msg.Payloads = decompressedPayloads + msg.Compressed = false } p.l.Debug().Msgf("(session: %d) got pubsub msg", session)