From 31f00ad1ae4f2fecbc7baa994ac0a7649b7498d3 Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Fri, 17 Jul 2026 09:47:44 -0400 Subject: [PATCH 1/3] feat: new Subscribe method with pub/sub for hot-path polling --- README.md | 94 ++++++++++-- examples/go.mod | 5 + examples/go.sum | 4 + examples/main.go | 33 +++- outbox.go | 68 ++++++++- outbox_bench_test.go | 128 +++++++++++----- pgpubsub.go | 314 ++++++++++++++++++++++++++++++++++++++ pubsub.go | 49 ++++++ subscribe.go | 177 +++++++++++++++++++++ subscribe_e2e_test.go | 346 ++++++++++++++++++++++++++++++++++++++++++ 10 files changed, 1163 insertions(+), 55 deletions(-) create mode 100644 pgpubsub.go create mode 100644 pubsub.go create mode 100644 subscribe.go create mode 100644 subscribe_e2e_test.go diff --git a/README.md b/README.md index 7588e20..65c59b1 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ Here's an example of flushing messages on `topic1` by simply printing them to th ```go type printFlusher struct{} -func (printFlusher) Flush(_ context.Context, msgs []*sqlc.Message) error { +func (printFlusher) Flush(_ pgoutbox.FlushContext, msgs []*sqlc.Message) error { for _, m := range msgs { fmt.Printf(" flushed id=%d topic=%s payload=%s\n", m.ID, m.Topic, string(m.Payload)) } @@ -97,16 +97,13 @@ outbox.ProcessMessages(ctx, "shipments") ## Atomic flush and delete -If your flusher writes to Postgres itself (e.g. into a relay table), implement `TxFlusher` instead of `Flusher`. `ProcessMessages` will pass the same transaction it uses to lock and delete messages, so your writes and the outbox delete commit or roll back together: +If your flusher writes to Postgres itself (e.g. into a relay table), use the transaction exposed by the `FlushContext` passed to `Flush`. It is the same transaction `ProcessMessages` uses to lock and delete messages, so your writes and the outbox delete commit or roll back together: ```go -type relayFlusher struct{ db *pgxpool.Pool } +type relayFlusher struct{} -func (f *relayFlusher) Flush(_ context.Context, _ []*sqlc.Message) error { - panic("not used; FlushWithTx is called instead") -} - -func (f *relayFlusher) FlushWithTx(ctx context.Context, tx pgx.Tx, msgs []*sqlc.Message) error { +func (f *relayFlusher) Flush(ctx pgoutbox.FlushContext, msgs []*sqlc.Message) error { + tx := ctx.Tx() for _, m := range msgs { if _, err := tx.Exec(ctx, "INSERT INTO relay (payload) VALUES ($1)", m.Payload); err != nil { return err @@ -116,9 +113,51 @@ func (f *relayFlusher) FlushWithTx(ctx context.Context, tx pgx.Tx, msgs []*sqlc. } ``` +`FlushContext` embeds `context.Context`, so flushers that don't need the transaction can ignore `Tx()` and treat it as a plain context. + +## Continuous processing with `Subscribe` + +Instead of calling `ProcessMessages` yourself, `Subscribe` runs it in a loop: it drains the topic, then waits until either the poll interval elapses or a new-message notification arrives (see below), and drains again. It blocks until its context is cancelled: + +```go +go func() { + err := outbox.Subscribe(ctx, "orders", + pgoutbox.WithPollInterval(5*time.Second), // default: 5s + pgoutbox.WithProcessOpts(pgoutbox.WithBatchSize(500)), // forwarded to every ProcessMessages call + ) + if err != nil && !errors.Is(err, context.Canceled) { + panic(err) + } +}() +``` + +Processing errors don't kill the loop — they're logged to the `WithLogger` logger and retried on the next wake-up. `Subscribe` returns an error immediately only if no flusher is registered for the topic or the subscription itself can't be established. + +### Waking on new messages with `LISTEN`/`NOTIFY` + +Without further configuration, `Subscribe` is purely poll-based. To wake subscribers the moment new messages commit, attach a `PubSub` — the built-in implementation rides Postgres `LISTEN`/`NOTIFY`: + +```go +ps, err := pgoutbox.NewPGPubSub(ctx, pool) +if err != nil { + panic(err) +} + +outbox, err := pgoutbox.NewOutbox(ctx, pool, pgoutbox.WithPubSub(ps)) +``` + +With a `PubSub` attached, `AddMessages` publishes a notification for each staged topic and `Subscribe` wakes on it instead of waiting out the poll interval. The pg-backed `PubSub` publishes *inside the `AddMessages` transaction*, so the notification is delivered exactly when the insert commits — and never for a transaction that rolls back. Postgres deduplicates identical notifications within a transaction, so any number of `AddMessages` calls for a topic in one transaction cost a single wake-up. + +Delivery is best-effort by design: if a notification is lost (for example while the listener reconnects), polling picks the messages up within one poll interval. The listener occupies a single dedicated connection (hijacked out of the pool so it doesn't consume a pool slot) no matter how many topics are subscribed. + +Two details worth knowing: + +- All notifications travel over one NOTIFY channel, `pgoutbox_pubsub` by default. Two outboxes sharing a database (e.g. different schemas) should use distinct channels via `pgoutbox.WithNotifyChannel("my_channel")` to avoid waking each other's subscribers. +- `PubSub` is an interface, so you can bring your own transport (e.g. Redis, NATS) instead of `LISTEN`/`NOTIFY`. If your implementation also implements `TxPublisher`, notifications are published transactionally as described above; otherwise `AddMessages` publishes best-effort at insert time. + ## Message expiration -Call `Start` to run background maintenance goroutines that delete old messages. Expiration is configured per topic, with an optional default for topics not explicitly named: +`NewOutbox` starts background maintenance goroutines that delete old messages; they run until the context passed to `NewOutbox` is cancelled. Expiration is configured per topic, with an optional default for topics not explicitly named: ```go ctx, cancel := context.WithCancel(context.Background()) @@ -134,14 +173,14 @@ if err != nil { } ``` -`Start` always runs a background scanner goroutine that polls the `topics` table, but maintenance loops are only launched for topics that actually have an expiration configured. Topics don't need to be declared at startup: the library tracks every topic that receives a message in a `topics` table (via a Postgres trigger) and applies the default expiration automatically. +The outbox always runs a background scanner goroutine that polls the `topics` table, but maintenance loops are only launched for topics that actually have an expiration configured. Topics don't need to be declared at startup: the library tracks every topic that receives a message in a `topics` table (via a Postgres trigger) and applies the default expiration automatically. Multiple outbox instances (e.g. replicas of the same service) coordinate cleanup using a per-topic maintenance lease, so only one instance runs the delete at a time. To log errors from the maintenance goroutines, pass a [zerolog](https://github.com/rs/zerolog) logger: ```go -outbox, err := pgoutbox.NewOutbox(pool, +outbox, err := pgoutbox.NewOutbox(ctx, pool, pgoutbox.WithTopicExpiration("orders", 24*time.Hour), pgoutbox.WithLogger(logger), ) @@ -176,8 +215,30 @@ if errors.Is(err, pgoutbox.ErrExclusiveLeaseHeld) { } ``` +An instance that never acquired the lease (or whose lease has expired) receives `pgoutbox.ErrExclusiveLeaseRequired` instead. + When the holder's context is cancelled, the lease expires naturally (within the lease duration, 30 s by default) and another instance's `AcquireTopic` call unblocks. +### Exclusive subscribers + +`Subscribe` composes with exclusive consumers: pass `WithExclusive()` and it manages the lease for you. + +```go +// Exactly one instance across the fleet drains "orders" at a time; the rest +// wait in line and take over on failure. +err := outbox.Subscribe(ctx, "orders", pgoutbox.WithExclusive()) +``` + +With `WithExclusive()`, `Subscribe`: + +1. acquires the lease before its first processing pass, blocking while another instance holds it (like `AcquireTopic`); +2. re-acquires it automatically if the lease is ever lost mid-subscribe (for example, a heartbeat lapse during a database blip); and +3. releases it on return, so a waiting instance takes over immediately instead of waiting out the lease's grace period. + +Several instances calling `Subscribe(..., WithExclusive())` on the same topic therefore form a failover group: one active consumer, the rest hot standbys. Nothing is missed during a handoff — the new holder's first drain pass covers any backlog that accumulated while the lease changed hands. + +Without `WithExclusive()`, subscribing to a topic whose exclusive lease is held elsewhere doesn't fail — every pass errors (visible via `WithLogger`) and is retried, so the subscriber sits idle until the lease frees up or is acquired. For an exclusive topic, either call `AcquireTopic` before `Subscribe` or pass `WithExclusive()`. + ## Benchmarks You can run benchmarks locally; for example, to write and flush 100k messages, you can run: @@ -186,7 +247,7 @@ You can run benchmarks locally; for example, to write and flush 100k messages, y go test -bench=. -benchtime=100000x ``` -On a local Macbook with an M3 Max core, this results in `8492 msgs/sec`: +`BenchmarkOutbox_WriteAndPublishThroughput` drains each topic with a busy-polling `ProcessMessages` loop; `BenchmarkOutbox_SubscribeThroughput` drains each topic with a `Subscribe` call woken by `LISTEN`/`NOTIFY` (its poll interval is set far above the benchmark runtime, so throughput is carried entirely by notifications — and each producer commit pays the in-transaction `pg_notify`). Both run at 1 and 10 topics, with producers spreading messages round-robin and one consumer per topic. On a local Macbook with an M3 Max core: ``` $ go test -bench=. -benchtime=100000x @@ -194,5 +255,12 @@ goos: darwin goarch: arm64 pkg: github.com/hatchet-dev/pgoutbox cpu: Apple M3 Max -BenchmarkOutbox_WriteAndPublishThroughput-14 100000 117757 ns/op 8492 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=1-14 100000 139432 ns/op 7172 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=1-14 100000 248696 ns/op 4021 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=10-14 100000 245562 ns/op 4072 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=10-14 100000 242325 ns/op 4127 msgs/sec +BenchmarkOutbox_SubscribeThroughput/Flush/topics=1-14 100000 188324 ns/op 5310 msgs/sec +BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=1-14 100000 305362 ns/op 3275 msgs/sec +BenchmarkOutbox_SubscribeThroughput/Flush/topics=10-14 100000 301837 ns/op 3313 msgs/sec +BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=10-14 100000 321406 ns/op 3111 msgs/sec ``` diff --git a/examples/go.mod b/examples/go.mod index e45bcd4..471aa57 100644 --- a/examples/go.mod +++ b/examples/go.mod @@ -10,13 +10,18 @@ require ( ) require ( + github.com/google/uuid v1.6.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/mattn/go-colorable v0.1.14 // indirect + github.com/mattn/go-isatty v0.0.21 // indirect github.com/mfridman/interpolate v0.0.2 // indirect github.com/pressly/goose/v3 v3.27.1 // indirect + github.com/rs/zerolog v1.35.1 // indirect github.com/sethvargo/go-retry v0.3.0 // indirect go.uber.org/multierr v1.11.0 // indirect golang.org/x/sync v0.20.0 // indirect + golang.org/x/sys v0.43.0 // indirect golang.org/x/text v0.36.0 // indirect ) diff --git a/examples/go.sum b/examples/go.sum index 2430756..e96a3a3 100644 --- a/examples/go.sum +++ b/examples/go.sum @@ -55,6 +55,8 @@ github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= github.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8SYxI99mE= github.com/magiconair/properties v1.8.10/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= +github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE= +github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8= github.com/mattn/go-isatty v0.0.21 h1:xYae+lCNBP7QuW4PUnNG61ffM4hVIfm+zUzDuSzYLGs= github.com/mattn/go-isatty v0.0.21/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4= github.com/mfridman/interpolate v0.0.2 h1:pnuTK7MQIxxFz1Gr+rjSIx9u7qVjf5VOoM/u6BbAxPY= @@ -91,6 +93,8 @@ github.com/pressly/goose/v3 v3.27.1 h1:6uEvcprBybDmW4hcz3gYujhARhye+GoWKhEWyzD5s github.com/pressly/goose/v3 v3.27.1/go.mod h1:maruOxsPnIG2yHHyo8UqKWXYKFcH7Q76csUV7+7KYoM= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/rs/zerolog v1.35.1 h1:m7xQeoiLIiV0BCEY4Hs+j2NG4Gp2o2KPKmhnnLiazKI= +github.com/rs/zerolog v1.35.1/go.mod h1:EjML9kdfa/RMA7h/6z6pYmq1ykOuA8/mjWaEvGI+jcw= github.com/sethvargo/go-retry v0.3.0 h1:EEt31A35QhrcRZtrYFDTBg91cqZVnFL2navjDrah2SE= github.com/sethvargo/go-retry v0.3.0/go.mod h1:mNX17F0C/HguQMyMyJxcnU471gOZGxCLyYaFyAZraas= github.com/shirou/gopsutil/v4 v4.26.3 h1:2ESdQt90yU3oXF/CdOlRCJxrP+Am1aBYubTMTfxJ1qc= diff --git a/examples/main.go b/examples/main.go index ab5e3fd..eca126f 100644 --- a/examples/main.go +++ b/examples/main.go @@ -3,9 +3,11 @@ package main import ( "context" "encoding/json" + "errors" "fmt" "log" "os" + "time" "github.com/hatchet-dev/pgoutbox" "github.com/hatchet-dev/pgoutbox/sqlc" @@ -19,7 +21,7 @@ const ( type printFlusher struct{} -func (printFlusher) Flush(_ context.Context, msgs []*sqlc.Message) error { +func (printFlusher) Flush(_ pgoutbox.FlushContext, msgs []*sqlc.Message) error { for _, m := range msgs { fmt.Printf(" flushed id=%d topic=%s payload=%s\n", m.ID, m.Topic, string(m.Payload)) } @@ -42,12 +44,27 @@ func main() { } defer pool.Close() - outbox, err := pgoutbox.NewOutbox(pool, pgoutbox.WithSchema(schema)) + // LISTEN/NOTIFY pubsub: Subscribe wakes as soon as AddMessages commits + // instead of waiting out its poll interval. + pubsub, err := pgoutbox.NewPGPubSub(ctx, pool) + if err != nil { + log.Fatalf("create pubsub: %v", err) + } + + outbox, err := pgoutbox.NewOutbox(ctx, pool, pgoutbox.WithSchema(schema), pgoutbox.WithPubSub(pubsub)) if err != nil { log.Fatalf("create outbox: %v", err) } outbox.AddFlusher(topic, printFlusher{}) + subCtx, cancel := context.WithCancel(ctx) + defer cancel() + + done := make(chan error, 1) + go func() { + done <- outbox.Subscribe(subCtx, topic, pgoutbox.WithPollInterval(30*time.Second)) + }() + tx, err := pool.Begin(ctx) if err != nil { log.Fatalf("begin: %v", err) @@ -67,11 +84,15 @@ func main() { log.Fatalf("commit: %v", err) } - fmt.Printf("staged %d messages on topic %q in schema %q\n", len(msgs), topic, schema) + fmt.Printf("staged %d messages on topic %q in schema %q; waiting for the subscriber...\n", len(msgs), topic, schema) + + // The commit above delivers the notification, so the subscriber flushes + // well within this window despite the 30s poll interval. + time.Sleep(2 * time.Second) + cancel() - fmt.Println("processing...") - if _, err := outbox.ProcessMessages(ctx, topic); err != nil { - log.Fatalf("process: %v", err) + if err := <-done; err != nil && !errors.Is(err, context.Canceled) { + log.Fatalf("subscribe: %v", err) } fmt.Println("done") } diff --git a/outbox.go b/outbox.go index 1b77446..468eab0 100644 --- a/outbox.go +++ b/outbox.go @@ -53,6 +53,19 @@ type Outbox interface { // the calling instance must hold the exclusive lease (via AcquireTopic) or an error is returned. ProcessMessages(ctx context.Context, topic string, opts ...ProcessOpt) ([]*sqlc.Message, error) + // Subscribe blocks and continuously drains the topic: it runs + // ProcessMessages until the topic is empty, then waits for the poll + // interval to elapse — or, when the outbox was built with WithPubSub, for + // a new-message notification — and drains again. Processing errors are + // logged to the WithLogger logger and retried on the next wake-up; as + // with ProcessMessages, topics with an active exclusive consumer require + // AcquireTopic first — either call it beforehand, or pass WithExclusive + // to have Subscribe acquire, re-acquire, and release the lease itself. + // Returns ctx.Err() when ctx ends, or an error immediately if no flusher + // is registered for the topic, the PubSub subscription cannot be + // established, or the WithExclusive initial acquisition fails. + Subscribe(ctx context.Context, topic string, opts ...SubscribeOpt) error + // AcquireTopic blocks until this instance holds the exclusive processing lease // for the named topic, then returns. A background goroutine automatically renews // the lease until ctx is cancelled or ReleaseTopic is called, at which point the @@ -74,6 +87,11 @@ type Outbox interface { // instance currently holds a valid exclusive lease for the topic. var ErrExclusiveLeaseHeld = errors.New("exclusive lease held by another instance") +// ErrExclusiveLeaseRequired is returned by ProcessMessages when the topic has +// an exclusive-consumer record but this instance does not hold a live lease — +// either AcquireTopic was never called or the lease has since expired. +var ErrExclusiveLeaseRequired = errors.New("exclusive lease required: call AcquireTopic first") + // defaultBatchSize is the number of messages ProcessMessages will pull per // call when the caller has not specified WithBatchSize. const defaultBatchSize = 1000 @@ -164,6 +182,7 @@ type outboxImplOpts struct { expirations map[string]time.Duration defaultExpiration time.Duration logger zerolog.Logger + pubsub PubSub } func defaultOpts() *outboxImplOpts { @@ -193,6 +212,11 @@ type outboxImpl struct { // logger receives error-level messages from the background maintenance goroutines. logger zerolog.Logger + // pubsub carries new-message notifications between AddMessages and + // Subscribe. Nil unless configured via WithPubSub; everything works + // without it, Subscribe just degrades to pure polling. + pubsub PubSub + flushers sync.Map // exclusiveLeaseWatcher holds one goroutine per held topic that waits for @@ -257,6 +281,19 @@ func WithLogger(l zerolog.Logger) OutboxOpt { } } +// WithPubSub attaches a PubSub used to cut end-to-end latency: AddMessages +// publishes a notification for each staged topic and Subscribe wakes on those +// notifications instead of waiting out its poll interval. Delivery is +// best-effort — Subscribe's polling remains the fallback for lost +// notifications. If ps also implements TxPublisher (NewPGPubSub does), the +// notification is published inside the AddMessages transaction and delivered +// exactly when it commits. +func WithPubSub(ps PubSub) OutboxOpt { + return func(opts *outboxImplOpts) { + opts.pubsub = ps + } +} + // NewOutbox creates an outbox backed by pool and starts the background // maintenance goroutines. The goroutines run until ctx is cancelled; pass a // context tied to your application lifetime (e.g. from signal.NotifyContext). @@ -292,6 +329,7 @@ func NewOutbox(ctx context.Context, pool *pgxpool.Pool, fs ...OutboxOpt) (Outbox expirations: expirations, defaultExpiration: opts.defaultExpiration, logger: opts.logger, + pubsub: opts.pubsub, managed: make(map[string]*managedTopic), } o.exclusiveLeaseWatcher = newLeaseWatcher(o.expireExclusiveLeaseAfterGrace) @@ -363,6 +401,32 @@ func (o *outboxImpl) AddMessages(ctx context.Context, tx pgx.Tx, topic string, m return fmt.Errorf("could not insert messages for topic %q: %w", topic, err) } + return o.publishNewMessageNotification(ctx, tx, topic) +} + +// publishNewMessageNotification wakes Subscribe callers after AddMessages +// stages messages. It prefers the TxPublisher path — publishing on the +// caller's transaction so the notification lands exactly when the insert +// commits — and falls back to an immediate best-effort Pub otherwise. +func (o *outboxImpl) publishNewMessageNotification(ctx context.Context, tx pgx.Tx, topic string) error { + if o.pubsub == nil { + return nil + } + + if txp, ok := o.pubsub.(TxPublisher); ok { + // This shares the caller's transaction: a failure has aborted it, so + // surface the error rather than pretending the insert succeeded. + if err := txp.PubInTx(ctx, tx, topic, nil); err != nil { + return fmt.Errorf("could not publish new-message notification for topic %q: %w", topic, err) + } + return nil + } + + // Out-of-band publish is best-effort: the messages are durably staged and + // pollers pick them up within a poll interval if the notification is lost. + if err := o.pubsub.Pub(ctx, topic, nil); err != nil { + o.logger.Error().Err(err).Str("topic", topic).Msg("subscribe: failed to publish new-message notification") + } return nil } @@ -584,7 +648,7 @@ func (o *outboxImpl) checkExclusiveAccessForUpdate(ctx context.Context, wrapped if row.ExclusiveConsumerExpiresAt.Valid && row.ExclusiveConsumerExpiresAt.Time.After(time.Now()) { return nil } - return fmt.Errorf("exclusive lease for topic %q has expired; call AcquireTopic to renew", topic) + return fmt.Errorf("exclusive lease for topic %q has expired: %w", topic, ErrExclusiveLeaseRequired) } if row.ExclusiveConsumerExpiresAt.Valid && row.ExclusiveConsumerExpiresAt.Time.After(time.Now()) { @@ -592,7 +656,7 @@ func (o *outboxImpl) checkExclusiveAccessForUpdate(ctx context.Context, wrapped } // Another instance held the lease but it has expired; require explicit acquire. - return fmt.Errorf("exclusive access required for topic %q: call AcquireTopic first", topic) + return fmt.Errorf("exclusive access required for topic %q: %w", topic, ErrExclusiveLeaseRequired) } func (o *outboxImpl) ProcessMessages(ctx context.Context, topic string, popts ...ProcessOpt) ([]*sqlc.Message, error) { diff --git a/outbox_bench_test.go b/outbox_bench_test.go index c0ea03b..9c9c7cc 100644 --- a/outbox_bench_test.go +++ b/outbox_bench_test.go @@ -2,6 +2,7 @@ package pgoutbox_test import ( "context" + "errors" "fmt" "sync" "sync/atomic" @@ -47,27 +48,52 @@ func (c *txCountingFlusher) Flush(ctx pgoutbox.FlushContext, msgs []*sqlc.Messag return nil } +// benchTopicCounts is the topic-count dimension of the benchmark matrix: +// producers spray messages round-robin across the topics and one consumer +// drains each topic. +var benchTopicCounts = []int{1, 10} + func BenchmarkOutbox_WriteAndPublishThroughput(b *testing.B) { - b.Run("Flush", func(b *testing.B) { - benchmarkThroughput(b, func(_ string, onFlush func(int)) pgoutbox.Flusher { - return &countingFlusher{onFlush: onFlush} + benchmarkThroughputMatrix(b, false) +} + +// BenchmarkOutbox_SubscribeThroughput is the notification-driven counterpart +// of BenchmarkOutbox_WriteAndPublishThroughput: producers stage messages as +// fast as they can (each commit publishing a pg_notify via the attached +// PGPubSub) while one Subscribe call per topic drains it. The poll interval +// is set far above the benchmark's runtime, so throughput here is carried by +// the LISTEN/NOTIFY wake-ups, not polling. +func BenchmarkOutbox_SubscribeThroughput(b *testing.B) { + benchmarkThroughputMatrix(b, true) +} + +// benchmarkThroughputMatrix runs the Flush/TxFlush × topic-count grid. +func benchmarkThroughputMatrix(b *testing.B, useSubscribe bool) { + for _, numTopics := range benchTopicCounts { + b.Run(fmt.Sprintf("Flush/topics=%d", numTopics), func(b *testing.B) { + benchmarkThroughput(b, func(_ string, onFlush func(int)) pgoutbox.Flusher { + return &countingFlusher{onFlush: onFlush} + }, useSubscribe, numTopics) }) - }) - b.Run("TxFlush", func(b *testing.B) { - benchmarkThroughput(b, func(schema string, onFlush func(int)) pgoutbox.Flusher { - table := pgx.Identifier{schema, "bench_side_log"}.Sanitize() - _, err := sharedPool.Exec(context.Background(), fmt.Sprintf("CREATE TABLE %s (msg_id bigint NOT NULL)", table)) - require.NoError(b, err) - return &txCountingFlusher{table: table, onFlush: onFlush} + b.Run(fmt.Sprintf("TxFlush/topics=%d", numTopics), func(b *testing.B) { + benchmarkThroughput(b, func(schema string, onFlush func(int)) pgoutbox.Flusher { + table := pgx.Identifier{schema, "bench_side_log"}.Sanitize() + _, err := sharedPool.Exec(context.Background(), fmt.Sprintf("CREATE TABLE %s (msg_id bigint NOT NULL)", table)) + require.NoError(b, err) + return &txCountingFlusher{table: table, onFlush: onFlush} + }, useSubscribe, numTopics) }) - }) + } } // benchmarkThroughput drives a producer/consumer loop against a freshly-built // outbox. newFlusher builds the flusher under test from the test schema and the // inFlight-releasing callback, letting us reuse the harness for both the plain -// Flush path and the FlushWithTx path. -func benchmarkThroughput(b *testing.B, newFlusher func(schema string, onFlush func(n int)) pgoutbox.Flusher) { +// Flush path and the FlushWithTx path. With useSubscribe each topic's consumer +// is a Subscribe call woken by pg_notify (and AddMessages pays the in-tx +// publish); otherwise it is a busy-polling ProcessMessages loop. Producers +// assign messages to the numTopics topics round-robin. +func benchmarkThroughput(b *testing.B, newFlusher func(schema string, onFlush func(n int)) pgoutbox.Flusher, useSubscribe bool, numTopics int) { const ( numWorkers = MAX_CONNS maxInFlight = 5000 @@ -79,61 +105,92 @@ func benchmarkThroughput(b *testing.B, newFlusher func(schema string, onFlush fu schema := uniqueSchema(b) - outbox, err := pgoutbox.NewOutbox( - ctx, - sharedPool, - pgoutbox.WithSchema(schema), - ) + opts := []pgoutbox.OutboxOpt{pgoutbox.WithSchema(schema)} + if useSubscribe { + // The schema name doubles as the NOTIFY channel so concurrent + // benchmarks don't wake each other's subscribers. + ps, err := pgoutbox.NewPGPubSub(ctx, sharedPool, pgoutbox.WithNotifyChannel(schema)) + require.NoError(b, err) + opts = append(opts, pgoutbox.WithPubSub(ps)) + } + + outbox, err := pgoutbox.NewOutbox(ctx, sharedPool, opts...) require.NoError(b, err) + topics := make([]string, numTopics) + for i := range topics { + topics[i] = fmt.Sprintf("bench_%d", i) + } + // inFlight bounds the number of un-flushed rows allowed in the table at // once: producers acquire a slot before INSERT, the flusher releases it // after the row is committed-deleted. inFlight := make(chan struct{}, maxInFlight) var flushed atomic.Int64 - outbox.AddFlusher("bench", newFlusher(schema, func(n int) { + flusher := newFlusher(schema, func(n int) { for range n { <-inFlight } flushed.Add(int64(n)) - })) + }) + for _, topic := range topics { + outbox.AddFlusher(topic, flusher) + } procCtx, stopProcessor := context.WithCancel(ctx) var procWg sync.WaitGroup - procWg.Go(func() { - for procCtx.Err() == nil { - msgs, err := outbox.ProcessMessages(procCtx, "bench", pgoutbox.WithBatchSize(batchSize)) - if err != nil { - if procCtx.Err() != nil { - return + for _, topic := range topics { + if useSubscribe { + procWg.Go(func() { + // The poll interval is a safety net well above the benchmark's + // runtime: every drain below must be triggered by a notification. + err := outbox.Subscribe(procCtx, topic, + pgoutbox.WithPollInterval(30*time.Second), + pgoutbox.WithProcessOpts(pgoutbox.WithBatchSize(batchSize)), + ) + if err != nil && !errors.Is(err, context.Canceled) { + b.Errorf("Subscribe %s: %v", topic, err) } - b.Logf("ProcessMessages: %v", err) - continue - } - if len(msgs) == 0 { - time.Sleep(time.Millisecond) - } + }) + } else { + procWg.Go(func() { + for procCtx.Err() == nil { + msgs, err := outbox.ProcessMessages(procCtx, topic, pgoutbox.WithBatchSize(batchSize)) + if err != nil { + if procCtx.Err() != nil { + return + } + b.Logf("ProcessMessages %s: %v", topic, err) + continue + } + if len(msgs) == 0 { + time.Sleep(time.Millisecond) + } + } + }) } - }) + } payload := []byte(`{"id":1}`) b.ResetTimer() work := make(chan struct{}) + var msgIdx atomic.Int64 var wg sync.WaitGroup for range numWorkers { wg.Go(func() { for range work { inFlight <- struct{}{} + topic := topics[msgIdx.Add(1)%int64(numTopics)] tx, err := sharedPool.Begin(ctx) if err != nil { b.Errorf("begin: %v", err) return } - if err := outbox.AddMessages(ctx, tx, "bench", []pgoutbox.MessageOpts{ + if err := outbox.AddMessages(ctx, tx, topic, []pgoutbox.MessageOpts{ {Payload: payload}, }); err != nil { _ = tx.Rollback(ctx) @@ -155,6 +212,9 @@ func benchmarkThroughput(b *testing.B, newFlusher func(schema string, onFlush fu wg.Wait() for flushed.Load() < int64(b.N) { + if ctx.Err() != nil || b.Failed() { + b.Fatalf("consumer stalled: flushed %d of %d messages", flushed.Load(), b.N) + } time.Sleep(time.Millisecond) } diff --git a/pgpubsub.go b/pgpubsub.go new file mode 100644 index 0000000..63d56d6 --- /dev/null +++ b/pgpubsub.go @@ -0,0 +1,314 @@ +package pgoutbox + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/rs/zerolog" +) + +// defaultNotifyChannel is the Postgres NOTIFY channel all pgPubSub messages +// are multiplexed over. Every message carries its pub/sub topic in the JSON +// envelope, so one LISTEN connection serves any number of topics. +const defaultNotifyChannel = "pgoutbox_pubsub" + +// subscriberBufferSize is the capacity of each subscription channel. When a +// subscriber falls this far behind, further messages are dropped rather than +// blocking the listener — acceptable because delivery is best-effort and the +// outbox's notifications are pure wake-up signals that coalesce naturally. +const subscriberBufferSize = 16 + +// pgNotifyReconnectDelay is how long the listener waits before redialing +// after its connection fails. Exposed as a var for tests. +var pgNotifyReconnectDelay = newAtomicDuration(5 * time.Second) + +type pgPubSubOpts struct { + channel string + logger zerolog.Logger +} + +func defaultPGPubSubOpts() *pgPubSubOpts { + return &pgPubSubOpts{ + channel: defaultNotifyChannel, + logger: zerolog.Nop(), + } +} + +// PGPubSubOpt configures the PubSub returned by NewPGPubSub. +type PGPubSubOpt func(*pgPubSubOpts) + +// WithNotifyChannel overrides the Postgres NOTIFY channel the PubSub +// multiplexes over. All messages on a channel are broadcast to every listener +// of that channel, so two outboxes sharing a database (e.g. different +// schemas) should use distinct channels to avoid spurious wake-ups. +func WithNotifyChannel(name string) PGPubSubOpt { + return func(opts *pgPubSubOpts) { + opts.channel = name + } +} + +// WithNotifyLogger attaches a zerolog logger that receives errors from the +// background listener (connection failures, malformed payloads). If not set, +// those errors are silent. +func WithNotifyLogger(l zerolog.Logger) PGPubSubOpt { + return func(opts *pgPubSubOpts) { + opts.logger = l + } +} + +// pgPubSub implements PubSub (and TxPublisher) on top of Postgres +// LISTEN/NOTIFY. All messages travel over a single NOTIFY channel wrapped in +// a JSON envelope carrying the pub/sub topic; a lone background listener +// dispatches them to in-process subscribers. +// +// The listener runs on a connection hijacked out of the pool: it holds LISTEN +// session state and blocks in WaitForNotification indefinitely, so it must +// neither occupy a pool slot nor ever be handed back for reuse. +type pgPubSub struct { + pool *pgxpool.Pool + channel string + logger zerolog.Logger + + // ctx bounds the background listener; it is the ctx passed to NewPGPubSub. + ctx context.Context + + // listenMu guards listening, the lazy one-time start of the listener + // goroutine. Held across the initial dial so that Sub only returns once + // LISTEN is active — messages published after a successful Sub are + // delivered (barring connection loss) rather than racing the setup. + listenMu sync.Mutex + listening bool + + subscribersMu sync.Mutex + subscribers map[string]map[chan *PubSubMessage]struct{} +} + +// NewPGPubSub returns a PubSub backed by Postgres LISTEN/NOTIFY on the given +// pool. The background listener starts lazily on the first Sub call and runs +// until ctx is cancelled; pass a context tied to your application lifetime. +// +// The returned PubSub implements TxPublisher, so an outbox configured with it +// publishes new-message notifications transactionally: subscribers wake when +// the staging transaction commits, and not at all if it rolls back. +// +// NOTIFY payloads are capped by Postgres at roughly 8000 bytes; Pub returns +// an error beyond that. The outbox's own notifications are empty. +func NewPGPubSub(ctx context.Context, pool *pgxpool.Pool, fs ...PGPubSubOpt) (PubSub, error) { + opts := defaultPGPubSubOpts() + + for _, f := range fs { + f(opts) + } + + // The channel name is interpolated into the LISTEN statement, so restrict + // it to identifier-safe characters just like the schema name. + if !schemaNameRE.MatchString(opts.channel) { + return nil, fmt.Errorf("invalid notify channel name %q: must match %s", opts.channel, schemaNameRE) + } + + return &pgPubSub{ + pool: pool, + channel: opts.channel, + logger: opts.logger, + ctx: ctx, + subscribers: make(map[string]map[chan *PubSubMessage]struct{}), + }, nil +} + +func marshalPubSubMessage(topic string, payload []byte) ([]byte, error) { + if topic == "" { + return nil, fmt.Errorf("topic must not be empty") + } + wrapped, err := json.Marshal(&PubSubMessage{Topic: topic, Payload: payload}) + if err != nil { + return nil, fmt.Errorf("could not marshal pubsub message for topic %q: %w", topic, err) + } + return wrapped, nil +} + +func (p *pgPubSub) Pub(ctx context.Context, topic string, payload []byte) error { + wrapped, err := marshalPubSubMessage(topic, payload) + if err != nil { + return err + } + if _, err := p.pool.Exec(ctx, "select pg_notify($1, $2)", p.channel, string(wrapped)); err != nil { + return fmt.Errorf("could not publish to topic %q: %w", topic, err) + } + return nil +} + +// PubInTx publishes on the caller's transaction, deferring delivery to commit +// time. Postgres deduplicates identical notifications within a transaction, +// so repeated AddMessages calls for one topic in one transaction cost a +// single wake-up. +func (p *pgPubSub) PubInTx(ctx context.Context, tx pgx.Tx, topic string, payload []byte) error { + wrapped, err := marshalPubSubMessage(topic, payload) + if err != nil { + return err + } + if _, err := tx.Exec(ctx, "select pg_notify($1, $2)", p.channel, string(wrapped)); err != nil { + return fmt.Errorf("could not publish to topic %q: %w", topic, err) + } + return nil +} + +func (p *pgPubSub) Sub(ctx context.Context, topic string) (<-chan *PubSubMessage, error) { + if topic == "" { + return nil, fmt.Errorf("topic must not be empty") + } + if err := p.ctx.Err(); err != nil { + return nil, fmt.Errorf("pubsub has shut down: %w", err) + } + + if err := p.ensureListening(ctx); err != nil { + return nil, err + } + + ch := make(chan *PubSubMessage, subscriberBufferSize) + + p.subscribersMu.Lock() + subs := p.subscribers[topic] + if subs == nil { + subs = make(map[chan *PubSubMessage]struct{}) + p.subscribers[topic] = subs + } + subs[ch] = struct{}{} + p.subscribersMu.Unlock() + + go func() { + select { + case <-ctx.Done(): + case <-p.ctx.Done(): + } + p.unsubscribe(topic, ch) + }() + + return ch, nil +} + +func (p *pgPubSub) unsubscribe(topic string, ch chan *PubSubMessage) { + p.subscribersMu.Lock() + defer p.subscribersMu.Unlock() + + subs := p.subscribers[topic] + if _, ok := subs[ch]; !ok { + return + } + delete(subs, ch) + if len(subs) == 0 { + delete(p.subscribers, topic) + } + // Closing under subscribersMu is what makes it safe: dispatch sends while + // holding the same lock, so a send on a closed channel is impossible. + close(ch) +} + +func (p *pgPubSub) dispatch(msg *PubSubMessage) { + p.subscribersMu.Lock() + defer p.subscribersMu.Unlock() + + for ch := range p.subscribers[msg.Topic] { + select { + case ch <- msg: + default: + // Subscriber buffer full — drop rather than block the listener. + } + } +} + +// ensureListening lazily starts the single background listener. The initial +// dial happens synchronously so a successful Sub means LISTEN is already +// active; reconnects after that are handled in the background. +func (p *pgPubSub) ensureListening(ctx context.Context) error { + p.listenMu.Lock() + defer p.listenMu.Unlock() + + if p.listening { + return nil + } + + conn, err := p.connectAndListen(ctx) + if err != nil { + return fmt.Errorf("could not start pubsub listener: %w", err) + } + + p.listening = true + go p.runListener(conn) + return nil +} + +// connectAndListen acquires a connection, hijacks it out of the pool, and +// issues LISTEN on it. On success the caller owns the connection and is +// responsible for closing it. +func (p *pgPubSub) connectAndListen(ctx context.Context) (*pgx.Conn, error) { + poolConn, err := p.pool.Acquire(ctx) + if err != nil { + return nil, fmt.Errorf("acquire listener connection: %w", err) + } + + conn := poolConn.Hijack() + + if _, err := conn.Exec(ctx, "listen "+pgx.Identifier{p.channel}.Sanitize()); err != nil { + _ = conn.Close(context.WithoutCancel(ctx)) + return nil, fmt.Errorf("listen on channel %q: %w", p.channel, err) + } + + return conn, nil +} + +// runListener serves notifications off conn until p.ctx ends, redialing with +// a delay whenever the connection fails. Notifications sent while +// disconnected are lost — subscribers are expected to treat delivery as +// best-effort (Subscribe covers gaps with its poll interval). +func (p *pgPubSub) runListener(conn *pgx.Conn) { + for { + err := p.serveConn(conn) + _ = conn.Close(context.WithoutCancel(p.ctx)) + + if p.ctx.Err() != nil { + return + } + p.logger.Error().Err(err).Msg("pubsub: listener connection failed; reconnecting") + + for { + select { + case <-p.ctx.Done(): + return + case <-time.After(pgNotifyReconnectDelay.Load()): + } + + next, err := p.connectAndListen(p.ctx) + if err != nil { + if p.ctx.Err() != nil { + return + } + p.logger.Error().Err(err).Msg("pubsub: failed to reconnect listener") + continue + } + conn = next + break + } + } +} + +func (p *pgPubSub) serveConn(conn *pgx.Conn) error { + for { + notification, err := conn.WaitForNotification(p.ctx) + if err != nil { + return err + } + + msg := &PubSubMessage{} + if err := json.Unmarshal([]byte(notification.Payload), msg); err != nil { + p.logger.Error().Err(err).Msg("pubsub: dropping notification with malformed payload") + continue + } + + p.dispatch(msg) + } +} diff --git a/pubsub.go b/pubsub.go new file mode 100644 index 0000000..3f96dc6 --- /dev/null +++ b/pubsub.go @@ -0,0 +1,49 @@ +package pgoutbox + +import ( + "context" + + "github.com/jackc/pgx/v5" +) + +// PubSubMessage is a single message delivered by a PubSub subscription. +type PubSubMessage struct { + // Topic is the pub/sub topic the message was published to. + Topic string `json:"topic"` + + // Payload is the opaque message body. It may be nil: the outbox's own + // new-message notifications carry no payload, since the notification + // itself is the signal to check the outbox. + Payload []byte `json:"payload,omitempty"` +} + +// PubSub is a minimal publish/subscribe transport for small notification +// messages. The outbox uses it (via WithPubSub) to wake Subscribe callers as +// soon as new messages are staged, instead of waiting out a poll interval. +// +// Delivery is expected to be best-effort: implementations may drop messages +// under load or while disconnected. The outbox tolerates both lost messages +// (Subscribe falls back to polling) and duplicate or spurious messages (an +// extra processing pass on an empty topic is a no-op). +type PubSub interface { + // Pub publishes payload to topic. + Pub(ctx context.Context, topic string, payload []byte) error + + // Sub subscribes to topic and returns a channel of messages published to + // it. The subscription lasts until ctx ends (or the PubSub itself shuts + // down), at which point the channel is closed. The channel should be + // buffered; implementations may drop messages rather than block when a + // slow consumer's buffer is full. + Sub(ctx context.Context, topic string) (<-chan *PubSubMessage, error) +} + +// TxPublisher is an optional interface a PubSub can implement to publish +// within a pgx transaction. When the PubSub configured via WithPubSub +// implements it, AddMessages publishes its new-message notification inside +// the caller's transaction, so the notification is delivered exactly when the +// insert commits — and never for a transaction that rolls back. Without it, +// AddMessages falls back to a best-effort Pub at insert time, which can wake +// subscribers before the messages are visible. +type TxPublisher interface { + PubInTx(ctx context.Context, tx pgx.Tx, topic string, payload []byte) error +} diff --git a/subscribe.go b/subscribe.go new file mode 100644 index 0000000..081b988 --- /dev/null +++ b/subscribe.go @@ -0,0 +1,177 @@ +package pgoutbox + +import ( + "context" + "errors" + "fmt" + "time" +) + +// defaultPollInterval is how often Subscribe re-checks the topic when the +// caller has not specified WithPollInterval. With a PubSub configured the +// poll is only a safety net for lost notifications, so it can be generous. +const defaultPollInterval = 5 * time.Second + +// SubscribeOpt is a per-call option for Subscribe. +type SubscribeOpt func(*subscribeOpts) + +type subscribeOpts struct { + pollInterval time.Duration + processOpts []ProcessOpt + exclusive bool +} + +func defaultSubscribeOpts() *subscribeOpts { + return &subscribeOpts{pollInterval: defaultPollInterval} +} + +// WithPollInterval sets how long Subscribe waits between processing passes +// when no new-message notification arrives. Must be > 0. +func WithPollInterval(d time.Duration) SubscribeOpt { + return func(opts *subscribeOpts) { + if d <= 0 { + return + } + opts.pollInterval = d + } +} + +// WithProcessOpts forwards per-call ProcessMessages options (e.g. +// WithBatchSize) to every processing pass Subscribe makes. +func WithProcessOpts(popts ...ProcessOpt) SubscribeOpt { + return func(opts *subscribeOpts) { + opts.processOpts = append(opts.processOpts, popts...) + } +} + +// WithExclusive makes Subscribe manage the topic's exclusive-consumer lease +// for the duration of the call: it acquires the lease before the first +// processing pass (blocking, like AcquireTopic, while another instance holds +// it), re-acquires it if it is ever lost mid-subscribe, and releases it on +// return so a waiting instance can take over immediately instead of waiting +// out the lease's grace period. Several instances calling Subscribe with +// WithExclusive on the same topic therefore form a failover group: exactly +// one drains the topic while the rest block in line behind the lease. +func WithExclusive() SubscribeOpt { + return func(opts *subscribeOpts) { + opts.exclusive = true + } +} + +func (o *outboxImpl) Subscribe(ctx context.Context, topic string, sopts ...SubscribeOpt) error { + if topic == "" { + return fmt.Errorf("topic must not be empty") + } + if _, ok := o.getFlusher(topic); !ok { + return fmt.Errorf("no flusher registered for topic %q", topic) + } + + so := defaultSubscribeOpts() + for _, opt := range sopts { + opt(so) + } + + if so.exclusive { + if err := o.AcquireTopic(ctx, topic); err != nil { + return fmt.Errorf("could not acquire exclusive lease for topic %q: %w", topic, err) + } + // Hand the lease off immediately on exit instead of letting it lapse + // through the grace period. Detached ctx because the usual way out is + // this ctx being cancelled; a failed release degrades to the + // grace-period lapse inside ReleaseTopic itself. + defer func() { + if err := o.ReleaseTopic(context.WithoutCancel(ctx), topic); err != nil { + o.logger.Error().Err(err).Str("topic", topic).Msg("subscribe: failed to release exclusive lease") + } + }() + } + + var notifications <-chan *PubSubMessage + if o.pubsub != nil { + ch, err := o.pubsub.Sub(ctx, topic) + if err != nil { + return fmt.Errorf("could not subscribe to notifications for topic %q: %w", topic, err) + } + notifications = ch + } + + for { + // Anything already buffered is covered by the pass below — drop it now + // so it doesn't immediately trigger a redundant pass. A notification + // arriving after this point stays in the buffer and wakes the select, + // which is correct: its messages may have committed after our pass. + drainNotifications(notifications) + + err := o.processTopicUntilEmpty(ctx, topic, so.processOpts) + + if so.exclusive && isExclusiveLeaseError(err) { + // The lease was lost mid-subscribe (session lapse, takeover after + // a grace expiry). Block until we hold it again, then drain right + // away: notifications that arrived meanwhile coalesced in the + // buffer, and the post-acquire pass covers the backlog regardless. + if acqErr := o.AcquireTopic(ctx, topic); acqErr != nil { + if ctx.Err() != nil { + return ctx.Err() + } + // Transient failure: fall through to the select and let the + // next wake-up retry the pass (and this re-acquisition). + o.logger.Error().Err(acqErr).Str("topic", topic).Msg("subscribe: failed to re-acquire exclusive lease") + } else { + continue + } + } + + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(so.pollInterval): + case _, ok := <-notifications: + if !ok { + // The PubSub shut down before our ctx did; degrade to polling. + notifications = nil + } + } + } +} + +// isExclusiveLeaseError reports whether a ProcessMessages failure means this +// instance does not (or no longer) hold the topic's exclusive lease. +func isExclusiveLeaseError(err error) bool { + return errors.Is(err, ErrExclusiveLeaseHeld) || errors.Is(err, ErrExclusiveLeaseRequired) +} + +// processTopicUntilEmpty runs ProcessMessages until the topic comes back +// empty, returning the error that ended the pass (nil when the topic was +// drained). Errors are logged here and implicitly retried when the next poll +// tick or notification wakes the Subscribe loop; the return value exists so +// Subscribe can recognize lease loss and re-acquire. +func (o *outboxImpl) processTopicUntilEmpty(ctx context.Context, topic string, popts []ProcessOpt) error { + for ctx.Err() == nil { + msgs, err := o.ProcessMessages(ctx, topic, popts...) + if err != nil { + if !errors.Is(err, context.Canceled) { + o.logger.Error().Err(err).Str("topic", topic).Msg("subscribe: failed to process messages") + } + return err + } + if len(msgs) == 0 { + return nil + } + } + return ctx.Err() +} + +// drainNotifications empties ch without blocking. A nil or closed channel is +// a no-op. +func drainNotifications(ch <-chan *PubSubMessage) { + for { + select { + case _, ok := <-ch: + if !ok { + return + } + default: + return + } + } +} diff --git a/subscribe_e2e_test.go b/subscribe_e2e_test.go new file mode 100644 index 0000000..88f4a51 --- /dev/null +++ b/subscribe_e2e_test.go @@ -0,0 +1,346 @@ +package pgoutbox_test + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/hatchet-dev/pgoutbox" +) + +// newTestPubSub returns a PGPubSub on the shared pool, isolated on its own +// NOTIFY channel so parallel tests don't wake each other up. Schema names are +// identifier-safe, so they double as channel names. +func newTestPubSub(t *testing.T, ctx context.Context, channel string) pgoutbox.PubSub { + t.Helper() + ps, err := pgoutbox.NewPGPubSub(ctx, sharedPool, pgoutbox.WithNotifyChannel(channel)) + require.NoError(t, err) + return ps +} + +func TestPGPubSub_PubAndSub(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + ps := newTestPubSub(t, ctx, uniqueSchema(t)) + + orders, err := ps.Sub(ctx, "orders") + require.NoError(t, err) + shipments, err := ps.Sub(ctx, "shipments") + require.NoError(t, err) + + require.NoError(t, ps.Pub(ctx, "orders", []byte(`{"hello":"world"}`))) + + select { + case msg := <-orders: + require.NotNil(t, msg) + assert.Equal(t, "orders", msg.Topic) + assert.JSONEq(t, `{"hello":"world"}`, string(msg.Payload)) + case <-time.After(15 * time.Second): + t.Fatal("timed out waiting for published message") + } + + // The shipments subscriber must not have seen the orders message. + select { + case msg := <-shipments: + t.Fatalf("shipments subscriber received message for topic %q", msg.Topic) + default: + } +} + +func TestPGPubSub_SubChannelClosesOnCtxCancel(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + ps := newTestPubSub(t, ctx, uniqueSchema(t)) + + subCtx, subCancel := context.WithCancel(ctx) + ch, err := ps.Sub(subCtx, "orders") + require.NoError(t, err) + + subCancel() + + select { + case _, ok := <-ch: + assert.False(t, ok, "channel should be closed, not carrying a message") + case <-time.After(15 * time.Second): + t.Fatal("timed out waiting for subscription channel to close") + } +} + +// TestOutbox_SubscribeWakesOnNotification verifies the notification fast +// path: with a poll interval far longer than the test, messages still flush +// promptly after commit — and, because the PGPubSub publishes inside the +// AddMessages transaction, not before it. +func TestOutbox_SubscribeWakesOnNotification(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + schema := uniqueSchema(t) + ps := newTestPubSub(t, ctx, schema) + + outbox, err := pgoutbox.NewOutbox(ctx, sharedPool, pgoutbox.WithSchema(schema), pgoutbox.WithPubSub(ps)) + require.NoError(t, err) + + flusher := &captureFlusher{} + outbox.AddFlusher("orders", flusher) + + subCtx, subCancel := context.WithCancel(ctx) + defer subCancel() + + done := make(chan error, 1) + go func() { + done <- outbox.Subscribe(subCtx, "orders", pgoutbox.WithPollInterval(time.Minute)) + }() + + tx, err := sharedPool.Begin(ctx) + require.NoError(t, err) + require.NoError(t, outbox.AddMessages(ctx, tx, "orders", []pgoutbox.MessageOpts{ + {Payload: mustPayload(t, map[string]int{"id": 1})}, + {Payload: mustPayload(t, map[string]int{"id": 2})}, + {Payload: mustPayload(t, map[string]int{"id": 3})}, + })) + + // The transaction is still open: no notification has been delivered and + // the messages are invisible, so nothing may flush yet. + time.Sleep(1500 * time.Millisecond) + assert.Empty(t, flusher.Received(), "nothing should flush before the staging tx commits") + + require.NoError(t, tx.Commit(ctx)) + + // Poll interval is one minute, so delivery within a few seconds proves + // the pg_notify wake-up worked. The flusher records messages before the + // processing tx commits, so also wait for the delete to land. + require.Eventually(t, func() bool { + return len(flusher.Received()) == 3 && countMessages(t, ctx, schema, "orders") == 0 + }, 15*time.Second, 50*time.Millisecond, "subscriber should flush promptly on notification") + + subCancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(15 * time.Second): + t.Fatal("Subscribe did not return after ctx cancellation") + } +} + +// TestOutbox_SubscribePollsWithoutPubSub verifies the polling fallback: with +// no PubSub configured, Subscribe still picks up messages within a poll +// interval. +func TestOutbox_SubscribePollsWithoutPubSub(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + schema := uniqueSchema(t) + + outbox, err := pgoutbox.NewOutbox(ctx, sharedPool, pgoutbox.WithSchema(schema)) + require.NoError(t, err) + + flusher := &captureFlusher{} + outbox.AddFlusher("orders", flusher) + + subCtx, subCancel := context.WithCancel(ctx) + defer subCancel() + + done := make(chan error, 1) + go func() { + done <- outbox.Subscribe(subCtx, "orders", pgoutbox.WithPollInterval(100*time.Millisecond)) + }() + + tx, err := sharedPool.Begin(ctx) + require.NoError(t, err) + require.NoError(t, outbox.AddMessages(ctx, tx, "orders", []pgoutbox.MessageOpts{ + {Payload: mustPayload(t, map[string]int{"id": 1})}, + {Payload: mustPayload(t, map[string]int{"id": 2})}, + })) + require.NoError(t, tx.Commit(ctx)) + + // As above, wait for both the flush and the committed delete. + require.Eventually(t, func() bool { + return len(flusher.Received()) == 2 && countMessages(t, ctx, schema, "orders") == 0 + }, 15*time.Second, 50*time.Millisecond, "subscriber should flush via polling") + + subCancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(15 * time.Second): + t.Fatal("Subscribe did not return after ctx cancellation") + } +} + +// addTestMessages stages and commits n messages for topic through ob. +func addTestMessages(t *testing.T, ctx context.Context, ob pgoutbox.Outbox, topic string, n int) { + t.Helper() + + tx, err := sharedPool.Begin(ctx) + require.NoError(t, err) + defer tx.Rollback(ctx) + + msgs := make([]pgoutbox.MessageOpts, n) + for i := range msgs { + msgs[i] = pgoutbox.MessageOpts{Payload: mustPayload(t, map[string]int{"id": i})} + } + require.NoError(t, ob.AddMessages(ctx, tx, topic, msgs)) + require.NoError(t, tx.Commit(ctx)) +} + +// TestOutbox_SubscribeExclusiveFailover verifies the WithExclusive failover +// group: instance A acquires the lease and drains; standby B blocks behind +// the lease and processes nothing; when A's Subscribe exits, its +// release-on-return hands the lease to B, which takes over promptly (well +// within the poll interval, via the acquire retry loop plus its post-acquire +// drain pass). +func TestOutbox_SubscribeExclusiveFailover(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + schema := uniqueSchema(t) + ps := newTestPubSub(t, ctx, schema) + + obA, err := pgoutbox.NewOutbox(ctx, sharedPool, pgoutbox.WithSchema(schema), pgoutbox.WithPubSub(ps)) + require.NoError(t, err) + obB, err := pgoutbox.NewOutbox(ctx, sharedPool, pgoutbox.WithSchema(schema), pgoutbox.WithPubSub(ps)) + require.NoError(t, err) + + flusherA := &captureFlusher{} + flusherB := &captureFlusher{} + obA.AddFlusher("orders", flusherA) + obB.AddFlusher("orders", flusherB) + + ctxA, cancelA := context.WithCancel(ctx) + defer cancelA() + doneA := make(chan error, 1) + go func() { + doneA <- obA.Subscribe(ctxA, "orders", pgoutbox.WithExclusive(), pgoutbox.WithPollInterval(time.Minute)) + }() + + addTestMessages(t, ctx, obA, "orders", 2) + require.Eventually(t, func() bool { + return len(flusherA.Received()) == 2 && countMessages(t, ctx, schema, "orders") == 0 + }, 15*time.Second, 50*time.Millisecond, "A should acquire the lease and drain the first batch") + + // B lines up as the standby: its Subscribe blocks inside AcquireTopic + // while A holds the lease. + ctxB, cancelB := context.WithCancel(ctx) + defer cancelB() + doneB := make(chan error, 1) + go func() { + doneB <- obB.Subscribe(ctxB, "orders", pgoutbox.WithExclusive(), pgoutbox.WithPollInterval(time.Minute)) + }() + + addTestMessages(t, ctx, obA, "orders", 2) + require.Eventually(t, func() bool { + return len(flusherA.Received()) == 4 && countMessages(t, ctx, schema, "orders") == 0 + }, 15*time.Second, 50*time.Millisecond, "A should keep draining while it holds the lease") + assert.Empty(t, flusherB.Received(), "standby B must not process while A holds the lease") + + // A steps down. Its release-on-return expires the lease immediately, so + // B's blocked AcquireTopic wins on its next retry. + cancelA() + select { + case err := <-doneA: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(15 * time.Second): + t.Fatal("A's Subscribe did not return after cancellation") + } + + addTestMessages(t, ctx, obB, "orders", 2) + require.Eventually(t, func() bool { + return len(flusherB.Received()) == 2 && countMessages(t, ctx, schema, "orders") == 0 + }, 15*time.Second, 50*time.Millisecond, "B should take over after A releases the lease") + assert.Len(t, flusherA.Received(), 4, "A must not process after stepping down") + + cancelB() + select { + case err := <-doneB: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(15 * time.Second): + t.Fatal("B's Subscribe did not return after cancellation") + } +} + +// TestOutbox_SubscribeExclusiveReacquiresLostLease verifies the self-heal +// path: when the lease is lost mid-subscribe, the next processing pass fails +// with ErrExclusiveLeaseRequired and Subscribe re-acquires on its own instead +// of spinning on errors. +func TestOutbox_SubscribeExclusiveReacquiresLostLease(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + schema := uniqueSchema(t) + ps := newTestPubSub(t, ctx, schema) + + outbox, err := pgoutbox.NewOutbox(ctx, sharedPool, pgoutbox.WithSchema(schema), pgoutbox.WithPubSub(ps)) + require.NoError(t, err) + + flusher := &captureFlusher{} + outbox.AddFlusher("orders", flusher) + + subCtx, subCancel := context.WithCancel(ctx) + defer subCancel() + done := make(chan error, 1) + go func() { + done <- outbox.Subscribe(subCtx, "orders", pgoutbox.WithExclusive(), pgoutbox.WithPollInterval(time.Minute)) + }() + + addTestMessages(t, ctx, outbox, "orders", 1) + require.Eventually(t, func() bool { + return len(flusher.Received()) == 1 + }, 15*time.Second, 50*time.Millisecond, "subscriber should drain while holding the lease") + + // Expire the lease behind the subscriber's back. The wake-up for the + // messages below then fails its pass and must trigger a re-acquire. + _, err = sharedPool.Exec(ctx, fmt.Sprintf( + "UPDATE %s.topics SET exclusive_consumer_expires_at = now() - interval '1 second' WHERE topic = $1", + pgx.Identifier{schema}.Sanitize()), "orders") + require.NoError(t, err) + + addTestMessages(t, ctx, outbox, "orders", 2) + require.Eventually(t, func() bool { + return len(flusher.Received()) == 3 && countMessages(t, ctx, schema, "orders") == 0 + }, 15*time.Second, 50*time.Millisecond, "subscriber should re-acquire the lost lease and keep draining") + + subCancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(15 * time.Second): + t.Fatal("Subscribe did not return after cancellation") + } +} + +func TestOutbox_SubscribeRequiresFlusher(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + schema := uniqueSchema(t) + + outbox, err := pgoutbox.NewOutbox(ctx, sharedPool, pgoutbox.WithSchema(schema)) + require.NoError(t, err) + + err = outbox.Subscribe(ctx, "unregistered") + require.Error(t, err) + assert.Contains(t, err.Error(), "unregistered") + assert.False(t, errors.Is(err, context.Canceled)) +} From 2f2f33e680189085b573d1040c96b59916be60d1 Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Fri, 17 Jul 2026 11:19:13 -0400 Subject: [PATCH 2/3] add lease write timeout --- outbox.go | 14 ++++++++++++++ subscribe.go | 9 ++++++--- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/outbox.go b/outbox.go index 468eab0..8e59b36 100644 --- a/outbox.go +++ b/outbox.go @@ -176,6 +176,15 @@ var exclusiveLeaseRenewInterval = newAtomicDuration(10 * time.Second) // to acquire the lease when another instance currently holds it. var exclusiveLeaseRetryInterval = newAtomicDuration(1 * time.Second) +// exclusiveLeaseWriteTimeout bounds the detached, best-effort lease writes: +// the immediate release when an exclusive Subscribe exits and the +// grace-period expiry written after a hold ends. These writes are detached +// from their caller's (usually already-cancelled) ctx, so without a bound a +// stalled database would block shutdown indefinitely; a write that cannot +// land within this window is better abandoned — the holder's consumer +// session lapsing covers the failure case. +const exclusiveLeaseWriteTimeout = 10 * time.Second + type outboxImplOpts struct { schema string autoMigrate bool @@ -576,7 +585,12 @@ func (o *outboxImpl) tryAcquireTopicLease(ctx context.Context, topic string) (bo // lapses one lease duration from now. Called by o.exclusiveLeaseWatcher with // an already-detached ctx once the AcquireTopic ctx has ended. The WHERE // clause makes this a silent no-op if another instance has since taken over. +// Self-bounded so that no caller — the watcher goroutine or the +// AcquireTopic/ReleaseTopic failure paths — can hang on a stalled database. func (o *outboxImpl) expireExclusiveLeaseAfterGrace(ctx context.Context, topic string) { + ctx, cancel := context.WithTimeout(ctx, exclusiveLeaseWriteTimeout) + defer cancel() + expiresAt := time.Now().UTC().Add(exclusiveLeaseDuration.Load()) err := o.queries.RenewTopicExclusiveConsumer(ctx, dbwrap.New(o.pool, o.schema), sqlc.RenewTopicExclusiveConsumerParams{ Topic: topic, diff --git a/subscribe.go b/subscribe.go index 081b988..941d272 100644 --- a/subscribe.go +++ b/subscribe.go @@ -76,11 +76,14 @@ func (o *outboxImpl) Subscribe(ctx context.Context, topic string, sopts ...Subsc return fmt.Errorf("could not acquire exclusive lease for topic %q: %w", topic, err) } // Hand the lease off immediately on exit instead of letting it lapse - // through the grace period. Detached ctx because the usual way out is - // this ctx being cancelled; a failed release degrades to the + // through the grace period. Detached from ctx (the usual way out is + // this ctx being cancelled) but bounded, so a stalled database can't + // hang subscriber shutdown; a missed release degrades to the // grace-period lapse inside ReleaseTopic itself. defer func() { - if err := o.ReleaseTopic(context.WithoutCancel(ctx), topic); err != nil { + releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), exclusiveLeaseWriteTimeout) + defer cancel() + if err := o.ReleaseTopic(releaseCtx, topic); err != nil { o.logger.Error().Err(err).Str("topic", topic).Msg("subscribe: failed to release exclusive lease") } }() From 51768ed25af3677e9286347ca9a2054522fa4bf8 Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Mon, 20 Jul 2026 11:43:13 -0700 Subject: [PATCH 3/3] update benchmarks with batching --- README.md | 36 ++++++++++++++++------ outbox_bench_test.go | 73 ++++++++++++++++++++++++++++---------------- 2 files changed, 73 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 65c59b1..1636dec 100644 --- a/README.md +++ b/README.md @@ -247,7 +247,7 @@ You can run benchmarks locally; for example, to write and flush 100k messages, y go test -bench=. -benchtime=100000x ``` -`BenchmarkOutbox_WriteAndPublishThroughput` drains each topic with a busy-polling `ProcessMessages` loop; `BenchmarkOutbox_SubscribeThroughput` drains each topic with a `Subscribe` call woken by `LISTEN`/`NOTIFY` (its poll interval is set far above the benchmark runtime, so throughput is carried entirely by notifications — and each producer commit pays the in-transaction `pg_notify`). Both run at 1 and 10 topics, with producers spreading messages round-robin and one consumer per topic. On a local Macbook with an M3 Max core: +`BenchmarkOutbox_WriteAndPublishThroughput` drains each topic with a busy-polling `ProcessMessages` loop; `BenchmarkOutbox_SubscribeThroughput` drains each topic with a `Subscribe` call woken by `LISTEN`/`NOTIFY` (its poll interval is set far above the benchmark runtime, so throughput is carried entirely by notifications — and each producer commit pays the in-transaction `pg_notify`). Both run a matrix of 1 and 10 topics (messages spread round-robin, one consumer per topic) and producer batch sizes of 1, 10, and 100 messages per `AddMessages` transaction. On a local Macbook with an M3 Max core: ``` $ go test -bench=. -benchtime=100000x @@ -255,12 +255,30 @@ goos: darwin goarch: arm64 pkg: github.com/hatchet-dev/pgoutbox cpu: Apple M3 Max -BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=1-14 100000 139432 ns/op 7172 msgs/sec -BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=1-14 100000 248696 ns/op 4021 msgs/sec -BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=10-14 100000 245562 ns/op 4072 msgs/sec -BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=10-14 100000 242325 ns/op 4127 msgs/sec -BenchmarkOutbox_SubscribeThroughput/Flush/topics=1-14 100000 188324 ns/op 5310 msgs/sec -BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=1-14 100000 305362 ns/op 3275 msgs/sec -BenchmarkOutbox_SubscribeThroughput/Flush/topics=10-14 100000 301837 ns/op 3313 msgs/sec -BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=10-14 100000 321406 ns/op 3111 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=1/batch=1-14 100000 125855 ns/op 7946 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=1/batch=1-14 100000 219546 ns/op 4555 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=1/batch=10-14 100000 16328 ns/op 61244 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=1/batch=10-14 100000 133446 ns/op 7494 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=1/batch=100-14 100000 6806 ns/op 146925 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=1/batch=100-14 100000 150929 ns/op 6626 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=10/batch=1-14 100000 293680 ns/op 3405 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=10/batch=1-14 100000 241538 ns/op 4140 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=10/batch=10-14 100000 25111 ns/op 39822 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=10/batch=10-14 100000 48755 ns/op 20511 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/Flush/topics=10/batch=100-14 100000 4468 ns/op 223817 msgs/sec +BenchmarkOutbox_WriteAndPublishThroughput/TxFlush/topics=10/batch=100-14 100000 36192 ns/op 27631 msgs/sec +BenchmarkOutbox_SubscribeThroughput/Flush/topics=1/batch=1-14 100000 209316 ns/op 4777 msgs/sec +BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=1/batch=1-14 100000 301901 ns/op 3312 msgs/sec +BenchmarkOutbox_SubscribeThroughput/Flush/topics=1/batch=10-14 100000 21391 ns/op 46749 msgs/sec +BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=1/batch=10-14 100000 153971 ns/op 6495 msgs/sec +BenchmarkOutbox_SubscribeThroughput/Flush/topics=1/batch=100-14 100000 7479 ns/op 133706 msgs/sec +BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=1/batch=100-14 100000 141189 ns/op 7083 msgs/sec +BenchmarkOutbox_SubscribeThroughput/Flush/topics=10/batch=1-14 100000 289379 ns/op 3456 msgs/sec +BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=10/batch=1-14 100000 924623 ns/op 1082 msgs/sec +BenchmarkOutbox_SubscribeThroughput/Flush/topics=10/batch=10-14 100000 33516 ns/op 29836 msgs/sec +BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=10/batch=10-14 100000 58475 ns/op 17101 msgs/sec +BenchmarkOutbox_SubscribeThroughput/Flush/topics=10/batch=100-14 100000 5352 ns/op 186840 msgs/sec +BenchmarkOutbox_SubscribeThroughput/TxFlush/topics=10/batch=100-14 100000 39021 ns/op 25627 msgs/sec ``` + +Batching producer writes is the single biggest lever: staging 100 messages per transaction reaches ~150-220k msgs/sec on the cheap-flusher path, roughly 20× the one-message-per-transaction rate. The `batch=1` cells run long enough to be sensitive to background database noise (checkpoints, autovacuum), so expect swings between runs — the 1082 msgs/sec outlier above measures ~3300 msgs/sec in isolation. diff --git a/outbox_bench_test.go b/outbox_bench_test.go index 9c9c7cc..5536436 100644 --- a/outbox_bench_test.go +++ b/outbox_bench_test.go @@ -53,6 +53,10 @@ func (c *txCountingFlusher) Flush(ctx pgoutbox.FlushContext, msgs []*sqlc.Messag // drains each topic. var benchTopicCounts = []int{1, 10} +// benchAddBatchSizes is the producer batch-size dimension: how many messages +// each AddMessages transaction stages. +var benchAddBatchSizes = []int{1, 10, 100} + func BenchmarkOutbox_WriteAndPublishThroughput(b *testing.B) { benchmarkThroughputMatrix(b, false) } @@ -67,22 +71,25 @@ func BenchmarkOutbox_SubscribeThroughput(b *testing.B) { benchmarkThroughputMatrix(b, true) } -// benchmarkThroughputMatrix runs the Flush/TxFlush × topic-count grid. +// benchmarkThroughputMatrix runs the Flush/TxFlush × topic-count × +// producer-batch-size grid. func benchmarkThroughputMatrix(b *testing.B, useSubscribe bool) { for _, numTopics := range benchTopicCounts { - b.Run(fmt.Sprintf("Flush/topics=%d", numTopics), func(b *testing.B) { - benchmarkThroughput(b, func(_ string, onFlush func(int)) pgoutbox.Flusher { - return &countingFlusher{onFlush: onFlush} - }, useSubscribe, numTopics) - }) - b.Run(fmt.Sprintf("TxFlush/topics=%d", numTopics), func(b *testing.B) { - benchmarkThroughput(b, func(schema string, onFlush func(int)) pgoutbox.Flusher { - table := pgx.Identifier{schema, "bench_side_log"}.Sanitize() - _, err := sharedPool.Exec(context.Background(), fmt.Sprintf("CREATE TABLE %s (msg_id bigint NOT NULL)", table)) - require.NoError(b, err) - return &txCountingFlusher{table: table, onFlush: onFlush} - }, useSubscribe, numTopics) - }) + for _, addBatch := range benchAddBatchSizes { + b.Run(fmt.Sprintf("Flush/topics=%d/batch=%d", numTopics, addBatch), func(b *testing.B) { + benchmarkThroughput(b, func(_ string, onFlush func(int)) pgoutbox.Flusher { + return &countingFlusher{onFlush: onFlush} + }, useSubscribe, numTopics, addBatch) + }) + b.Run(fmt.Sprintf("TxFlush/topics=%d/batch=%d", numTopics, addBatch), func(b *testing.B) { + benchmarkThroughput(b, func(schema string, onFlush func(int)) pgoutbox.Flusher { + table := pgx.Identifier{schema, "bench_side_log"}.Sanitize() + _, err := sharedPool.Exec(context.Background(), fmt.Sprintf("CREATE TABLE %s (msg_id bigint NOT NULL)", table)) + require.NoError(b, err) + return &txCountingFlusher{table: table, onFlush: onFlush} + }, useSubscribe, numTopics, addBatch) + }) + } } } @@ -91,9 +98,10 @@ func benchmarkThroughputMatrix(b *testing.B, useSubscribe bool) { // inFlight-releasing callback, letting us reuse the harness for both the plain // Flush path and the FlushWithTx path. With useSubscribe each topic's consumer // is a Subscribe call woken by pg_notify (and AddMessages pays the in-tx -// publish); otherwise it is a busy-polling ProcessMessages loop. Producers -// assign messages to the numTopics topics round-robin. -func benchmarkThroughput(b *testing.B, newFlusher func(schema string, onFlush func(n int)) pgoutbox.Flusher, useSubscribe bool, numTopics int) { +// publish); otherwise it is a busy-polling ProcessMessages loop. Each producer +// transaction stages addBatchSize messages via one AddMessages call, with +// batches assigned to the numTopics topics round-robin. +func benchmarkThroughput(b *testing.B, newFlusher func(schema string, onFlush func(n int)) pgoutbox.Flusher, useSubscribe bool, numTopics, addBatchSize int) { const ( numWorkers = MAX_CONNS maxInFlight = 5000 @@ -176,23 +184,32 @@ func benchmarkThroughput(b *testing.B, newFlusher func(schema string, onFlush fu b.ResetTimer() - work := make(chan struct{}) - var msgIdx atomic.Int64 + // Each work item is one producer transaction staging that many messages + // (addBatchSize, or less for the final remainder batch) on one topic. + work := make(chan int) + var batchIdx atomic.Int64 var wg sync.WaitGroup for range numWorkers { wg.Go(func() { - for range work { - inFlight <- struct{}{} - topic := topics[msgIdx.Add(1)%int64(numTopics)] + batch := make([]pgoutbox.MessageOpts, addBatchSize) + for i := range batch { + batch[i] = pgoutbox.MessageOpts{Payload: payload} + } + for n := range work { + // Workers hold at most numWorkers*addBatchSize slots while + // blocked here, well under maxInFlight, so partial + // acquisition cannot deadlock. + for range n { + inFlight <- struct{}{} + } + topic := topics[batchIdx.Add(1)%int64(numTopics)] tx, err := sharedPool.Begin(ctx) if err != nil { b.Errorf("begin: %v", err) return } - if err := outbox.AddMessages(ctx, tx, topic, []pgoutbox.MessageOpts{ - {Payload: payload}, - }); err != nil { + if err := outbox.AddMessages(ctx, tx, topic, batch[:n]); err != nil { _ = tx.Rollback(ctx) b.Errorf("AddMessages: %v", err) return @@ -205,8 +222,10 @@ func benchmarkThroughput(b *testing.B, newFlusher func(schema string, onFlush fu }) } - for range b.N { - work <- struct{}{} + for remaining := b.N; remaining > 0; { + n := min(addBatchSize, remaining) + work <- n + remaining -= n } close(work) wg.Wait()