diff --git a/cmd/hatchet-engine/engine/run.go b/cmd/hatchet-engine/engine/run.go index 139e6feb7b..396bfdad36 100644 --- a/cmd/hatchet-engine/engine/run.go +++ b/cmd/hatchet-engine/engine/run.go @@ -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), @@ -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/rabbitmq/gzip.go b/internal/msgqueue/gzip.go similarity index 67% rename from internal/msgqueue/rabbitmq/gzip.go rename to internal/msgqueue/gzip.go index da7cda3ca1..d13bdcb2cd 100644 --- a/internal/msgqueue/rabbitmq/gzip.go +++ b/internal/msgqueue/gzip.go @@ -1,14 +1,25 @@ -package rabbitmq +package msgqueue import ( "bytes" "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 @@ -23,7 +34,7 @@ type CompressionResult struct { // 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. +// concurrently compressing, which is small for a publish path. var gzipWriterPool = sync.Pool{ New: func() any { return gzip.NewWriter(nil) @@ -38,15 +49,16 @@ func getPayloadSize(payloads [][]byte) int { 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 *MessageQueueImpl) compressPayloads(payloads [][]byte) (*CompressionResult, error) { +// 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 { + if !t.Enabled || len(payloads) == 0 { return result, nil } @@ -55,7 +67,7 @@ func (t *MessageQueueImpl) compressPayloads(payloads [][]byte) (*CompressionResu result.OriginalSize = totalSize // Only compress if total size exceeds threshold - if totalSize < t.compressionThreshold { + if totalSize < t.Threshold { result.CompressedSize = totalSize result.CompressionRatio = 1.0 return result, nil @@ -97,8 +109,12 @@ func (t *MessageQueueImpl) compressPayloads(payloads [][]byte) (*CompressionResu return result, nil } -// decompressPayloads decompresses message payloads using gzip. -func (t *MessageQueueImpl) decompressPayloads(payloads [][]byte) ([][]byte, error) { +// 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 } diff --git a/internal/msgqueue/rabbitmq/gzip_test.go b/internal/msgqueue/gzip_test.go similarity index 74% rename from internal/msgqueue/rabbitmq/gzip_test.go rename to internal/msgqueue/gzip_test.go index 60e1378a80..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" @@ -25,23 +25,23 @@ func generatePayloads(count, size int) [][]byte { return payloads } -func newMQ() *MessageQueueImpl { - return &MessageQueueImpl{ - 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 := mq.decompressPayloads(result.Payloads) + decompressed, err := DecompressPayloads(result.Payloads) assert.NoError(t, err) @@ -52,69 +52,69 @@ func TestCompressDecompressRoundtrip(t *testing.T) { } func TestCompressPayloadsDisabled(t *testing.T) { - mq := &MessageQueueImpl{ - 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 := &MessageQueueImpl{ - 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/mq_pub_buffer_test.go b/internal/msgqueue/mq_pub_buffer_test.go index 28d3ef74ae..de3f6ad8f9 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. @@ -24,11 +22,10 @@ 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) 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..e20bfe6f4f 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 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. 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 MsgHandler, postAck MsgHandler) (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/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) } 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/msgqueue.go b/internal/msgqueue/msgqueue.go index 570e59580f..861ac46a71 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,74 +151,10 @@ 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 +// 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. @@ -247,13 +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) - - // 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 + 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/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..5708379ce6 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,14 +164,10 @@ 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 } -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 { @@ -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..c7e16de689 --- /dev/null +++ b/internal/msgqueue/postgres/pubsub.go @@ -0,0 +1,299 @@ +package postgres + +import ( + "context" + "encoding/json" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "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 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] +} + +type PubSubOpt func(*PubSubOpts) + +type PubSubOpts struct { + l *zerolog.Logger + pool *pgxpool.Pool +} + +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 + } +} + +// 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() + + for _, f := range fs { + f(opts) + } + + c := cache.NewTTL[string, bool]() + + p := &PubSub{ + repo: repo, + l: opts.l, + pool: opts.pool, + ttlCache: c, + } + + return func() error { + c.Stop() + return nil + }, 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 { + 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 +// 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. +// +// 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) + + 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") + } + } + + // 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) + } + } + + // 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..6979a53fe9 --- /dev/null +++ b/internal/msgqueue/postgres/pubsub_test.go @@ -0,0 +1,335 @@ +//go:build integration + +package postgres + +import ( + "bytes" + "compress/gzip" + "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()) +} + +// 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) { + 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..d573d319ff --- /dev/null +++ b/internal/msgqueue/pubsub.go @@ -0,0 +1,162 @@ +package msgqueue + +import ( + "context" + "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 + kind TopicKind +} + +func (t Topic) Name() string { + return t.name +} + +func (t Topic) Kind() TopicKind { + return t.kind +} + +// 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", + kind: TopicKindTenantStream, + } +} + +// 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), + kind: TopicKindSchedulerPartition, + } +} + +// 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. + // + // 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. + 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. +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.Kind() == TopicKindTenantStream && msg.ID != MsgIDTaskStreamEvent { + return nil + } + + return g.inner.Pub(ctx, topic, msg) +} + +func (g *gatedPubSub) Sub(topic Topic, handler MsgHandler) (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..8f75490b51 --- /dev/null +++ b/internal/msgqueue/pubsub_test.go @@ -0,0 +1,157 @@ +package msgqueue + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/rs/zerolog" + "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.Equal(t, TopicKindTenantStream, tenantTopic.Kind()) + + schedulerTopic := SchedulerPartitionTopic("mypartition") + assert.Equal(t, "mypartition_scheduler_v1", schedulerTopic.Name()) + assert.Equal(t, TopicKindSchedulerPartition, schedulerTopic.Kind()) +} + +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 MsgHandler) (func() error, error) { + return func() error { return nil }, nil +} + +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() + + 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/pubsub.go b/internal/msgqueue/rabbitmq/pubsub.go new file mode 100644 index 0000000000..b1b50512f1 --- /dev/null +++ b/internal/msgqueue/rabbitmq/pubsub.go @@ -0,0 +1,465 @@ +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 msgqueue.Compressor +} + +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 = msgqueue.DefaultCompressionThreshold + } + + 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: msgqueue.Compressor{ + Enabled: opts.compressionEnabled, + Threshold: opts.compressionThreshold, + }, + } + + 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.compressor.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.Clone() + 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) > maxPayloadSize { + if len(msg.Payloads) == 1 { + 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 + // 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) { + msgCp := msg.Clone() + msgCp.Payloads = chunk + + err := p.Pub(ctx, topic, msgCp) + + if err != nil { + return err + } + } + + return nil + } + + exchange := "" + routingKey := topic.Name() + + 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 { + 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.MsgHandler) (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 := msgqueue.DecompressPayloads(msg.Payloads) + + if err != nil { + p.l.Error().Msgf("error decompressing payloads: %v", err) + return + } + + msg.Payloads = decompressedPayloads + msg.Compressed = false + } + + 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"] = consumerTimeout.Milliseconds() + + name := topic.Name() + + if topic.Kind() == msgqueue.TopicKindTenantStream { + 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.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 { + 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..2d822388ab 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" ) @@ -32,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 @@ -42,19 +47,13 @@ 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 - maxPayloadSize int - compressionEnabled bool - compressionThreshold int + compressor msgqueue.Compressor + enableMessageRejection bool maxDeathCount int } @@ -66,28 +65,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 +118,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 @@ -138,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 @@ -198,24 +189,21 @@ 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: msgqueue.Compressor{ + Enabled: opts.compressionEnabled, + Threshold: opts.compressionThreshold, + }, + enableMessageRejection: opts.enableMessageRejection, + maxDeathCount: opts.maxDeathCount, + } // init the queues in a blocking fashion poolCh, err := subChannelPool.Acquire(ctx) @@ -298,20 +286,25 @@ 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 + // 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) + 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 { - msg.Payloads = compressionResult.Payloads - msg.Compressed = true + msgCp := msg.Clone() + 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) @@ -366,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 @@ -391,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 @@ -469,28 +456,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 @@ -499,8 +464,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()) @@ -550,73 +515,15 @@ 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 } args := make(amqp.Table) - args["x-consumer-timeout"] = 300000 // 5 minutes + args["x-consumer-timeout"] = consumerTimeout.Milliseconds() 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 +569,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 } @@ -679,8 +576,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 @@ -787,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/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..83edf370d6 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" @@ -12,7 +11,7 @@ import ( ) type sharedTenantSub struct { - fs *syncx.Map[int, AckHook] + fs *syncx.Map[int, MsgHandler] counter int isRunning bool mu sync.Mutex @@ -21,19 +20,19 @@ 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, } } -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() @@ -48,18 +47,10 @@ 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 { + t.fs.Range(func(key int, f MsgHandler) bool { if err := f(task); err != nil { innerErr = multierror.Append(innerErr, err) } @@ -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 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 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 2e818dfd89..66eecfe0e6 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..8772c803fb 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, @@ -1158,7 +1171,7 @@ 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 } } diff --git a/internal/services/controllers/olap/process_dag_status_updates.go b/internal/services/controllers/olap/process_dag_status_updates.go index e364eb3169..e80b74d847 100644 --- a/internal/services/controllers/olap/process_dag_status_updates.go +++ b/internal/services/controllers/olap/process_dag_status_updates.go @@ -103,12 +103,8 @@ 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 + 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 82d511cfdf..e90d5465a9 100644 --- a/internal/services/controllers/olap/process_task_status_updates.go +++ b/internal/services/controllers/olap/process_task_status_updates.go @@ -103,12 +103,8 @@ 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 + 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 8ee6166af3..00ca076e13 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 } + + // 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") + } } 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..74fcbcef54 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, @@ -810,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, ) @@ -991,14 +1007,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 +1040,12 @@ func (tc *TasksControllerImpl) notifyQueuesOnCompletion(ctx context.Context, ten return } - err = tc.mq.SendMessage( - ctx, - msgqueue.TenantEventConsumerQueue(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 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..fb402a9d25 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{ @@ -944,7 +957,7 @@ 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) 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..386a327c52 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") } } }() @@ -1347,7 +1347,7 @@ 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 @@ -1407,7 +1407,7 @@ 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 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..977c5b20d5 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: 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 545e236118..8ddd7ef6de 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, ) @@ -605,11 +617,7 @@ 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)) @@ -681,11 +689,7 @@ 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") @@ -759,11 +763,7 @@ 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") @@ -918,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 diff --git a/internal/services/ticker/schedule_workflow_v1_test.go b/internal/services/ticker/schedule_workflow_v1_test.go index bbab76a341..422fa94c1b 100644 --- a/internal/services/ticker/schedule_workflow_v1_test.go +++ b/internal/services/ticker/schedule_workflow_v1_test.go @@ -72,11 +72,10 @@ 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) 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 7d57733ce2..bf0f49dc8a 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,15 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers } } + cleanupPubSub, pubsubv1, err = createPubSubV1(dc, cf, &l) + + if err != nil { + return nil, nil, err + } + ing, err = ingestor.NewIngestor( ingestor.WithMessageQueueV1(mqv1), + ingestor.WithPubSub(pubsubv1), ingestor.WithRepositoryV1(dc.V1), ) @@ -854,6 +865,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") } @@ -904,6 +919,7 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers Encryption: encryptionSvc, Layer: dc, MessageQueueV1: mqv1, + PubSubV1: pubsubv1, Services: services, PausedControllers: pausedControllers, InternalClientFactory: internalClientFactory, @@ -932,6 +948,113 @@ 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), + pgmq.WithPubSubPool(pubsubPool), + ) + + 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 +// 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) }