├── .gitignore ├── go.mod ├── internal ├── test │ └── test.go └── util.go ├── go.sum ├── config.go ├── .github └── workflows │ └── ci.yml ├── Makefile ├── topology_int_test.go ├── hang_int_test.go ├── publisher_int_test.go ├── README.md ├── healthcheck_int_test.go ├── topology.go ├── benchmark_int_test.go ├── conn_test.go ├── consumer_int_test.go ├── connection.go ├── consumer.go ├── LICENSE └── publisher.go /.gitignore: -------------------------------------------------------------------------------- 1 | .env 2 | bin/* 3 | .coverage* -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/danlock/rmq 2 | 3 | go 1.21 4 | 5 | require github.com/rabbitmq/amqp091-go v1.10.0 6 | -------------------------------------------------------------------------------- /internal/test/test.go: -------------------------------------------------------------------------------- 1 | package test 2 | 3 | import "testing" 4 | 5 | func FailOnError(t testing.TB, err error) { 6 | if err != nil { 7 | t.Helper() 8 | t.Fatalf("%+v", err) 9 | } 10 | } 11 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/rabbitmq/amqp091-go v1.10.0 h1:STpn5XsHlHGcecLmMFCtg7mqq0RnD+zFr4uzukfVhBw= 2 | github.com/rabbitmq/amqp091-go v1.10.0/go.mod h1:Hy4jKW5kQART1u+JkDTF9YYOQUHXqMuhrgxOEeS7G4o= 3 | go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= 4 | go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= 5 | -------------------------------------------------------------------------------- /config.go: -------------------------------------------------------------------------------- 1 | package rmq 2 | 3 | import ( 4 | "context" 5 | "log/slog" 6 | "time" 7 | 8 | "github.com/danlock/rmq/internal" 9 | ) 10 | 11 | // Args contains common options shared by danlock/rmq classes. 12 | type Args struct { 13 | // AMQPTimeout sets a timeout on AMQP operations. Defaults to 1 minute. 14 | AMQPTimeout time.Duration 15 | // Delay returns the delay between retry attempts. Defaults to FibonacciDelay. 16 | Delay func(attempt int) time.Duration 17 | // Log can be left nil, set with slog.Log or wrapped around your favorite logging library 18 | Log func(ctx context.Context, level slog.Level, msg string, args ...any) 19 | } 20 | 21 | func (cfg *Args) setDefaults() { 22 | if cfg.AMQPTimeout == 0 { 23 | cfg.AMQPTimeout = time.Minute 24 | } 25 | 26 | if cfg.Delay == nil { 27 | cfg.Delay = FibonacciDelay 28 | } 29 | 30 | internal.WrapLogFunc(&cfg.Log) 31 | } 32 | 33 | func FibonacciDelay(attempt int) time.Duration { 34 | if attempt < len(internal.FibonacciDurations) { 35 | return internal.FibonacciDurations[attempt] 36 | } else { 37 | return internal.FibonacciDurations[len(internal.FibonacciDurations)-1] 38 | } 39 | } 40 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: danlock/rmq CI 2 | 3 | on: [push] 4 | 5 | jobs: 6 | test: 7 | runs-on: ubuntu-latest 8 | services: 9 | rabbitmq: 10 | image: rabbitmq:3.10 11 | env: 12 | RABBITMQ_DEFAULT_USER: guest 13 | RABBITMQ_DEFAULT_PASS: guest 14 | ports: 15 | - 5672 16 | 17 | steps: 18 | - uses: actions/checkout@v3 19 | with: 20 | # default fetch-depth is insufficent to find previous coverage notes for gwatts/go-coverage-action@v1 21 | fetch-depth: 10 22 | 23 | - name: Set up Go 24 | uses: actions/setup-go@v4 25 | with: 26 | go-version: '1.21' 27 | 28 | - name: Dependencies 29 | run: make deps 30 | 31 | - name: Unit test 32 | run: make unit-test 33 | 34 | - name: Test 35 | run: TEST_AMQP_URI=amqp://guest:guest@127.0.0.1:${{ job.services.rabbitmq.ports['5672'] }}/ make test 36 | 37 | - uses: gwatts/go-coverage-action@v1 38 | id: coverage 39 | env: 40 | TEST_AMQP_URI: amqp://guest:guest@127.0.0.1:${{ job.services.rabbitmq.ports['5672'] }}/ 41 | with: 42 | coverage-threshold: 80 43 | test-args: '["-tags=rabbit"]' 44 | 45 | # - name: update coverage badge 46 | # env: 47 | # COVERAGE_PATH: ${{ steps.coverage.outputs.gocov-agg-pathname }} 48 | # run: make update-readme-badge 49 | 50 | # - uses: stefanzweifel/git-auto-commit-action@v4 51 | # with: 52 | # file_pattern: 'README.md' 53 | # commit_message: "Update coverage badge" 54 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | #! /usr/bin/make 2 | SHELL = /bin/bash 3 | BUILDTIME = $(shell date -u --rfc-3339=seconds) 4 | GITHASH = $(shell git describe --dirty --always) 5 | GITCOMMITNO = $(shell git rev-list --all --count) 6 | SHORTBUILDTAG = v0.0.$(GITCOMMITNO)-$(GITHASH) 7 | BUILDINFO = Build Time:$(BUILDTIME) 8 | LDFLAGS = -X 'main.buildTag=$(SHORTBUILDTAG)' -X 'main.buildInfo=$(BUILDINFO)' 9 | 10 | TEST_AMQP_URI ?= amqp://guest:guest@localhost:5672/ 11 | COVERAGE_PATH ?= .coverage 12 | 13 | 14 | depend: deps 15 | deps: 16 | go get ./... 17 | go mod tidy 18 | 19 | version: 20 | @echo $(SHORTBUILDTAG) 21 | 22 | unit-test: 23 | @go test -failfast -race -count=3 ./... 24 | 25 | test: 26 | @TEST_AMQP_URI=$(TEST_AMQP_URI) go test -failfast -v -race -count=2 -tags="rabbit" ./... 27 | 28 | bench: 29 | @TEST_AMQP_URI=$(TEST_AMQP_URI) go test -failfast -benchmem -run=^$ -v -count=2 -tags="rabbit" -bench . ./... 30 | 31 | coverage: 32 | @TEST_AMQP_URI=$(TEST_AMQP_URI) go test -failfast -covermode=count -tags="rabbit" -coverprofile=$(COVERAGE_PATH) 33 | 34 | coverage-html: 35 | @rm $(COVERAGE_PATH) || true 36 | @$(MAKE) coverage 37 | @rm $(COVERAGE_PATH).html || true 38 | @go tool cover -html=$(COVERAGE_PATH) -o $(COVERAGE_PATH).html 39 | 40 | coverage-browser: 41 | @rm $(COVERAGE_PATH) || true 42 | @$(MAKE) coverage 43 | @go tool cover -html=$(COVERAGE_PATH) 44 | 45 | update-readme-badge: 46 | @go tool cover -func=$(COVERAGE_PATH) -o=$(COVERAGE_PATH).badge 47 | @go run github.com/AlexBeauchemin/gobadge@v0.3.0 -filename=$(COVERAGE_PATH).badge 48 | 49 | # pkg.go.dev documentation is updated via go get updating the google proxy 50 | update-godocs: 51 | @cd ../regex-img; \ 52 | GOPROXY=https://proxy.golang.org go get -u github.com/danlock/rmq; \ 53 | go mod tidy 54 | 55 | release: 56 | @$(MAKE) deps 57 | ifeq ($(findstring dirty,$(SHORTBUILDTAG)),dirty) 58 | @echo "Version $(SHORTBUILDTAG) is filthy, commit to clean it" && exit 1 59 | endif 60 | @read -t 5 -p "$(SHORTBUILDTAG) will be the new released version. Hit enter to proceed, CTRL-C to cancel." 61 | @$(MAKE) test 62 | @$(MAKE) bench 63 | @git tag $(SHORTBUILDTAG) 64 | @git push origin $(SHORTBUILDTAG) 65 | -------------------------------------------------------------------------------- /topology_int_test.go: -------------------------------------------------------------------------------- 1 | //go:build rabbit 2 | 3 | package rmq_test 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "log/slog" 9 | "os" 10 | "testing" 11 | "time" 12 | 13 | "github.com/danlock/rmq" 14 | amqp "github.com/rabbitmq/amqp091-go" 15 | ) 16 | 17 | func TestDeclareTopology(t *testing.T) { 18 | ctx := context.Background() 19 | baseCfg := rmq.Args{Log: slog.Log} 20 | rmqConn := rmq.ConnectWithAMQPConfig(ctx, rmq.ConnectArgs{Args: baseCfg}, os.Getenv("TEST_AMQP_URI"), amqp.Config{}) 21 | suffix := fmt.Sprintf("%s|%p", time.Now(), t) 22 | baseTopology := rmq.Topology{ 23 | Exchanges: []rmq.Exchange{{ 24 | Name: "temporary", 25 | Kind: amqp.ExchangeTopic, 26 | AutoDelete: true, 27 | }, { 28 | Name: "ephemeral", 29 | Kind: amqp.ExchangeTopic, 30 | AutoDelete: true, 31 | }, { 32 | Name: "ephemeral", 33 | Kind: amqp.ExchangeTopic, 34 | AutoDelete: true, 35 | Passive: true, 36 | }}, 37 | ExchangeBindings: []rmq.ExchangeBinding{{ 38 | Destination: "temporary", 39 | RoutingKey: "dopeopleevenuseexchangebindings", 40 | Source: "ephemeral", 41 | }}, 42 | Queues: []rmq.Queue{{ 43 | // DeclareTopology skips queues without Names. Those are for Consumers instead 44 | }, { 45 | Name: "transient" + suffix, 46 | Durable: true, 47 | Exclusive: true, 48 | Args: amqp.Table{amqp.QueueTTLArg: time.Minute.Milliseconds()}, 49 | }, { 50 | Name: "transient" + suffix, 51 | Durable: true, 52 | Exclusive: true, 53 | Passive: true, 54 | Args: amqp.Table{amqp.QueueTTLArg: time.Minute.Milliseconds()}, 55 | }}, 56 | QueueBindings: []rmq.QueueBinding{{ 57 | QueueName: "transient" + suffix, 58 | ExchangeName: "temporary", 59 | RoutingKey: "route66", 60 | }}, 61 | } 62 | 63 | amqpConn, err := rmqConn.CurrentConnection(ctx) 64 | if err != nil { 65 | t.Fatalf("failed to CurrentConnection %v", err) 66 | } 67 | 68 | tests := []struct { 69 | name string 70 | timeout time.Duration 71 | topology rmq.Topology 72 | wantErr bool 73 | }{ 74 | { 75 | "success", 76 | time.Minute, 77 | baseTopology, 78 | false, 79 | }, 80 | { 81 | "empty success", 82 | time.Minute, 83 | rmq.Topology{}, 84 | false, 85 | }, 86 | { 87 | "failure due to timeout", 88 | time.Millisecond, 89 | baseTopology, 90 | true, 91 | }, 92 | } 93 | for _, tt := range tests { 94 | t.Run(tt.name, func(t *testing.T) { 95 | ctx, cancel := context.WithTimeout(ctx, tt.timeout) 96 | defer cancel() 97 | tt.topology.Log = slog.Log 98 | if err := rmq.DeclareTopology(ctx, amqpConn, tt.topology); (err != nil) != tt.wantErr { 99 | t.Errorf("DeclareTopology() error = %v, wantErr %v", err, tt.wantErr) 100 | } 101 | }) 102 | } 103 | } 104 | -------------------------------------------------------------------------------- /internal/util.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "time" 7 | 8 | "log/slog" 9 | ) 10 | 11 | // ChanReq and ChanResp are used to send and receive resources over a channel. 12 | // The Ctx is sent so that the listener can use it for a timeout if necessary. 13 | // RespChan should be buffered to at least 1 to not block the listening goroutine. 14 | type ChanReq[T any] struct { 15 | Ctx context.Context 16 | RespChan chan ChanResp[T] 17 | } 18 | type ChanResp[T any] struct { 19 | Val T 20 | Err error 21 | } 22 | 23 | var FibonacciDurations = [...]time.Duration{ 24 | 0, time.Second, time.Second, 2 * time.Second, 3 * time.Second, 5 * time.Second, 25 | 8 * time.Second, 13 * time.Second, 21 * time.Second, 34 * time.Second, 26 | } 27 | 28 | // slog.Log's function signature. Useful for context aware logging and simpler to wrap than an interface like slog.Handler 29 | type SlogLog = func(context.Context, slog.Level, string, ...any) 30 | 31 | // WrapLogFunc runs fmt.Sprintf on the msg, args parameters so the end user can use slog.Log or any other logging library more interchangeably. 32 | // The slog.Log func signature is an improvement over the usual func(string, any...). 33 | // The end user can take advantage of context for log tracing, slog.Level to ignore warnings, and we only depend on the stdlib. 34 | // This does mean calldepth loggers will need a +1 however. 35 | func WrapLogFunc(logFunc *SlogLog) { 36 | if logFunc == nil { 37 | panic("WrapLogFunc called with nil") 38 | } else if *logFunc == nil { 39 | *logFunc = func(context.Context, slog.Level, string, ...any) {} 40 | } else { 41 | userLog := *logFunc 42 | *logFunc = func(ctx context.Context, level slog.Level, msg string, args ...any) { 43 | userLog(ctx, level, fmt.Sprintf(msg, args...)) 44 | } 45 | } 46 | } 47 | 48 | // AMQP091Logger wraps the amqp091 Logger interface with a little boilerplate. 49 | type AMQP091Logger struct { 50 | Ctx context.Context 51 | Log SlogLog 52 | } 53 | 54 | func (l AMQP091Logger) Printf(format string, v ...interface{}) { 55 | l.Log(l.Ctx, slog.LevelError, "rabbitmq/amqp091-go: "+fmt.Sprintf(format, v...)) 56 | } 57 | 58 | const healthyLifetime = 20 * time.Millisecond 59 | 60 | // Retry attempts do repeatedly until it's ctx ends. If do returns false delayForAttempt is used to backoff retries. 61 | // If do is true and ran longer than healthyLifetime or it's last delay the backoff is reset. do returns it's own lifetime, since it may do some setup beforehand. 62 | func Retry(ctx context.Context, delayForAttempt func(int) time.Duration, do func(time.Duration) (time.Duration, bool)) { 63 | var delay time.Duration 64 | var attempt = 0 65 | for { 66 | delay = delayForAttempt(attempt) 67 | attempt++ 68 | lifetime, ok := do(delay) 69 | if ok && lifetime >= max(healthyLifetime, delay) { 70 | delay, attempt = 0, 0 71 | } 72 | 73 | select { 74 | case <-ctx.Done(): 75 | return 76 | case <-time.After(delay): 77 | } 78 | } 79 | } 80 | -------------------------------------------------------------------------------- /hang_int_test.go: -------------------------------------------------------------------------------- 1 | //go:build rabbit 2 | 3 | package rmq_test 4 | 5 | import ( 6 | "context" 7 | "errors" 8 | "log" 9 | "log/slog" 10 | "net" 11 | "os" 12 | "testing" 13 | "time" 14 | 15 | "github.com/danlock/rmq" 16 | amqp "github.com/rabbitmq/amqp091-go" 17 | ) 18 | 19 | func TestHanging(t *testing.T) { 20 | Example_hanging() 21 | } 22 | 23 | func panicOnErr(err error) { 24 | if err != nil { 25 | panic(err) 26 | } 27 | } 28 | 29 | func Example_hanging() { 30 | var tcpConn *net.TCPConn 31 | amqp091Config := amqp.Config{ 32 | // Set dial so we have access to the net.Conn 33 | // This is the same as amqp091.DefaultDial(time.Second) except we also grab the connection 34 | Dial: func(network, addr string) (net.Conn, error) { 35 | conn, err := net.DialTimeout(network, addr, time.Second) 36 | if err != nil { 37 | return nil, err 38 | } 39 | if err := conn.SetDeadline(time.Now().Add(time.Second)); err != nil { 40 | return nil, err 41 | } 42 | tcpConn = conn.(*net.TCPConn) 43 | return conn, nil 44 | }, 45 | } 46 | // Create an innocent, unsuspecting amqp091 connection 47 | amqp091Conn, err := amqp.DialConfig(os.Getenv("TEST_AMQP_URI"), amqp091Config) 48 | panicOnErr(err) 49 | // Create a channel to ensure the connection's working. 50 | amqp091Chan, err := amqp091Conn.Channel() 51 | panicOnErr(err) 52 | panicOnErr(amqp091Chan.Close()) 53 | // Betray amqp091Conn expectations by dawdling. While this is unnatural API usage, the intention is to emulate a connection hang. 54 | dawdlingBegins := make(chan struct{}, 1) 55 | // hangTime is 3 seconds for faster tests, but this could easily be much longer... 56 | hangTime := 3 * time.Second 57 | hangConnection := func(tcpConn *net.TCPConn) { 58 | go func() { 59 | sysConn, err := tcpConn.SyscallConn() 60 | panicOnErr(err) 61 | // sysConn.Write blocks the whole connection until it finishes 62 | err = sysConn.Write(func(fd uintptr) bool { 63 | dawdlingBegins <- struct{}{} 64 | time.Sleep(hangTime) 65 | return true 66 | }) 67 | panicOnErr(err) 68 | }() 69 | select { 70 | case <-time.After(time.Second): 71 | panic("sysConn.Write took too long!") 72 | case <-dawdlingBegins: 73 | } 74 | } 75 | hangConnection(tcpConn) 76 | // The unsuspecting amqp091Conn.Channel() dutifully waits for hangTime. 77 | // Doesn't matter what amqp.DefaultDial(connectionTimeout) was (only 1 second...) 78 | chanStart := time.Now() 79 | amqp091Chan, err = amqp091Conn.Channel() 80 | panicOnErr(err) 81 | panicOnErr(amqp091Chan.Close()) 82 | panicOnErr(amqp091Conn.Close()) 83 | // test our expectation that amqp091Conn.Channel hung for at least 90% of hangTime, to prevent flaky tests. 84 | if time.Since(chanStart) < (hangTime - (hangTime / 10)) { 85 | panic("amqp091Conn.Channel returned faster than expected") 86 | } 87 | // The above demonstrates one of the biggest issues with amqp091, since your applications stuck if the connection hangs, 88 | // and you don't have any options to prevent this. 89 | ctx := context.Background() 90 | // danlock/rmq gives you 2 ways to prevent unbound hangs, Args.AMQPTimeout and the context passed into each function call. 91 | rmqConnCfg := rmq.ConnectArgs{Args: rmq.Args{Log: slog.Log, AMQPTimeout: time.Second}} 92 | // Create a paranoid AMQP connection 93 | rmqConn := rmq.ConnectWithAMQPConfig(ctx, rmqConnCfg, os.Getenv("TEST_AMQP_URI"), amqp091Config) 94 | // Grab a channel to ensure the connection is working 95 | amqp091Chan, err = rmqConn.Channel(ctx) 96 | panicOnErr(err) 97 | panicOnErr(amqp091Chan.Close()) 98 | // a hung connection, just like we've always feared 99 | hangConnection(tcpConn) 100 | // However we will simply error long before hangTime. 101 | chanStart = time.Now() 102 | _, err = rmqConn.Channel(ctx) 103 | if !errors.Is(err, context.DeadlineExceeded) { 104 | log.Fatalf("rmqConn.Channel returned unexpected error %v", err) 105 | } 106 | chanDur := time.Since(chanStart) 107 | // rmqConn is too paranoid to hang for 90% of hangTime, but double check anyway 108 | if chanDur > (hangTime - (hangTime / 10)) { 109 | log.Fatalf("rmqConn.Channel hung for (%s)", chanDur) 110 | } 111 | // A caveat here is that rmqConn has leaked a goroutine that blocks until the connection sorts itself out. 112 | // If amqp091-go ever fixes https://github.com/rabbitmq/amqp091-go/issues/225 then we can improve this situation. 113 | } 114 | -------------------------------------------------------------------------------- /publisher_int_test.go: -------------------------------------------------------------------------------- 1 | //go:build rabbit 2 | 3 | package rmq_test 4 | 5 | import ( 6 | "context" 7 | "log/slog" 8 | "os" 9 | "reflect" 10 | "testing" 11 | "time" 12 | 13 | "github.com/danlock/rmq" 14 | amqp "github.com/rabbitmq/amqp091-go" 15 | ) 16 | 17 | func TestPublisher(t *testing.T) { 18 | ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) 19 | defer cancel() 20 | 21 | baseCfg := rmq.Args{Log: slog.Log} 22 | rmqConn := rmq.ConnectWithAMQPConfig(ctx, rmq.ConnectArgs{Args: baseCfg}, os.Getenv("TEST_AMQP_URI"), amqp.Config{}) 23 | 24 | unreliableRMQPub := rmq.NewPublisher(ctx, rmqConn, rmq.PublisherArgs{DontConfirm: true}) 25 | _, err := unreliableRMQPub.PublishUntilConfirmed(ctx, time.Minute, rmq.Publishing{}) 26 | if err == nil { 27 | t.Fatalf("PublishUntilConfirmed succeeded despite the publisher set dont confirm") 28 | } 29 | 30 | returnChan := make(chan amqp.Return, 5) 31 | rmqPub := rmq.NewPublisher(ctx, rmqConn, rmq.PublisherArgs{ 32 | Args: baseCfg, 33 | NotifyReturn: returnChan, 34 | LogReturns: true, 35 | }) 36 | ForceRedial(ctx, rmqConn) 37 | pubCtx, pubCancel := context.WithTimeout(ctx, 10*time.Second) 38 | defer pubCancel() 39 | 40 | wantedPub := rmq.Publishing{Exchange: "amq.topic"} 41 | wantedPub.Body = []byte("Testrmq.Publisher") 42 | _, err = rmqPub.PublishUntilConfirmed(pubCtx, time.Minute, wantedPub) 43 | if err != nil { 44 | t.Fatalf("PublishUntilConfirmed failed with %v", err) 45 | } 46 | // Publish a message to get it returned 47 | retPub := rmq.Publishing{Exchange: "amq.topic", RoutingKey: "nowhere", Mandatory: true} 48 | retPub.Body = []byte("return me") 49 | 50 | err = rmqPub.PublishUntilAcked(pubCtx, 0, retPub) 51 | if err != nil { 52 | t.Fatalf("PublishUntilAcked failed with %v", err) 53 | } 54 | ForceRedial(ctx, rmqConn) 55 | select { 56 | case <-pubCtx.Done(): 57 | t.Fatalf("didnt get return") 58 | case ret := <-returnChan: 59 | if !reflect.DeepEqual(retPub.Body, ret.Body) { 60 | t.Fatalf("got different return message") 61 | } 62 | } 63 | 64 | // Publish a few messages at the same time. 65 | pubCount := 10 66 | errChan := make(chan error, pubCount) 67 | for i := 0; i < pubCount; i++ { 68 | go func() { 69 | errChan <- rmqPub.PublishUntilAcked(pubCtx, 0, wantedPub) 70 | }() 71 | } 72 | for i := 0; i < pubCount; i++ { 73 | if err := <-errChan; err != nil { 74 | t.Fatalf("PublishUntilAcked returned unexpected error %v", err) 75 | } 76 | } 77 | // Publishing mandatory and immediate messages to nonexistent queues should get confirmed, just not acked. 78 | mandatoryPub := rmq.Publishing{Exchange: "Idontexist", Mandatory: true} 79 | immediatePub := rmq.Publishing{Exchange: "Idontexist", Immediate: true} 80 | mandatoryPub.Body = wantedPub.Body 81 | immediatePub.Body = wantedPub.Body 82 | 83 | defConf, err := rmqPub.PublishUntilConfirmed(pubCtx, 0, mandatoryPub) 84 | if err != nil { 85 | t.Fatalf("PublishUntilConfirmed returned unexpected error %v", err) 86 | } 87 | if defConf.Acked() { 88 | t.Fatalf("PublishUntilConfirmed returned unexpected ack for mandatory pub") 89 | } 90 | 91 | defConf, err = rmqPub.PublishUntilConfirmed(pubCtx, 0, immediatePub) 92 | if err != nil { 93 | t.Fatalf("PublishUntilConfirmed returned unexpected error %v", err) 94 | } 95 | if defConf.Acked() { 96 | t.Fatalf("PublishUntilConfirmed returned unexpected ack for immediate pub") 97 | } 98 | 99 | returnedPub := rmq.Publishing{Exchange: "amq.topic", RoutingKey: "whereverhueyislooking", Mandatory: true} 100 | returnedPub.Body = []byte("oops") 101 | 102 | err = rmqPub.PublishUntilAcked(pubCtx, 0, returnedPub) 103 | if err != nil { 104 | t.Fatalf("PublishUntilAcked got err for return %v", err) 105 | } 106 | 107 | select { 108 | case <-pubCtx.Done(): 109 | t.Fatalf("didnt get return") 110 | case ret := <-returnChan: 111 | if !reflect.DeepEqual(returnedPub.Body, ret.Body) { 112 | t.Fatalf("got different return message") 113 | } 114 | } 115 | 116 | err = rmqPub.PublishUntilAcked(pubCtx, 0, wantedPub) 117 | if err != nil { 118 | t.Fatalf("PublishUntilAcked returned unexpected error %v", err) 119 | } 120 | 121 | // Cancel everything, now the publisher stopped processing 122 | cancel() 123 | _, err = rmqPub.Publish(context.Background(), wantedPub) 124 | if err == nil { 125 | t.Fatalf("publish shouldn't succeed") 126 | } 127 | _, err = rmqPub.PublishUntilConfirmed(ctx, 0, wantedPub) 128 | if err == nil { 129 | t.Fatalf("PublishUntilConfirmed shouldn't succeed") 130 | } 131 | 132 | err = rmqPub.PublishUntilAcked(ctx, 0, wantedPub) 133 | if err == nil { 134 | t.Fatalf("PublishUntilAcked shouldn't succeed") 135 | } 136 | } 137 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # rmq 2 | ![Coverage](https://img.shields.io/badge/Coverage-85.6%25-brightgreen) 3 | [![Go Report Card](https://goreportcard.com/badge/github.com/danlock/rmq)](https://goreportcard.com/report/github.com/danlock/rmq) 4 | [![Go Reference](https://pkg.go.dev/badge/github.com/danlock/rmq.svg)](https://pkg.go.dev/github.com/danlock/rmq) 5 | 6 | An AMQP library for Go, built on top of amqp091. 7 | 8 | [streadway/amqp](https://github.com/streadway/amqp), the library the RabbitMQ maintainers forked to [amqp-091](https://github.com/rabbitmq/amqp091-go), is a stable, thin client for communicating to RabbitMQ, but lacks many of the features present in RabbitMQ libraries from other languages. Many redialable AMQP connections have been reinvented in Go codebases everywhere. 9 | 10 | This package attempts to provide a wrapper of useful features on top of amqp091, in the hopes of preventing at least one more unnecessary reinvention (other than itself!) 11 | 12 | # Design Goals 13 | 14 | - Minimal API that doesn't get in the way of lower level access. The amqp091.Connection is there if you need it. amqp-091 knowledge is more transferable since danlock/rmq builds on top of those concepts rather than encapsulating things it doesn't need to. 15 | 16 | - Network aware message delivery. Networks fail so danlock/rmq uses context.Context and default timeouts wherever possible, and tries to redeliver across network failures, unlike amqp091-go. 17 | 18 | - One dependency (rabbitmq/amqp091-go). 19 | 20 | - Prioritize readability. This means no functions with 5 boolean args. 21 | 22 | # Examples 23 | 24 | Using an AMQP publisher to publish a message with at least once delivery, that retries for up to a minute on failures. 25 | 26 | ``` 27 | ctx, cancel := context.WithTimeout(context.TODO(), time.Minute) 28 | defer cancel() 29 | cfg := rmq.Args{Log: slog.Log} 30 | 31 | rmqConn := rmq.ConnectWithURLs(ctx, rmq.ConnectArgs{Args: cfg}, os.Getenv("AMQP_URL_1"), os.Getenv("AMQP_URL_2")) 32 | 33 | rmqPub := rmq.NewPublisher(ctx, rmqConn, rmq.PublisherArgs{Args: cfg}) 34 | 35 | msg := rmq.Publishing{Exchange: "amq.topic", RoutingKey: "somewhere", Mandatory: true} 36 | msg.Body = []byte(`{"life": 42}`) 37 | 38 | if err := rmqPub.PublishUntilAcked(ctx, time.Minute, msg); err != nil { 39 | return fmt.Errorf("PublishUntilAcked timed out because %w", err) 40 | } 41 | ``` 42 | 43 | Using a reliable AMQP consumer that receives deliveries through transient network failures while processing work concurrently with bounded goroutines. 44 | 45 | ``` 46 | ctx, := context.TODO() 47 | cfg := rmq.Args{Log: slog.Log} 48 | 49 | rmqConn := rmq.ConnectWithAMQPConfig(ctx, rmq.ConnectArgs{Args: cfg}, os.Getenv("AMQP_URL"), amqp.Config{}) 50 | 51 | consCfg := rmq.ConsumerArgs{ 52 | Args: cfg, 53 | Queue: rmq.Queue{Name: "q2d2", AutoDelete: true}, 54 | Qos: rmq.Qos{PrefetchCount: 1000}, 55 | } 56 | 57 | rmq.NewConsumer(rmqConn, consCfg).ConsumeConcurrently(ctx, 100, func(ctx context.Context, msg amqp.Delivery) { 58 | process(msg) 59 | handleAckErr(msg.Ack(false)) 60 | }) 61 | ``` 62 | 63 | Creating an AMQP topology that is automatically applied on reconnections as seen in the Java and C# RabbitMQ client drivers. 64 | 65 | ``` 66 | ctx, := context.TODO() 67 | cfg := rmq.Args{Log: slog.Log} 68 | 69 | topology := rmq.Topology{ 70 | Args: cfg, 71 | Exchanges: []rmq.Exchange{{Name: "xchg", Kind: amqp.ExchangeDirect, AutoDelete: true}}, 72 | Queues: []rmq.Queue{{Name: "huehue", Durable: true, AutoDelete: true}}, 73 | QueueBindings: []rmq.QueueBinding{{QueueName: "huehue", ExchangeName: "xchg"}}, 74 | } 75 | 76 | // It may be desired to read your AMQP topology from disk as JSON or some other config format. rmq.Topology is a simple struct so it can be done like so. 77 | // err := json.NewDecoder(topologyFile).Decode(&topology) 78 | // topology.Args = cfg 79 | 80 | rmqConn := rmq.ConnectWithURLs(ctx, rmq.ConnectArgs{Args: cfg, Topology: topology}, os.Getenv("AMQP_URL")) 81 | ``` 82 | 83 | Take a look at healthcheck_int_test.go for a more complete example of using all of danlock/rmq together, or hang_int_test.go for an example of danlock/rmq being more network-aware than amqp091-go. 84 | 85 | # Logging 86 | 87 | danlock/rmq sometimes handles errors by retrying instead of returning. In situations like this, danlock/rmq logs if you allow it to for easier debugging. 88 | 89 | All classes accept a Log function pointer that can be ignored entirely, set easily with slog.Log, or wrapped around your favorite logging library. 90 | 91 | Here is an example logrus wrapper. danlock/rmq only uses the predefined slog.Level's, and doesn't send any args. 92 | ``` 93 | Args{ 94 | Log: func(ctx context.Context, level slog.Level, msg string, _ ...any) { 95 | logruslevel, _ := logrus.ParseLevel(level.String()) 96 | logrus.StandardLogger().WithContext(ctx).Logf(logruslevel, msg) 97 | } 98 | } 99 | ``` -------------------------------------------------------------------------------- /healthcheck_int_test.go: -------------------------------------------------------------------------------- 1 | //go:build rabbit 2 | 3 | package rmq_test 4 | 5 | import ( 6 | "context" 7 | "errors" 8 | "log" 9 | "log/slog" 10 | "os" 11 | "reflect" 12 | "testing" 13 | "time" 14 | 15 | "github.com/danlock/rmq" 16 | amqp "github.com/rabbitmq/amqp091-go" 17 | ) 18 | 19 | func TestHealthcheck(t *testing.T) { 20 | Example_healthcheck() 21 | } 22 | 23 | // Example shows how to write an unsophisticated healthcheck for a service intending to ensure it's rmq.Connection is capable of processing messages. 24 | // Even though rmq.Connection reconnects on errors, there can always be unforeseen networking/DNS/RNGesus issues 25 | // that necessitate a docker/kubernetes healthcheck restarting the service when unhealthy. 26 | // While this is useful as an example, it wouldn't be used on production for several reasons, only one of which is the lack of reuse of AMQP connections and AMQP channels. 27 | func Example_healthcheck() { 28 | // Real applications should use a real context. If this healthcheck was called via HTTP request for example, 29 | // that HTTP request's context would be a good candidate. 30 | ctx, cancel := context.WithTimeout(context.TODO(), time.Minute) 31 | defer cancel() 32 | // If we want to use a different log library instead of log/slog.Log, wrap the function instead. 33 | // If call depth is being logged, add to it so it doesn't just print this log function. 34 | // Here we use log instead of slog 35 | customLog := func(ctx context.Context, level slog.Level, msg string, _ ...any) { 36 | log.Printf("[%s] trace_id=%v msg="+msg, level, ctx.Value("your_embedded_trace_id")) 37 | } 38 | commonCfg := rmq.Args{Log: customLog} 39 | // Create an AMQP topology for our healthcheck, which uses a temporary exchange. 40 | // Design goals of danlock/rmq include reducing the amount of naked booleans in function signatures. 41 | topology := rmq.Topology{ 42 | Args: commonCfg, 43 | Exchanges: []rmq.Exchange{{Name: "healthcheck", Kind: amqp.ExchangeDirect, AutoDelete: true}}, 44 | Queues: []rmq.Queue{{Name: "healthcheck", AutoDelete: true}}, 45 | QueueBindings: []rmq.QueueBinding{{QueueName: "healthcheck", ExchangeName: "healthcheck"}}, 46 | } 47 | // danlock/rmq best practice is including your applications topology in your ConnectConfig 48 | cfg := rmq.ConnectArgs{Args: commonCfg, Topology: topology} 49 | // RabbitMQ best practice is to pub and sub on different AMQP connections to avoid TCP backpressure causing issues with message consumption. 50 | pubRMQConn := rmq.ConnectWithURLs(ctx, cfg, os.Getenv("TEST_AMQP_URI")) 51 | subRMQConn := rmq.ConnectWithURLs(ctx, cfg, os.Getenv("TEST_AMQP_URI")) 52 | 53 | // A rudimentary healthcheck of a rmq.Connection is to ensure it can get a Channel, but we can do better 54 | _, err := subRMQConn.MustChannel(ctx) 55 | if err != nil { 56 | panic("couldn't get a channel") 57 | } 58 | 59 | rmqCons := rmq.NewConsumer(subRMQConn, rmq.ConsumerArgs{ 60 | Args: commonCfg, 61 | Queue: topology.Queues[0], 62 | Qos: rmq.Qos{PrefetchCount: 10}, 63 | }) 64 | // Now we have a RabbitMQ queue with messages incoming on the deliveries channel, even if the network flakes. 65 | deliveries := rmqCons.Consume(ctx) 66 | 67 | rmqPub := rmq.NewPublisher(ctx, pubRMQConn, rmq.PublisherArgs{Args: commonCfg}) 68 | // Now we have an AMQP publisher that can sends messages with at least once delivery. 69 | // Generate "unique" messages for our healthchecker to check later 70 | baseMsg := rmq.Publishing{Exchange: topology.Exchanges[0].Name, Mandatory: true} 71 | msgOne := baseMsg 72 | msgOne.Body = []byte(time.Now().String()) 73 | msgTwo := baseMsg 74 | msgTwo.Body = []byte(time.Now().String()) 75 | 76 | pubCtx, pubCtxCancel := context.WithTimeoutCause(ctx, 10*time.Second, errors.New("He's dead, Jim")) 77 | defer pubCtxCancel() 78 | 79 | conf, err := rmqPub.PublishUntilConfirmed(pubCtx, 0, msgOne) 80 | if err != nil { 81 | panic("uh oh, context timed out?") 82 | } 83 | // PublishUntilConfirmed only returns once the amqp.DeferredConfirmation is Done(), 84 | // so you can check Acked() without fear that the return value is simply telling you that it's not Done() yet. 85 | if !conf.Acked() { 86 | panic("uh oh, nacked") 87 | } 88 | // PublishUntilAcked resends on nacks. A Healthcheck may want to do that instead, since restarting a service whenever a nack happens probably won't help. 89 | if err = rmqPub.PublishUntilAcked(ctx, 0, msgTwo); err != nil { 90 | panic("uh oh, context timed out?") 91 | } 92 | 93 | // Now that we've sent, make sure we can receive. 94 | for i := 0; i < 2; i++ { 95 | select { 96 | case <-time.After(time.Second): 97 | panic("where's my message?") 98 | case msg := <-deliveries: 99 | if !reflect.DeepEqual(msg.Body, msgOne.Body) && !reflect.DeepEqual(msg.Body, msgTwo.Body) { 100 | panic("realistically this would probably be an error with another instance using this healthcheck simultaneously. Prevent this with an unique exchange or topic exchange with unique routing keys.") 101 | } 102 | } 103 | } 104 | 105 | // We sent and received 2 messages, so we're probably healthy enough to survive until the next docker/kubernetes health check. 106 | } 107 | -------------------------------------------------------------------------------- /topology.go: -------------------------------------------------------------------------------- 1 | package rmq 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "log/slog" 8 | "time" 9 | 10 | amqp "github.com/rabbitmq/amqp091-go" 11 | ) 12 | 13 | // Exchange contains args for amqp.Channel.ExchangeDeclare 14 | type Exchange struct { 15 | Name string // Name is required by ExchangeDeclare. 16 | Kind string // Kind is required by ExchangeDeclare. amqp091-go exports valid values like amqp.ExchangeDirect, etc 17 | Durable bool 18 | AutoDelete bool 19 | Internal bool 20 | NoWait bool 21 | Passive bool 22 | Args amqp.Table 23 | } 24 | 25 | // ExchangeBinding contains args for amqp.Channel.ExchangeBind 26 | type ExchangeBinding struct { 27 | Destination string 28 | RoutingKey string 29 | Source string 30 | NoWait bool 31 | Args amqp.Table 32 | } 33 | 34 | // DeclareTopology declares an AMQP topology once. 35 | // If you want this to be redeclared automatically on connections, add your Topology to ConnectArgs instead. 36 | func DeclareTopology(ctx context.Context, amqpConn AMQPConnection, topology Topology) error { 37 | logPrefix := fmt.Sprintf("rmq.DeclareTopology for AMQPConnection (%s -> %s)", amqpConn.LocalAddr(), amqpConn.RemoteAddr()) 38 | 39 | if topology.empty() { 40 | return nil 41 | } 42 | dontLog := topology.Log == nil 43 | topology.setDefaults() 44 | ctx, cancel := context.WithTimeout(ctx, topology.AMQPTimeout) 45 | defer cancel() 46 | 47 | // amqp091 currently does not use contexts all throughout, and therefore any call could block forever if the network is temperamental that day. 48 | // Call them in a goroutine so we can bail if necessary 49 | start := time.Now() 50 | errChan := make(chan error, 1) 51 | shouldLog := make(chan struct{}) 52 | 53 | go func() { 54 | mqChan, err := amqpConn.Channel() 55 | if err != nil { 56 | errChan <- fmt.Errorf(logPrefix+" failed to get amqp.Channel due to %w", err) 57 | return 58 | } 59 | err = topology.declare(ctx, mqChan) 60 | // An amqp.Channel must not be used from multiple goroutines simultaneously, so close it inside this goroutine to prevent cryptic RabbitMQ errors. 61 | mqChanErr := mqChan.Close() 62 | // Should we join mqChanErr if err is nil? When declare succeeeds a Close error is fairly inconsequential. Unless it leaves the channel in a bad state... 63 | // But we don't actually use the channel after this. Maybe just log it in that case? Food for thought. 64 | if mqChanErr != nil && !errors.Is(mqChanErr, amqp.ErrClosed) { 65 | err = errors.Join(err, mqChanErr) 66 | } 67 | 68 | select { 69 | case <-shouldLog: 70 | topology.Log(ctx, slog.LevelWarn, logPrefix+" completed after it's context finished. It took %s. Err: %+v", time.Since(start), err) 71 | default: 72 | errChan <- err 73 | } 74 | }() 75 | 76 | select { 77 | case <-ctx.Done(): 78 | // Log our leaked goroutine's response whenever it finally finishes since it may have useful debugging information. 79 | if !dontLog { 80 | close(shouldLog) 81 | } 82 | return fmt.Errorf(logPrefix+" unable to complete before context due to %w", context.Cause(ctx)) 83 | case err := <-errChan: 84 | return err 85 | } 86 | } 87 | 88 | // Topology contains all the exchange, queue and binding information needed for your application to use RabbitMQ. 89 | type Topology struct { 90 | Args 91 | 92 | Exchanges []Exchange 93 | ExchangeBindings []ExchangeBinding 94 | Queues []Queue 95 | QueueBindings []QueueBinding 96 | } 97 | 98 | func (t *Topology) empty() bool { 99 | return len(t.Exchanges) == 0 && len(t.Queues) == 0 && 100 | len(t.ExchangeBindings) == 0 && len(t.QueueBindings) == 0 101 | } 102 | 103 | func (t *Topology) declare(ctx context.Context, mqChan *amqp.Channel) (err error) { 104 | logPrefix := fmt.Sprintf("rmq.Topology.declare ") 105 | 106 | for _, e := range t.Exchanges { 107 | exchangeDeclare := mqChan.ExchangeDeclare 108 | if e.Passive { 109 | exchangeDeclare = mqChan.ExchangeDeclarePassive 110 | } 111 | err = exchangeDeclare(e.Name, e.Kind, e.Durable, e.AutoDelete, e.Internal, e.NoWait, e.Args) 112 | 113 | if err != nil { 114 | return fmt.Errorf(logPrefix+" failed to declare exchange %s due to %w", e.Name, err) 115 | } else if err = context.Cause(ctx); err != nil { 116 | return fmt.Errorf(logPrefix+" failed to declare exchanges before context ended due to %w", err) 117 | } 118 | } 119 | 120 | for _, eb := range t.ExchangeBindings { 121 | err = mqChan.ExchangeBind(eb.Destination, eb.RoutingKey, eb.Source, eb.NoWait, eb.Args) 122 | if err != nil { 123 | return fmt.Errorf(logPrefix+" failed to bind exchange %s to %s due to %w", eb.Destination, eb.Source, err) 124 | } else if err = context.Cause(ctx); err != nil { 125 | return fmt.Errorf(logPrefix+" failed to declare exchange bindings before context ended due to %w", err) 126 | } 127 | } 128 | 129 | for _, q := range t.Queues { 130 | if q.Name == "" { 131 | // Anonymous Queues auto generate different names on different amqp.Channel's. 132 | // Queues like this must be declared by the Consumer instead so it can receive messages from the queue, so we skip them here. 133 | // Should we log a warning? Seems like that would get annoying if you just wanted to pass the same Queue struct around. 134 | continue 135 | } 136 | 137 | queueDeclare := mqChan.QueueDeclare 138 | if q.Passive { 139 | queueDeclare = mqChan.QueueDeclarePassive 140 | } 141 | _, err = queueDeclare(q.Name, q.Durable, q.AutoDelete, q.Exclusive, q.NoWait, q.Args) 142 | 143 | if err != nil { 144 | return fmt.Errorf(logPrefix+" failed to declare queue due to %w", err) 145 | } else if err = context.Cause(ctx); err != nil { 146 | return fmt.Errorf(logPrefix+" failed to declare queues before context ended due to %w", err) 147 | } 148 | } 149 | 150 | for _, b := range t.QueueBindings { 151 | err = mqChan.QueueBind(b.QueueName, b.RoutingKey, b.ExchangeName, b.NoWait, b.Args) 152 | if err != nil { 153 | return fmt.Errorf(logPrefix+" unable to bind queue to exchange '%s' via key '%s' due to %w", b.ExchangeName, b.RoutingKey, err) 154 | } else if err = context.Cause(ctx); err != nil { 155 | return fmt.Errorf(logPrefix+" failed to declare queue bindings before context ended due to %w", err) 156 | } 157 | } 158 | 159 | return nil 160 | } 161 | -------------------------------------------------------------------------------- /benchmark_int_test.go: -------------------------------------------------------------------------------- 1 | //go:build rabbit 2 | 3 | package rmq_test 4 | 5 | import ( 6 | "bytes" 7 | "context" 8 | "fmt" 9 | "log/slog" 10 | "os" 11 | "strconv" 12 | "testing" 13 | "time" 14 | 15 | "github.com/danlock/rmq" 16 | amqp "github.com/rabbitmq/amqp091-go" 17 | ) 18 | 19 | const benchNumPubs = 1000 20 | 21 | func generatePublishings(num int, routingKey string) []rmq.Publishing { 22 | publishings := make([]rmq.Publishing, num) 23 | for i := range publishings { 24 | publishings[i] = rmq.Publishing{ 25 | RoutingKey: routingKey, 26 | Mandatory: true, 27 | Publishing: amqp.Publishing{ 28 | Body: []byte(fmt.Sprintf("%d.%d", i, time.Now().UnixNano())), 29 | }, 30 | } 31 | } 32 | return publishings 33 | } 34 | 35 | func BenchmarkPublishAndConsumeMany(b *testing.B) { 36 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute) 37 | defer cancel() 38 | 39 | randSuffix := fmt.Sprintf("%d.%p", time.Now().UnixNano(), b) 40 | 41 | queueName := "BenchmarkPublishAndConsumeMany" + randSuffix 42 | baseCfg := rmq.Args{Log: slog.Log} 43 | topology := rmq.Topology{ 44 | Args: baseCfg, 45 | Queues: []rmq.Queue{{ 46 | Name: queueName, 47 | Args: amqp.Table{ 48 | amqp.QueueTTLArg: time.Minute.Milliseconds(), 49 | }, 50 | }}, 51 | } 52 | 53 | subRMQConn := rmq.ConnectWithURLs(ctx, rmq.ConnectArgs{Args: baseCfg, Topology: topology}, os.Getenv("TEST_AMQP_URI")) 54 | pubRMQConn := rmq.ConnectWithURLs(ctx, rmq.ConnectArgs{Args: baseCfg, Topology: topology}, os.Getenv("TEST_AMQP_URI")) 55 | 56 | consumer := rmq.NewConsumer(subRMQConn, rmq.ConsumerArgs{ 57 | Args: baseCfg, 58 | Queue: topology.Queues[0], 59 | }) 60 | 61 | publisher := rmq.NewPublisher(ctx, pubRMQConn, rmq.PublisherArgs{ 62 | Args: baseCfg, 63 | LogReturns: true, 64 | }) 65 | 66 | // publisher2, publisher3 := rmq.NewPublisher(ctx, pubRMQConn, rmq.PublisherArgs{ 67 | // Args: baseCfg, 68 | // LogReturns: true, 69 | // }), rmq.NewPublisher(ctx, pubRMQConn, rmq.PublisherArgs{ 70 | // Args: baseCfg, 71 | // LogReturns: true, 72 | // }) 73 | 74 | dot := []byte(".") 75 | errChan := make(chan error) 76 | consumeChan := consumer.Consume(ctx) 77 | 78 | publishings := generatePublishings(benchNumPubs, queueName) 79 | 80 | cases := []struct { 81 | name string 82 | publishFunc func(b *testing.B) 83 | }{ 84 | { 85 | "PublishBatchUntilAcked", 86 | func(b *testing.B) { 87 | if err := publisher.PublishBatchUntilAcked(ctx, 0, publishings...); err != nil { 88 | b.Fatalf("PublishBatchUntilAcked err %v", err) 89 | } 90 | }, 91 | }, 92 | // { 93 | // "PublishBatchUntilAcked into thirds", 94 | // func(b *testing.B) { 95 | // errChan := make(chan error) 96 | // publishers := []*rmq.Publisher{publisher, publisher, publisher} 97 | // for i := range publishers { 98 | // go func(i int) { 99 | // errChan <- publishers[i].PublishBatchUntilAcked(ctx, 0, publishings[i:i+1]...) 100 | // }(i) 101 | // } 102 | // successes := 0 103 | // for { 104 | // select { 105 | // case err := <-errChan: 106 | // if err != nil { 107 | // b.Fatalf("PublishBatchUntilAcked err %v", err) 108 | // } 109 | // successes++ 110 | // if successes == len(publishers) { 111 | // return 112 | // } 113 | // case <-ctx.Done(): 114 | // b.Fatalf("PublishBatchUntilAcked timed out") 115 | // } 116 | // } 117 | // }, 118 | // }, 119 | // { 120 | // "PublishBatchUntilAcked on three Publishers", 121 | // func(b *testing.B) { 122 | // errChan := make(chan error) 123 | // publishers := []*rmq.Publisher{publisher, publisher2, publisher3} 124 | // for i := range publishers { 125 | // go func(i int) { 126 | // errChan <- publishers[i].PublishBatchUntilAcked(ctx, 0, publishings[i:i+1]...) 127 | // }(i) 128 | // } 129 | // successes := 0 130 | // for { 131 | // select { 132 | // case err := <-errChan: 133 | // if err != nil { 134 | // b.Fatalf("PublishBatchUntilAcked err %v", err) 135 | // } 136 | // successes++ 137 | // if successes == len(publishers) { 138 | // return 139 | // } 140 | // case <-ctx.Done(): 141 | // b.Fatalf("PublishBatchUntilAcked timed out") 142 | // } 143 | // } 144 | // }, 145 | // }, 146 | // { 147 | // "Concurrent PublishUntilAcked", 148 | // func(b *testing.B) { 149 | // errChan := make(chan error) 150 | // for i := range publishings { 151 | // go func(i int) { 152 | // errChan <- publisher.PublishUntilAcked(ctx, 0, publishings[i]) 153 | // }(i) 154 | // } 155 | // successes := 0 156 | // for { 157 | // select { 158 | // case err := <-errChan: 159 | // if err != nil { 160 | // b.Fatalf("PublishUntilAcked err %v", err) 161 | // } 162 | // successes++ 163 | // if successes == len(publishings) { 164 | // return 165 | // } 166 | // case <-ctx.Done(): 167 | // b.Fatalf("PublishUntilAcked timed out") 168 | // } 169 | // } 170 | // }, 171 | // }, 172 | } 173 | 174 | for _, bb := range cases { 175 | b.Run(bb.name, func(b *testing.B) { 176 | for i := 0; i < b.N; i++ { 177 | go func(i int) (err error) { 178 | received := make(map[uint64]struct{}, len(publishings)) 179 | defer func() { errChan <- err }() 180 | for { 181 | select { 182 | case msg := <-consumeChan: 183 | rawIndex := bytes.Split(msg.Body, dot)[0] 184 | index, err := strconv.ParseUint(string(rawIndex), 10, 64) 185 | if err != nil { 186 | return fmt.Errorf("strconv.ParseUint err %w", err) 187 | } 188 | received[index] = struct{}{} 189 | if err := msg.Ack(false); err != nil { 190 | return fmt.Errorf("msg.Ack err %w", err) 191 | } 192 | if len(received) == len(publishings) { 193 | return nil 194 | } 195 | case <-ctx.Done(): 196 | return fmt.Errorf("timed out after consuming %d publishings on bench run %d", len(received), i) 197 | } 198 | } 199 | }(i) 200 | 201 | bb.publishFunc(b) 202 | 203 | select { 204 | case <-ctx.Done(): 205 | b.Fatalf("timed out on bench run %d", i) 206 | case err := <-errChan: 207 | if err != nil { 208 | b.Fatalf("on bench run %d consumer err %v", i, err) 209 | 210 | } 211 | } 212 | } 213 | }) 214 | } 215 | } 216 | -------------------------------------------------------------------------------- /conn_test.go: -------------------------------------------------------------------------------- 1 | package rmq 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "log/slog" 8 | "net" 9 | "testing" 10 | "time" 11 | 12 | amqp "github.com/rabbitmq/amqp091-go" 13 | ) 14 | 15 | type MockAMQPConnection struct { 16 | ChannelFn func() (*amqp.Channel, error) 17 | NotifyCloseChan chan *amqp.Error 18 | } 19 | 20 | func (m *MockAMQPConnection) Channel() (*amqp.Channel, error) { 21 | if m.ChannelFn != nil { 22 | return m.ChannelFn() 23 | } 24 | return nil, fmt.Errorf("sike") 25 | } 26 | func (m *MockAMQPConnection) NotifyClose(receiver chan *amqp.Error) chan *amqp.Error { 27 | if m.NotifyCloseChan != nil { 28 | return m.NotifyCloseChan 29 | } 30 | return receiver 31 | } 32 | func (m *MockAMQPConnection) LocalAddr() net.Addr { 33 | return &net.UnixAddr{"MockAMQPConnection", "unix"} 34 | } 35 | func (m *MockAMQPConnection) RemoteAddr() net.Addr { 36 | return &net.UnixAddr{"MockAMQPConnection", "unix"} 37 | } 38 | func (m *MockAMQPConnection) CloseDeadline(time.Time) error { 39 | return nil 40 | } 41 | func (m *MockAMQPConnection) IsClosed() bool { 42 | return false 43 | } 44 | 45 | func TestConnect(t *testing.T) { 46 | errs := make(chan any, 1) 47 | go func() { 48 | defer func() { 49 | errs <- recover() 50 | }() 51 | _ = Connect(nil, ConnectArgs{}, nil) 52 | }() 53 | result := <-errs 54 | if result == nil { 55 | t.Fatalf("Connect should panic when missing required arguments") 56 | } 57 | rmqConn := Connect(context.Background(), ConnectArgs{}, func() (AMQPConnection, error) { return nil, fmt.Errorf("sike") }) 58 | if rmqConn == nil { 59 | t.Fatalf("Connect failed to return a rmq.Connection") 60 | } 61 | } 62 | 63 | func TestConnection_CurrentConnection(t *testing.T) { 64 | canceledCtx, cancel := context.WithCancel(context.Background()) 65 | cancel() 66 | tests := []struct { 67 | name string 68 | connCtx, reqCtx context.Context 69 | connDialFn func() (AMQPConnection, error) 70 | wantErr error 71 | }{ 72 | { 73 | "success", 74 | context.Background(), context.Background(), 75 | func() (AMQPConnection, error) { return &MockAMQPConnection{}, nil }, 76 | nil, 77 | }, 78 | { 79 | "failed due to request context canceled", 80 | context.Background(), canceledCtx, 81 | func() (AMQPConnection, error) { return &MockAMQPConnection{}, nil }, 82 | context.Canceled, 83 | }, 84 | { 85 | "failed due to connection context canceled", 86 | canceledCtx, context.Background(), 87 | func() (AMQPConnection, error) { return &MockAMQPConnection{}, nil }, 88 | context.Canceled, 89 | }, 90 | } 91 | for _, tt := range tests { 92 | t.Run(tt.name, func(t *testing.T) { 93 | rmqConn := Connect(tt.connCtx, ConnectArgs{}, tt.connDialFn) 94 | got, err := rmqConn.CurrentConnection(tt.reqCtx) 95 | if !errors.Is(err, tt.wantErr) { 96 | t.Fatalf("rmq.Connection.CurrentConnection() error = %v, wantErr %v", err, tt.wantErr) 97 | } 98 | if err == nil && got == nil { 99 | t.Errorf("rmq.Connection.CurrentConnection() should have returned an AMQPConnection") 100 | } 101 | }) 102 | } 103 | } 104 | 105 | func TestConnection_Channel(t *testing.T) { 106 | closeChan := make(chan *amqp.Error, 5) 107 | badErr := fmt.Errorf("shucks") 108 | // close the connection, redialer should grab a new one in time 109 | closeChan <- amqp.ErrClosed 110 | goodMockAMQP := &MockAMQPConnection{ 111 | ChannelFn: func() (*amqp.Channel, error) { return nil, nil }, 112 | NotifyCloseChan: closeChan} 113 | badMockAMQP := &MockAMQPConnection{ 114 | ChannelFn: func() (*amqp.Channel, error) { return nil, badErr }} 115 | slowMockAMQP := &MockAMQPConnection{ 116 | ChannelFn: func() (*amqp.Channel, error) { 117 | time.Sleep(time.Second / 4) 118 | return nil, nil 119 | }} 120 | 121 | ctx, cancel := context.WithCancel(context.Background()) 122 | defer cancel() 123 | baseCfg := Args{Log: slog.Log} 124 | connConf := ConnectArgs{ 125 | Args: baseCfg, 126 | } 127 | goodRMQConn := Connect(ctx, connConf, func() (AMQPConnection, error) { 128 | return goodMockAMQP, nil 129 | }) 130 | badRMQConn := Connect(ctx, connConf, func() (AMQPConnection, error) { 131 | return badMockAMQP, nil 132 | }) 133 | slowRMQConn := Connect(ctx, connConf, func() (AMQPConnection, error) { 134 | return slowMockAMQP, nil 135 | }) 136 | slowConfig := ConnectArgs{Args: Args{Log: slog.Log, AMQPTimeout: 50 * time.Millisecond}} 137 | slowUsingTimeoutRMQConn := Connect(ctx, slowConfig, func() (AMQPConnection, error) { 138 | return slowMockAMQP, nil 139 | }) 140 | 141 | shortCtx, shortCancel := context.WithTimeout(ctx, 50*time.Millisecond) 142 | defer shortCancel() 143 | tests := []struct { 144 | name string 145 | ctx context.Context 146 | rmqConn *Connection 147 | wantErr error 148 | }{ 149 | { 150 | "success", 151 | ctx, 152 | goodRMQConn, 153 | nil, 154 | }, { 155 | "failed", 156 | ctx, 157 | badRMQConn, 158 | badErr, 159 | }, { 160 | "ctx finished before slow channel", 161 | shortCtx, 162 | slowRMQConn, 163 | context.DeadlineExceeded, 164 | }, { 165 | "AMQPChannelTimeout before slow channel", 166 | ctx, 167 | slowUsingTimeoutRMQConn, 168 | context.DeadlineExceeded, 169 | }, 170 | } 171 | for _, tt := range tests { 172 | t.Run(tt.name, func(t *testing.T) { 173 | ctx, cancel := context.WithTimeout(tt.ctx, time.Second) 174 | defer cancel() 175 | _, err := tt.rmqConn.Channel(ctx) 176 | if !errors.Is(err, tt.wantErr) { 177 | t.Errorf("rmq.Connection.Channel() error = %v, wantErr %v", err, tt.wantErr) 178 | return 179 | } 180 | }) 181 | } 182 | 183 | connConf = ConnectArgs{ 184 | Args: Args{ 185 | Log: connConf.Log, 186 | Delay: func(attempt int) time.Duration { return time.Duration(attempt) * time.Millisecond }, 187 | }, 188 | } 189 | 190 | flakyCount := 0 191 | dialErr := fmt.Errorf("dial fail") 192 | flakyRMQConn := Connect(ctx, connConf, func() (AMQPConnection, error) { 193 | flakyCount++ 194 | if flakyCount == 5 { 195 | return badMockAMQP, nil 196 | } else if flakyCount > 6 { 197 | return goodMockAMQP, nil 198 | } 199 | return nil, dialErr 200 | }) 201 | 202 | _, err := flakyRMQConn.Channel(ctx) 203 | if !errors.Is(err, badErr) { 204 | t.Fatalf("rmq.Connection.Channel() error = %v, wantErr %v", err, badErr) 205 | } 206 | _, err = flakyRMQConn.Channel(ctx) 207 | if err != nil { 208 | t.Fatalf("rmq.Connection.Channel() error = %v", err) 209 | } 210 | 211 | midChanCtx, midChanCancel := context.WithCancel(context.Background()) 212 | 213 | midChanMockAMQP := &MockAMQPConnection{ 214 | ChannelFn: func() (*amqp.Channel, error) { 215 | midChanCancel() 216 | time.Sleep(time.Second / 4) 217 | return nil, nil 218 | }} 219 | midChanRMQ := Connect(midChanCtx, connConf, func() (AMQPConnection, error) { 220 | return midChanMockAMQP, nil 221 | }) 222 | if _, err = midChanRMQ.Channel(ctx); !errors.Is(err, context.Canceled) { 223 | t.Fatalf("rmq.Connection.Channel() error = %v, wantErr %v", err, context.Canceled) 224 | } 225 | } 226 | -------------------------------------------------------------------------------- /consumer_int_test.go: -------------------------------------------------------------------------------- 1 | //go:build rabbit 2 | 3 | package rmq_test 4 | 5 | import ( 6 | "bytes" 7 | "context" 8 | "fmt" 9 | "log/slog" 10 | "math/rand" 11 | "os" 12 | "reflect" 13 | "strconv" 14 | "sync" 15 | "testing" 16 | "time" 17 | 18 | "github.com/danlock/rmq" 19 | amqp "github.com/rabbitmq/amqp091-go" 20 | ) 21 | 22 | func ForceRedial(ctx context.Context, rmqConn *rmq.Connection) error { 23 | amqpConn, err := rmqConn.CurrentConnection(ctx) 24 | if err != nil { 25 | return fmt.Errorf("rmqConn.CurrentConnection failed because %w", err) 26 | } 27 | // close the current connection to force a redial 28 | return amqpConn.CloseDeadline(time.Now().Add(time.Minute)) 29 | } 30 | 31 | func TestConsumer(t *testing.T) { 32 | ctx, cancel := context.WithCancel(context.Background()) 33 | defer cancel() 34 | slogLog := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug})).Log 35 | connectCfg := rmq.ConnectArgs{Args: rmq.Args{Log: slogLog}} 36 | rmqConn := rmq.ConnectWithURLs(ctx, connectCfg, "amqp://dont.exist", os.Getenv("TEST_AMQP_URI")) 37 | 38 | baseConsConfig := rmq.ConsumerArgs{ 39 | Args: connectCfg.Args, 40 | Queue: rmq.Queue{ 41 | Name: fmt.Sprintf("TestRMQConsumer.%p", t), 42 | Args: amqp.Table{amqp.QueueTTLArg: time.Minute.Milliseconds()}, 43 | }, 44 | Consume: rmq.Consume{ 45 | Consumer: "TestConsumer", 46 | }, 47 | Qos: rmq.Qos{ 48 | PrefetchCount: 2000, 49 | }, 50 | } 51 | 52 | baseConsumer := rmq.NewConsumer(rmqConn, baseConsConfig) 53 | 54 | canceledCtx, canceledCancel := context.WithCancel(ctx) 55 | canceledCancel() 56 | // ConsumeConcurrently should exit immediately on canceled contexts. 57 | baseConsumer.ConsumeConcurrently(canceledCtx, 0, nil) 58 | 59 | rmqBaseConsMessages := make(chan amqp.Delivery, 10) 60 | go baseConsumer.ConsumeConcurrently(ctx, 0, func(ctx context.Context, msg amqp.Delivery) { 61 | rmqBaseConsMessages <- msg 62 | _ = msg.Ack(false) 63 | }) 64 | time.Sleep(time.Second / 10) 65 | unreliableRMQPub := rmq.NewPublisher(ctx, rmqConn, rmq.PublisherArgs{DontConfirm: true}) 66 | unreliableRMQPub.Publish(ctx, rmq.Publishing{Exchange: "amq.fanout"}) 67 | rmqPub := rmq.NewPublisher(ctx, rmqConn, rmq.PublisherArgs{Args: connectCfg.Args}) 68 | 69 | ForceRedial(ctx, rmqConn) 70 | pubCtx, pubCancel := context.WithTimeout(ctx, 20*time.Second) 71 | defer pubCancel() 72 | 73 | wantedPub := rmq.Publishing{RoutingKey: baseConsConfig.Queue.Name} 74 | wantedPub.Body = []byte("TestRMQPublisher") 75 | 76 | pubCount := 3 77 | errChan := make(chan error, pubCount) 78 | for i := 0; i < pubCount; i++ { 79 | go func() { 80 | errChan <- rmqPub.PublishUntilAcked(pubCtx, 0, wantedPub) 81 | }() 82 | } 83 | ForceRedial(ctx, rmqConn) 84 | 85 | for i := 0; i < pubCount; i++ { 86 | if err := <-errChan; err != nil { 87 | t.Fatalf("PublishUntilAcked returned unexpected error %v", err) 88 | } 89 | if i%2 == 0 { 90 | ForceRedial(ctx, rmqConn) 91 | } 92 | } 93 | 94 | for i := 0; i < pubCount; i++ { 95 | var msg amqp.Delivery 96 | select { 97 | case <-pubCtx.Done(): 98 | t.Fatalf("timed out waiting for published message %d", i) 99 | case msg = <-rmqBaseConsMessages: 100 | } 101 | 102 | if !reflect.DeepEqual(msg.Body, wantedPub.Body) { 103 | t.Fatalf("Received unexpected message %s", string(msg.Body)) 104 | } 105 | } 106 | } 107 | 108 | func TestConsumer_Load(t *testing.T) { 109 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute/2) 110 | defer cancel() 111 | logf := slog.Log 112 | 113 | baseName := fmt.Sprint("TestRMQConsumer_Load_Base_", rand.Uint64()) 114 | prefetchName := fmt.Sprint("TestRMQConsumer_Load_Prefetch_", rand.Uint64()) 115 | baseCfg := rmq.Args{Log: logf} 116 | topology := rmq.Topology{ 117 | Args: baseCfg, 118 | Queues: []rmq.Queue{{ 119 | Name: baseName, 120 | Args: amqp.Table{amqp.QueueTTLArg: time.Minute.Milliseconds()}, 121 | }, { 122 | Name: prefetchName, 123 | Args: amqp.Table{amqp.QueueTTLArg: time.Minute.Milliseconds()}, 124 | }}, 125 | } 126 | 127 | rmqConn := rmq.ConnectWithAMQPConfig(ctx, rmq.ConnectArgs{Args: baseCfg, Topology: topology}, os.Getenv("TEST_AMQP_URI"), amqp.Config{}) 128 | 129 | periodicallyCloseConn := func() { 130 | for { 131 | select { 132 | case <-ctx.Done(): 133 | return 134 | case <-time.After(time.Second): 135 | ForceRedial(ctx, rmqConn) 136 | } 137 | } 138 | } 139 | go periodicallyCloseConn() 140 | 141 | baseConsConfig := rmq.ConsumerArgs{ 142 | Args: baseCfg, 143 | Queue: topology.Queues[0], 144 | Consume: rmq.Consume{ 145 | Consumer: baseName, 146 | }, 147 | } 148 | 149 | prefetchConsConfig := baseConsConfig 150 | prefetchConsConfig.Queue = topology.Queues[1] 151 | prefetchConsConfig.Qos.PrefetchCount = 10 152 | prefetchConsConfig.Queue.Name = prefetchName 153 | 154 | consumers := []rmq.ConsumerArgs{baseConsConfig, prefetchConsConfig} 155 | publisher := rmq.NewPublisher(ctx, rmqConn, rmq.PublisherArgs{Args: baseCfg}) 156 | 157 | msgCount := 5_000 158 | errChan := make(chan error, (msgCount/2+1)*len(consumers)) 159 | for _, c := range consumers { 160 | c := c 161 | go func() { 162 | ctx, cancel := context.WithCancel(ctx) 163 | receives := make(map[int]struct{}) 164 | var msgRecv uint64 165 | var consMu sync.Mutex 166 | rmq.NewConsumer(rmqConn, c).ConsumeConcurrently(ctx, 0, func(ctx context.Context, msg amqp.Delivery) { 167 | if !c.Consume.AutoAck { 168 | defer msg.Ack(false) 169 | } 170 | indexBytes := bytes.TrimPrefix(msg.Body, []byte(c.Queue.Name+":")) 171 | index, err := strconv.Atoi(string(indexBytes)) 172 | consMu.Lock() 173 | defer consMu.Unlock() 174 | if err != nil { 175 | logf(ctx, slog.LevelError, "%s got %d msgs. Last msg %s", c.Queue.Name, msgRecv, string(msg.Body)) 176 | errChan <- err 177 | cancel() 178 | } else { 179 | msgRecv++ 180 | receives[index] = struct{}{} 181 | if len(receives) == msgCount { 182 | logf(ctx, slog.LevelError, "%s got %d msgs", c.Queue.Name, msgRecv) 183 | errChan <- nil 184 | cancel() 185 | } 186 | } 187 | }) 188 | }() 189 | go func() { 190 | // Send half of the messages with an incredibly inefficient use of goroutines, and the rest in a PublishBatchUntilAcked. 191 | // Publishing all of this stuff in different goroutines should not cause any races. 192 | for i := 0; i < msgCount/2; i++ { 193 | go func(i int) { 194 | errChan <- publisher.PublishUntilAcked(ctx, 0, rmq.Publishing{ 195 | RoutingKey: c.Queue.Name, 196 | Mandatory: true, 197 | Publishing: amqp.Publishing{ 198 | Body: []byte(fmt.Sprint(c.Queue.Name, ":", i)), 199 | }, 200 | }) 201 | }(i) 202 | } 203 | pubs := make([]rmq.Publishing, msgCount/2) 204 | for i := range pubs { 205 | pubs[i] = rmq.Publishing{ 206 | RoutingKey: c.Queue.Name, 207 | Mandatory: true, 208 | Publishing: amqp.Publishing{ 209 | Body: []byte(fmt.Sprint(c.Queue.Name, ":", i+len(pubs))), 210 | }, 211 | } 212 | } 213 | errChan <- publisher.PublishBatchUntilAcked(ctx, 0, pubs...) 214 | }() 215 | } 216 | 217 | for i := 0; i < cap(errChan); i++ { 218 | select { 219 | case <-ctx.Done(): 220 | t.Fatalf("timed out after %d receives waiting for consumers to finish", i) 221 | case err := <-errChan: 222 | if err != nil { 223 | t.Fatalf("after %d receives got err from consumer %+v", i, err) 224 | } 225 | } 226 | } 227 | } 228 | 229 | // RabbitMQ behaviour around auto generated names and restricting declaring queues with amq prefix 230 | func TestRMQConsumer_AutogeneratedQueueNames(t *testing.T) { 231 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) 232 | defer cancel() 233 | slog.SetDefault(slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug}))) 234 | baseCfg := rmq.Args{Log: slog.Log} 235 | rmqConn := rmq.ConnectWithAMQPConfig(ctx, rmq.ConnectArgs{Args: baseCfg}, os.Getenv("TEST_AMQP_URI"), amqp.Config{}) 236 | 237 | // NewConsumer with an empty Queue.Name will declare a queue with a RabbitMQ generated name 238 | // This is useless unless the config also includes QueueBindings, since reconnections cause RabbitMQ to generate a different name anyway 239 | cons := rmq.NewConsumer(rmqConn, rmq.ConsumerArgs{ 240 | Args: baseCfg, 241 | QueueBindings: []rmq.QueueBinding{ 242 | {ExchangeName: "amq.fanout", RoutingKey: "TestRMQConsumer_AutogeneratedQueueNames"}, 243 | }, 244 | Qos: rmq.Qos{PrefetchCount: 1}, 245 | }) 246 | deliveries := cons.Consume(ctx) 247 | // Wait a sec for Consume to actually bring up the queue, since otherwise a published message could happen before a queue is declared. 248 | // danlock/rmq best practice to only use queues named in your Topology so you won't have to remember this. 249 | time.Sleep(time.Second / 3) 250 | amqpConn, err := rmqConn.CurrentConnection(ctx) 251 | if err != nil { 252 | t.Fatalf("failed getting current connection %v", err) 253 | } 254 | amqpConn.CloseDeadline(time.Now().Add(time.Minute)) 255 | 256 | // Declaring again should work without errors, but it will create a different queue rather than consuming from the first one. 257 | // rmq.Consumer could remember the last queue name to consume from it again, but that wouldn't be reliable with auto-deleted or expiring queues. 258 | // It's simpler to disallow that use case by not making RabbitMQ generated queue names available from rmq.Consumer. 259 | secondDeliveries := cons.Consume(ctx) 260 | publisher := rmq.NewPublisher(ctx, rmqConn, rmq.PublisherArgs{Args: baseCfg, LogReturns: true}) 261 | pubCount := 10 262 | time.Sleep(time.Second / 3) 263 | 264 | for i := 0; i < pubCount; i++ { 265 | go publisher.PublishUntilAcked(ctx, 0, rmq.Publishing{Exchange: "amq.fanout", RoutingKey: "TestRMQConsumer_AutogeneratedQueueNames", Mandatory: true}) 266 | } 267 | 268 | for i := 0; i < pubCount; i++ { 269 | select { 270 | case msg := <-deliveries: 271 | msg.Ack(false) 272 | case msg := <-secondDeliveries: 273 | msg.Ack(false) 274 | case <-ctx.Done(): 275 | t.Fatalf("timed out on delivery %d", i) 276 | } 277 | } 278 | } 279 | -------------------------------------------------------------------------------- /connection.go: -------------------------------------------------------------------------------- 1 | package rmq 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "log/slog" 8 | "net" 9 | "sync" 10 | "time" 11 | 12 | "github.com/danlock/rmq/internal" 13 | amqp "github.com/rabbitmq/amqp091-go" 14 | ) 15 | 16 | // AMQPConnection abstracts an amqp091.Connection 17 | type AMQPConnection interface { 18 | Channel() (*amqp.Channel, error) 19 | NotifyClose(receiver chan *amqp.Error) chan *amqp.Error 20 | LocalAddr() net.Addr 21 | RemoteAddr() net.Addr 22 | CloseDeadline(time.Time) error 23 | IsClosed() bool 24 | } 25 | 26 | type ConnectArgs struct { 27 | Args 28 | // Topology will be declared each connection to mitigate downed RabbitMQ nodes. Recommended to set, but not required. 29 | Topology Topology 30 | // DisableAMQP091Logs ensures the Connection's logs will not include rabbitmq/amqp091-go log's. 31 | // Note only the first Connection created will include rabbitmq/amqp091-go logs, due to rabbitmq/amqp091-go using a single global logger. 32 | DisableAMQP091Logs bool 33 | } 34 | 35 | func ConnectWithURLs(ctx context.Context, conf ConnectArgs, amqpURLs ...string) *Connection { 36 | if len(amqpURLs) == 0 { 37 | panic("ConnectWithURLs needs amqpURLs!") 38 | } 39 | return Connect(ctx, conf, func() (AMQPConnection, error) { 40 | errs := make([]error, 0, len(amqpURLs)) 41 | for _, amqpURL := range amqpURLs { 42 | amqpConn, err := amqp.Dial(amqpURL) 43 | if err == nil { 44 | return amqpConn, nil 45 | } 46 | errs = append(errs, err) 47 | } 48 | return nil, errors.Join(errs...) 49 | }) 50 | } 51 | 52 | func ConnectWithAMQPConfig(ctx context.Context, conf ConnectArgs, amqpURL string, amqpConf amqp.Config) *Connection { 53 | return Connect(ctx, conf, func() (AMQPConnection, error) { 54 | return amqp.DialConfig(amqpURL, amqpConf) 55 | }) 56 | } 57 | 58 | var setAMQP091Logger sync.Once 59 | 60 | // Connect returns a resilient, redialable AMQP connection that runs until it's context is canceled. 61 | // Each Channel() call triggers rmq.Connection to return an amqp.Channel from it's CurrentConnection() or redial with the provided dialFn for a new AMQP Connection. 62 | // ConnectWith* functions provide a few simple dialFn's for ease of use. They can be a simple wrapper around an amqp.Dial or much more complicated. 63 | // If you want to ensure the Connection is working, call MustChannel with a timeout. 64 | func Connect(ctx context.Context, conf ConnectArgs, dialFn func() (AMQPConnection, error)) *Connection { 65 | if dialFn == nil || ctx == nil { 66 | panic("Connect requires a ctx and a dialFn") 67 | } 68 | // Thread safely set the amqp091 logger once so it's included within danlock/rmq Connection Logs. 69 | // We use a sync.Once to avoid races, but this also means just the first Connection logs these errors. 70 | // It could be useful to have every Connection log these but practically that would lead to unneccessary log spam. 71 | // If this behaviour causes an issue your only solution is to set DisableAMQP091Logs and call amqp.SetLogger yourself 72 | if conf.Log != nil && !conf.DisableAMQP091Logs { 73 | setAMQP091Logger.Do(func() { 74 | amqp.SetLogger(internal.AMQP091Logger{ctx, conf.Log}) 75 | }) 76 | } 77 | conf.setDefaults() 78 | 79 | conn := Connection{ 80 | ctx: ctx, 81 | chanReqChan: make(chan internal.ChanReq[*amqp.Channel]), 82 | currentConReqChan: make(chan internal.ChanReq[AMQPConnection]), 83 | 84 | config: conf, 85 | } 86 | 87 | go conn.redial(dialFn) 88 | 89 | return &conn 90 | } 91 | 92 | func request[T any](connCtx, ctx context.Context, reqChan chan internal.ChanReq[T]) (t T, _ error) { 93 | if ctx == nil { 94 | return t, fmt.Errorf("nil context") 95 | } 96 | respChan := make(chan internal.ChanResp[T], 1) 97 | select { 98 | case <-connCtx.Done(): 99 | return t, fmt.Errorf("rmq.Connection context timed out because %w", context.Cause(connCtx)) 100 | case <-ctx.Done(): 101 | return t, fmt.Errorf("request context timed out because %w", context.Cause(ctx)) 102 | case reqChan <- internal.ChanReq[T]{Ctx: ctx, RespChan: respChan}: 103 | } 104 | 105 | select { 106 | case <-connCtx.Done(): 107 | return t, fmt.Errorf("rmq.Connection context timed out because %w", context.Cause(connCtx)) 108 | case <-ctx.Done(): 109 | return t, fmt.Errorf("request context timed out because %w", context.Cause(ctx)) 110 | case resp := <-respChan: 111 | return resp.Val, resp.Err 112 | } 113 | } 114 | 115 | // Connection is a threadsafe, redialable wrapper around an amqp091.Connection 116 | type Connection struct { 117 | ctx context.Context 118 | 119 | chanReqChan chan internal.ChanReq[*amqp.Channel] 120 | currentConReqChan chan internal.ChanReq[AMQPConnection] 121 | 122 | config ConnectArgs 123 | } 124 | 125 | // Channel requests an AMQP channel from the current AMQP Connection. 126 | // On errors the rmq.Connection will redial, and the caller is expected to call Channel() again for a new connection. 127 | func (c *Connection) Channel(ctx context.Context) (*amqp.Channel, error) { 128 | ctx, cancel := context.WithTimeout(ctx, c.config.AMQPTimeout) 129 | defer cancel() 130 | return request(c.ctx, ctx, c.chanReqChan) 131 | } 132 | 133 | // MustChannel calls Channel on an active Connection until it's context times out or it successfully gets a Channel. 134 | // Recommended to use context.WithTimeout. 135 | func (c *Connection) MustChannel(ctx context.Context) (*amqp.Channel, error) { 136 | logPrefix := "rmq.Connection.MustChannel" 137 | errs := make([]error, 0) 138 | for { 139 | select { 140 | case <-ctx.Done(): 141 | errs = append(errs, fmt.Errorf(logPrefix+" timed out due to %w", context.Cause(ctx))) 142 | return nil, errors.Join(errs...) 143 | case <-c.ctx.Done(): 144 | errs = append(errs, fmt.Errorf(logPrefix+" Connection timed out due to %w", context.Cause(c.ctx))) 145 | return nil, errors.Join(errs...) 146 | default: 147 | } 148 | 149 | mqChan, err := c.Channel(ctx) 150 | if err != nil { 151 | errs = append(errs, err) 152 | } else { 153 | return mqChan, nil 154 | } 155 | } 156 | } 157 | 158 | // CurrentConnection requests the current AMQPConnection being used by rmq.Connection. 159 | // It can be typecasted into an *amqp091.Connection. 160 | // Useful for making NotifyClose or NotifyBlocked channels for example. 161 | // If the CurrentConnection is closed, this function will return amqp.ErrClosed 162 | // until rmq.Connection dials successfully for another one. 163 | func (c *Connection) CurrentConnection(ctx context.Context) (AMQPConnection, error) { 164 | ctx, cancel := context.WithTimeout(ctx, c.config.AMQPTimeout) 165 | defer cancel() 166 | conn, err := request(c.ctx, ctx, c.currentConReqChan) 167 | if err != nil { 168 | return nil, err 169 | } 170 | if conn.IsClosed() { 171 | return nil, amqp.ErrClosed 172 | } 173 | return conn, nil 174 | } 175 | 176 | func (c *Connection) redial(dialFn func() (AMQPConnection, error)) { 177 | internal.Retry(c.ctx, c.config.Delay, func(delay time.Duration) (time.Duration, bool) { 178 | logPrefix := "rmq.Connection.redial" 179 | amqpConn, err := dialFn() 180 | if err != nil { 181 | c.config.Log(c.ctx, slog.LevelError, logPrefix+" failed, retrying after %s due to %+v", delay.String(), err) 182 | return 0, false 183 | } 184 | logPrefix = fmt.Sprintf("rmq.Connection.redial's AMQPConnection (%s -> %s)", amqpConn.LocalAddr(), amqpConn.RemoteAddr()) 185 | // Redeclare Topology if we have one. This has the bonus aspect of making sure the connection is actually usable, better than a Ping. 186 | if err := DeclareTopology(c.ctx, amqpConn, c.config.Topology); err != nil { 187 | c.config.Log(c.ctx, slog.LevelError, logPrefix+" DeclareTopology failed, retrying after %s due to %+v", delay.String(), err) 188 | return 0, false 189 | } 190 | start := time.Now() 191 | c.listen(amqpConn) 192 | return time.Since(start), true 193 | }) 194 | } 195 | 196 | // listen listens and responds to Channel and Connection requests. It returns on any failure to prompt another redial. 197 | func (c *Connection) listen(amqpConn AMQPConnection) { 198 | logPrefix := fmt.Sprintf("rmq.Connection's AMQPConnection (%s -> %s)", amqpConn.LocalAddr(), amqpConn.RemoteAddr()) 199 | notifyClose := amqpConn.NotifyClose(make(chan *amqp.Error, 1)) 200 | for { 201 | select { 202 | case <-c.ctx.Done(): 203 | if err := amqpConn.CloseDeadline(time.Now().Add(c.config.AMQPTimeout)); err != nil && !errors.Is(err, amqp.ErrClosed) { 204 | c.config.Log(c.ctx, slog.LevelError, logPrefix+" failed to close due to err: %+v", err) 205 | } 206 | return 207 | case err := <-notifyClose: 208 | if err != nil { 209 | c.config.Log(c.ctx, slog.LevelError, logPrefix+" received close notification err: %+v", err) 210 | } 211 | return 212 | case connReq := <-c.currentConReqChan: 213 | connReq.RespChan <- internal.ChanResp[AMQPConnection]{Val: amqpConn} 214 | case chanReq := <-c.chanReqChan: 215 | var resp internal.ChanResp[*amqp.Channel] 216 | resp.Val, resp.Err = c.safeChannel(chanReq.Ctx, amqpConn) 217 | chanReq.RespChan <- resp 218 | if resp.Err != nil { 219 | c.config.Log(c.ctx, slog.LevelDebug, logPrefix+" redialing due to %+v", resp.Err) 220 | return 221 | } 222 | } 223 | } 224 | } 225 | 226 | // safeChannel calls amqp.Connection.Channel with a timeout by launching it in a separate goroutine and waiting for the response. 227 | // This is inefficient, results in a leaked goroutine on timeout, but is the best we can do until amqp091 adds a context to the function. 228 | func (c *Connection) safeChannel(ctx context.Context, amqpConn AMQPConnection) (*amqp.Channel, error) { 229 | logPrefix := "rmq.Connection.safeChannel" 230 | respChan := make(chan internal.ChanResp[*amqp.Channel], 1) 231 | go func() { 232 | var resp internal.ChanResp[*amqp.Channel] 233 | resp.Val, resp.Err = amqpConn.Channel() 234 | // If our contexts timed out, close successfully created channels within this goroutine. 235 | if resp.Err == nil && resp.Val != nil && (c.ctx.Err() != nil || ctx.Err() != nil) { 236 | if err := resp.Val.Close(); err != nil && !errors.Is(err, amqp.ErrClosed) { 237 | c.config.Log(ctx, slog.LevelError, logPrefix+" failed to close channel due to err: %+v", err) 238 | } 239 | } else { 240 | respChan <- resp 241 | } 242 | }() 243 | 244 | select { 245 | case <-c.ctx.Done(): 246 | return nil, fmt.Errorf(logPrefix+" unable to complete before Connection %w", context.Cause(c.ctx)) 247 | case <-ctx.Done(): 248 | return nil, fmt.Errorf(logPrefix+" unable to complete before %w", context.Cause(ctx)) 249 | case resp := <-respChan: 250 | return resp.Val, resp.Err 251 | } 252 | } 253 | -------------------------------------------------------------------------------- /consumer.go: -------------------------------------------------------------------------------- 1 | package rmq 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "log/slog" 8 | "time" 9 | 10 | "github.com/danlock/rmq/internal" 11 | amqp "github.com/rabbitmq/amqp091-go" 12 | ) 13 | 14 | // ConsumerArgs contains information needed to declare and consume deliveries from a queue. 15 | type ConsumerArgs struct { 16 | Args 17 | 18 | Queue Queue 19 | QueueBindings []QueueBinding 20 | Consume Consume 21 | Qos Qos 22 | } 23 | 24 | // Queue contains args for amqp.Channel.QueueDeclare 25 | type Queue struct { 26 | Name string // If empty RabbitMQ will generate an unique name 27 | Durable bool 28 | AutoDelete bool 29 | Exclusive bool 30 | NoWait bool 31 | Passive bool // Recommended for Consumer's to passively declare a queue previously declared by a Topology 32 | Args amqp.Table 33 | } 34 | 35 | // QueueBinding contains args for amqp.Channel.QueueBind 36 | type QueueBinding struct { 37 | QueueName string // If empty, RabbitMQ will use the previously generated unique name 38 | ExchangeName string 39 | RoutingKey string 40 | NoWait bool 41 | Args amqp.Table 42 | } 43 | 44 | // Consume contains args for amqp.Channel.Consume 45 | type Consume struct { 46 | AutoAck bool // AutoAck should not be set if you want to actually receive all your messages 47 | Consumer string 48 | Exclusive bool 49 | NoLocal bool 50 | NoWait bool 51 | Args amqp.Table 52 | } 53 | 54 | // Qos contains args for amqp.Channel.Qos 55 | type Qos struct { 56 | PrefetchCount int // Recommended to be set, 2000 is a decent enough default but it heavily depends on your message size. 57 | PrefetchSize int 58 | Global bool 59 | } 60 | 61 | // Consumer enables reliable AMQP Queue consumption. 62 | type Consumer struct { 63 | config ConsumerArgs 64 | conn *Connection 65 | } 66 | 67 | // NewConsumer takes in a ConsumerArgs that describes the AMQP topology of a single queue, 68 | // and returns a rmq.Consumer that can redeclare this topology on any errors during queue consumption. 69 | // This enables robust reconnections even on unreliable networks. 70 | func NewConsumer(rmqConn *Connection, config ConsumerArgs) *Consumer { 71 | config.setDefaults() 72 | return &Consumer{config: config, conn: rmqConn} 73 | } 74 | 75 | // safeDeclareAndConsume safely declares and consumes from an amqp.Queue 76 | // Closes the amqp.Channel on errors. 77 | func (c *Consumer) safeDeclareAndConsume(ctx context.Context) (_ *amqp.Channel, _ <-chan amqp.Delivery, err error) { 78 | logPrefix := fmt.Sprintf("rmq.Consumer.safeDeclareAndConsume for queue %s", c.config.Queue.Name) 79 | timeoutCtx, cancel := context.WithTimeout(ctx, c.config.AMQPTimeout) 80 | defer cancel() 81 | 82 | mqChan, err := c.conn.Channel(timeoutCtx) 83 | if err != nil { 84 | return nil, nil, fmt.Errorf(logPrefix+" failed to get a channel due to err %w", err) 85 | } 86 | // Network calls that don't take a context can block indefinitely. 87 | // Call them in a goroutine so we can timeout if necessary 88 | 89 | respChan := make(chan internal.ChanResp[<-chan amqp.Delivery], 1) 90 | start := time.Now() 91 | go func() { 92 | var r internal.ChanResp[<-chan amqp.Delivery] 93 | r.Val, r.Err = c.declareAndConsume(ctx, mqChan) 94 | ctxDone := timeoutCtx.Err() != nil 95 | // Close the channel on errors or if the context times out, so the amqp channel isn't leaked 96 | if r.Err != nil || ctxDone { 97 | mqChanErr := mqChan.Close() 98 | if mqChanErr != nil && !errors.Is(mqChanErr, amqp.ErrClosed) { 99 | r.Err = errors.Join(r.Err, mqChanErr) 100 | } 101 | } 102 | if ctxDone { 103 | // Log our leaked goroutine's response whenever it finally finishes in case it has useful information. 104 | c.config.Log(ctx, slog.LevelWarn, logPrefix+" completed after it's context finished. It took %s. Err: %+v", time.Since(start), r.Err) 105 | } else { 106 | respChan <- r 107 | } 108 | }() 109 | 110 | select { 111 | case <-timeoutCtx.Done(): 112 | return nil, nil, fmt.Errorf(logPrefix+" unable to complete before context did due to %w", context.Cause(ctx)) 113 | case r := <-respChan: 114 | return mqChan, r.Val, r.Err 115 | } 116 | } 117 | 118 | func (c *Consumer) declareAndConsume(ctx context.Context, mqChan *amqp.Channel) (_ <-chan amqp.Delivery, err error) { 119 | logPrefix := fmt.Sprintf("rmq.Consumer.declareAndConsume for queue (%s)", c.config.Queue.Name) 120 | 121 | if c.config.Qos != (Qos{}) { 122 | err = mqChan.Qos(c.config.Qos.PrefetchCount, c.config.Qos.PrefetchSize, c.config.Qos.Global) 123 | if err != nil { 124 | return nil, fmt.Errorf(logPrefix+" unable to set prefetch due to %w", err) 125 | } else if err = context.Cause(ctx); err != nil { 126 | return nil, fmt.Errorf(logPrefix+" failed to set Qos before context ended due to %w", err) 127 | } 128 | } 129 | 130 | queueDeclare := mqChan.QueueDeclare 131 | if c.config.Queue.Passive { 132 | queueDeclare = mqChan.QueueDeclarePassive 133 | } 134 | _, err = queueDeclare( 135 | c.config.Queue.Name, 136 | c.config.Queue.Durable, 137 | c.config.Queue.AutoDelete, 138 | c.config.Queue.Exclusive, 139 | c.config.Queue.NoWait, 140 | c.config.Queue.Args, 141 | ) 142 | if err != nil { 143 | return nil, fmt.Errorf(logPrefix+" failed to declare queue due to %w", err) 144 | } else if err = context.Cause(ctx); err != nil { 145 | return nil, fmt.Errorf(logPrefix+" failed to declare queue before context ended due to %w", err) 146 | } 147 | 148 | for _, qb := range c.config.QueueBindings { 149 | err = mqChan.QueueBind(qb.QueueName, qb.RoutingKey, qb.ExchangeName, qb.NoWait, qb.Args) 150 | if err != nil { 151 | return nil, fmt.Errorf(logPrefix+" unable to bind queue (%s) to %s due to %w", qb.QueueName, qb.ExchangeName, err) 152 | } else if err = context.Cause(ctx); err != nil { 153 | return nil, fmt.Errorf(logPrefix+" failed to bind queues before context ended due to %w", err) 154 | } 155 | } 156 | 157 | // https://github.com/rabbitmq/amqp091-go/pull/192 Channel.ConsumeWithContext doesn't hold up under scrutiny. The actual network call (ch.call()) doesn't respect the passed in context. 158 | // As of amqp091-go 1.9.0 it doesn't look like we can use ConsumeWithContext to timeout network calls, so we're stuck with this wrapper. 159 | // Now ConsumeWithContext cancels itself when the context is finished, which seems unneccessary since callers can call Cancel, or in danlock/rmq, Close(), themselves. 160 | deliveries, err := mqChan.ConsumeWithContext( 161 | ctx, 162 | c.config.Queue.Name, 163 | c.config.Consume.Consumer, 164 | c.config.Consume.AutoAck, 165 | c.config.Consume.Exclusive, 166 | c.config.Consume.NoLocal, 167 | c.config.Consume.NoWait, 168 | c.config.Consume.Args, 169 | ) 170 | if err != nil { 171 | return nil, fmt.Errorf(logPrefix+" unable to consume due to %w", err) 172 | } 173 | 174 | return deliveries, nil 175 | } 176 | 177 | // Consume uses the rmq.Consumer config to declare and consume from an AMQP queue, forwarding deliveries to it's returned channel. 178 | // On errors Consume reconnects to AMQP, redeclares and resumes consumption and forwarding of deliveries. 179 | // Consume returns an unbuffered channel, and will block on sending to it if no ones listening. 180 | // The returned channel is closed only after the context finishes and the amqp.Channel.Consume's Go channel delivers it's messages. 181 | func (c *Consumer) Consume(ctx context.Context) <-chan amqp.Delivery { 182 | outChan := make(chan amqp.Delivery) 183 | go func() { 184 | internal.Retry(ctx, c.config.Delay, func(delay time.Duration) (time.Duration, bool) { 185 | logPrefix := fmt.Sprintf("rmq.Consumer.Consume for queue (%s)", c.config.Queue.Name) 186 | mqChan, inChan, err := c.safeDeclareAndConsume(ctx) 187 | if err != nil { 188 | c.config.Log(ctx, slog.LevelError, logPrefix+" failed to safeDeclareAndConsume, retrying in %s due to %v", delay.String(), err) 189 | return 0, false 190 | } 191 | 192 | start := time.Now() 193 | c.forwardDeliveries(ctx, mqChan, inChan, outChan) 194 | return time.Since(start), true 195 | }) 196 | close(outChan) 197 | }() 198 | return outChan 199 | } 200 | 201 | // forwardDeliveries forwards from inChan until it closes. If the context finishes it closes the amqp Channel so that the delivery channel will close after sending it's deliveries. 202 | func (c *Consumer) forwardDeliveries(ctx context.Context, mqChan *amqp.Channel, inChan <-chan amqp.Delivery, outChan chan<- amqp.Delivery) { 203 | logPrefix := fmt.Sprintf("rmq.Consumer.forwardDeliveries for queue (%s)", c.config.Queue.Name) 204 | closeNotifier := mqChan.NotifyClose(make(chan *amqp.Error, 6)) 205 | for { 206 | select { 207 | case <-ctx.Done(): 208 | if err := mqChan.Close(); err != nil && !errors.Is(err, amqp.ErrClosed) { 209 | c.config.Log(ctx, slog.LevelError, logPrefix+" failed to Close it's AMQP channel due to %v", err) 210 | // Typically we exit processDeliveries by waiting for inChan to close, but if we can't close even close the AMQP channel then abandon ship 211 | return 212 | } 213 | case err := <-closeNotifier: 214 | if err != nil { 215 | c.config.Log(ctx, slog.LevelError, logPrefix+" got an AMQP Channel Close error %+v", err) 216 | } 217 | case msg, ok := <-inChan: 218 | if !ok { 219 | c.config.Log(ctx, slog.LevelDebug, logPrefix+" amqp.Channel.ConsumeWithContext channel closed") 220 | return 221 | } 222 | // If the client never listens to outChan, this blocks forever 223 | // Other options include using select with a default and dropping the message if the client doesn't listen, dropping the message after a timeout, 224 | // or buffering messages and sending them again later. Of course the buffer could grow forever in that case without listeners. 225 | // The only thing blocked would be the rmq.Consumer.Consume goroutine listening for reconnects and logging errors, which seem unnecessary without a listener anyway. 226 | // Alls well since we don't lock up the entire amqp.Connection like streadway/amqp with Notify* channels... 227 | outChan <- msg 228 | } 229 | } 230 | } 231 | 232 | // ConsumeConcurrently simply runs the provided deliveryProcessor on each delivery from Consume in a new goroutine. 233 | // maxGoroutines limits the amounts of goroutines spawned and defaults to 500. 234 | // Qos.PrefetchCount can also limit goroutines spawned if deliveryProcessor properly Acks messages. 235 | // Blocks until the context is finished and the Consume channel closes. 236 | func (c *Consumer) ConsumeConcurrently(ctx context.Context, maxGoroutines uint64, deliveryProcessor func(ctx context.Context, msg amqp.Delivery)) { 237 | if maxGoroutines == 0 { 238 | maxGoroutines = 500 239 | } 240 | // We use a simple semaphore here and a new goroutine each time. 241 | // It may be more efficient to use a goroutine pool for small amounts of work, but a concerned caller can probably do it better themselves. 242 | semaphore := make(chan struct{}, maxGoroutines) 243 | deliverAndReleaseSemaphore := func(msg amqp.Delivery) { deliveryProcessor(ctx, msg); <-semaphore } 244 | for msg := range c.Consume(ctx) { 245 | semaphore <- struct{}{} 246 | go deliverAndReleaseSemaphore(msg) 247 | } 248 | } 249 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /publisher.go: -------------------------------------------------------------------------------- 1 | package rmq 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "log/slog" 8 | "time" 9 | 10 | "github.com/danlock/rmq/internal" 11 | amqp "github.com/rabbitmq/amqp091-go" 12 | ) 13 | 14 | type PublisherArgs struct { 15 | Args 16 | // NotifyReturn will receive amqp.Return's from any amqp.Channel this rmq.Publisher sends on. 17 | // Recommended to use a buffered channel. Closed after the publisher's context is done. 18 | NotifyReturn chan<- amqp.Return 19 | // LogReturns without their amqp.Return.Body using Args.Log when true 20 | LogReturns bool 21 | 22 | // DontConfirm means the Publisher's amqp.Channel won't be in Confirm mode. Methods except for Publish will throw an error. 23 | DontConfirm bool 24 | } 25 | 26 | type Publisher struct { 27 | ctx context.Context 28 | config PublisherArgs 29 | in chan *Publishing 30 | } 31 | 32 | // NewPublisher creates a rmq.Publisher that will publish messages to AMQP on a single amqp.Channel at a time. 33 | // On error it reconnects via rmq.Connection. Shuts down when it's context is finished. 34 | func NewPublisher(ctx context.Context, rmqConn *Connection, config PublisherArgs) *Publisher { 35 | if ctx == nil || rmqConn == nil { 36 | panic("rmq.NewPublisher called with nil ctx or rmqConn") 37 | } 38 | config.setDefaults() 39 | 40 | pub := &Publisher{ 41 | ctx: ctx, 42 | config: config, 43 | in: make(chan *Publishing), 44 | } 45 | 46 | go pub.connect(rmqConn) 47 | return pub 48 | } 49 | 50 | // connect grabs an amqp.Channel from rmq.Connection. It does so repeatedly on any error until it's context finishes. 51 | func (p *Publisher) connect(rmqConn *Connection) { 52 | internal.Retry(p.ctx, p.config.Delay, func(delay time.Duration) (time.Duration, bool) { 53 | logPrefix := "rmq.Publisher.connect" 54 | mqChan, err := rmqConn.Channel(p.ctx) 55 | if err != nil { 56 | p.config.Log(p.ctx, slog.LevelError, logPrefix+" failed to get amqp.Channel, retrying in %s due to %+v", delay.String(), err) 57 | return 0, false 58 | } 59 | if !p.config.DontConfirm { 60 | if err := mqChan.Confirm(false); err != nil { 61 | p.config.Log(p.ctx, slog.LevelError, logPrefix+" failed to put amqp.Channel in confirm mode, retrying in %s due to %+v", delay.String(), err) 62 | return 0, false 63 | } 64 | } 65 | 66 | start := time.Now() 67 | p.handleReturns(mqChan) 68 | p.listen(mqChan) 69 | return time.Since(start), true 70 | }) 71 | } 72 | 73 | const dropReturnsAfter = 10 * time.Millisecond 74 | 75 | // handleReturns echos the amqp.Channel's Return's until it closes 76 | func (p *Publisher) handleReturns(mqChan *amqp.Channel) { 77 | logPrefix := "rmq.Publisher.handleReturns" 78 | if p.config.NotifyReturn == nil && !p.config.LogReturns { 79 | return 80 | } 81 | notifyReturns := mqChan.NotifyReturn(make(chan amqp.Return)) 82 | go func() { 83 | dropTimer := time.NewTimer(0) 84 | <-dropTimer.C 85 | for r := range notifyReturns { 86 | if p.config.LogReturns { 87 | // A Body can be arbitrarily large and/or contain sensitve info. Don't log it by default. 88 | rBody := r.Body 89 | r.Body = nil 90 | p.config.Log(p.ctx, slog.LevelWarn, logPrefix+" got %+v", r) 91 | r.Body = rBody 92 | } 93 | if p.config.NotifyReturn == nil { 94 | continue 95 | } 96 | 97 | dropTimer.Reset(dropReturnsAfter) 98 | // Try not to repeat streadway/amqp's mistake of deadlocking if a client isn't listening to their Notify* channel. 99 | // (https://github.com/rabbitmq/amqp091-go/issues/18) 100 | // If they aren't listening to p.config.NotifyReturn, just drop the amqp.Return instead of deadlocking and leaking goroutines 101 | select { 102 | case p.config.NotifyReturn <- r: 103 | // Why is reusing a timer so bloody complicated... It's almost worth the timer leak just to reduce complexity 104 | if !dropTimer.Stop() { 105 | <-dropTimer.C 106 | } 107 | case <-dropTimer.C: 108 | } 109 | } 110 | // Close when the context is done, since we wont be sending anymore returns 111 | if p.config.NotifyReturn != nil && p.ctx.Err() != nil { 112 | close(p.config.NotifyReturn) 113 | } 114 | }() 115 | } 116 | 117 | // listen sends publishes on a amqp.Channel until it's closed. 118 | func (p *Publisher) listen(mqChan *amqp.Channel) { 119 | logPrefix := "rmq.Publisher.listen" 120 | finishedPublishing := make(chan struct{}, 1) 121 | ctx, cancel := context.WithCancel(p.ctx) 122 | defer cancel() 123 | // Handle publishes in a separate goroutine so a slow publish won't lock up listen() 124 | go func() { 125 | for { 126 | select { 127 | case <-ctx.Done(): 128 | close(finishedPublishing) 129 | return 130 | case pub := <-p.in: 131 | pub.publish(mqChan) 132 | } 133 | } 134 | }() 135 | 136 | notifyClose := mqChan.NotifyClose(make(chan *amqp.Error, 2)) 137 | for { 138 | select { 139 | case <-p.ctx.Done(): 140 | // Wait for publishing to finish since closing the channel in the middle of another channel request 141 | // tends to kill the entire connection with a "504 CHANNEL ERROR expected 'channel.open'" 142 | <-finishedPublishing 143 | if err := mqChan.Close(); err != nil && !errors.Is(err, amqp.ErrClosed) { 144 | p.config.Log(p.ctx, slog.LevelError, logPrefix+" got an error while closing channel %v", err) 145 | return 146 | } 147 | case err, ok := <-notifyClose: 148 | if !ok { 149 | return 150 | } else if err != nil { 151 | p.config.Log(p.ctx, slog.LevelError, logPrefix+" got an amqp.Channel close err %v", err) 152 | } 153 | } 154 | } 155 | } 156 | 157 | type Publishing struct { 158 | amqp.Publishing 159 | Exchange string 160 | RoutingKey string 161 | Mandatory bool 162 | Immediate bool 163 | 164 | // req is internal and private, which means it can't be set by callers. 165 | // This means it has the nice side effect of forcing callers to set struct fields when instantiating Publishing 166 | req internal.ChanReq[*amqp.DeferredConfirmation] 167 | } 168 | 169 | func (p *Publishing) publish(mqChan *amqp.Channel) { 170 | var resp internal.ChanResp[*amqp.DeferredConfirmation] 171 | resp.Val, resp.Err = mqChan.PublishWithDeferredConfirmWithContext( 172 | p.req.Ctx, p.Exchange, p.RoutingKey, p.Mandatory, p.Immediate, p.Publishing) 173 | p.req.RespChan <- resp 174 | } 175 | 176 | // Publish send a Publishing on rmq.Publisher's current amqp.Channel. 177 | // Returns amqp.DefferedConfirmation's only if the rmq.Publisher has Confirm set. 178 | // If an error is returned, rmq.Publisher will grab another amqp.Channel from rmq.Connection, which itself will redial AMQP if necessary. 179 | // This means simply retrying Publish on errors will send Publishing's even on flaky connections. 180 | func (p *Publisher) Publish(ctx context.Context, pub Publishing) (*amqp.DeferredConfirmation, error) { 181 | pub.req.Ctx = ctx 182 | pub.req.RespChan = make(chan internal.ChanResp[*amqp.DeferredConfirmation], 1) 183 | select { 184 | case <-ctx.Done(): 185 | return nil, fmt.Errorf("rmq.Publisher.Publish context done before publish sent %w", context.Cause(ctx)) 186 | case <-p.ctx.Done(): 187 | return nil, fmt.Errorf("rmq.Publisher context done before publish sent %w", context.Cause(p.ctx)) 188 | case p.in <- &pub: 189 | } 190 | 191 | select { 192 | case <-ctx.Done(): 193 | return nil, fmt.Errorf("rmq.Publisher.Publish context done before publish completed %w", context.Cause(ctx)) 194 | case <-p.ctx.Done(): 195 | return nil, fmt.Errorf("rmq.Publisher context done before publish completed %w", context.Cause(p.ctx)) 196 | case r := <-pub.req.RespChan: 197 | return r.Val, r.Err 198 | } 199 | } 200 | 201 | // PublishUntilConfirmed calls Publish and waits for Publishing to be confirmed. 202 | // It republishes if a message isn't confirmed after confirmTimeout, or if Publish returns an error. 203 | // Returns *amqp.DeferredConfirmation so the caller can check if it's Acked(). 204 | // confirmTimeout defaults to 1 minute. Recommended to call with context.WithTimeout. 205 | func (p *Publisher) PublishUntilConfirmed(ctx context.Context, confirmTimeout time.Duration, pub Publishing) (*amqp.DeferredConfirmation, error) { 206 | logPrefix := "rmq.Publisher.PublishUntilConfirmed" 207 | 208 | if p.config.DontConfirm { 209 | return nil, fmt.Errorf(logPrefix + " called on a rmq.Publisher that's not in Confirm mode") 210 | } 211 | 212 | if confirmTimeout <= 0 { 213 | confirmTimeout = 15 * time.Second 214 | } 215 | 216 | var pubDelay time.Duration 217 | attempt := 0 218 | errs := make([]error, 0) 219 | 220 | for { 221 | defConf, err := p.Publish(ctx, pub) 222 | if err != nil { 223 | pubDelay = p.config.Delay(attempt) 224 | attempt++ 225 | errs = append(errs, err) 226 | select { 227 | case <-ctx.Done(): 228 | err = fmt.Errorf(logPrefix+" context done before the publish was sent %w", context.Cause(ctx)) 229 | return defConf, errors.Join(append(errs, err)...) 230 | case <-time.After(pubDelay): 231 | continue 232 | } 233 | } 234 | attempt = 0 235 | 236 | confirmTimeout := time.NewTimer(confirmTimeout) 237 | defer confirmTimeout.Stop() 238 | 239 | select { 240 | case <-confirmTimeout.C: 241 | errs = append(errs, errors.New(logPrefix+" timed out waiting for confirm, republishing")) 242 | continue 243 | case <-ctx.Done(): 244 | err = fmt.Errorf("rmq.Publisher.PublishUntilConfirmed context done before the publish was confirmed %w", context.Cause(ctx)) 245 | return defConf, errors.Join(append(errs, err)...) 246 | case <-defConf.Done(): 247 | return defConf, nil 248 | } 249 | } 250 | } 251 | 252 | // PublishUntilAcked is like PublishUntilConfirmed, but it also republishes nacks. User discretion is advised. 253 | // 254 | // Nacks can happen for a variety of reasons, ranging from user error (mistyped exchange) to RabbitMQ internal errors. 255 | // 256 | // PublishUntilAcked will republish a Mandatory Publishing with a nonexistent exchange forever (until the exchange exists), as one example. 257 | // RabbitMQ acks Publishing's so monitor the NotifyReturn chan to ensure your Publishing's are being delivered. 258 | // 259 | // PublishUntilAcked is intended for ensuring a Publishing with a known destination queue will get acked despite flaky connections or temporary RabbitMQ node failures. 260 | // Recommended to call with context.WithTimeout. 261 | func (p *Publisher) PublishUntilAcked(ctx context.Context, confirmTimeout time.Duration, pub Publishing) error { 262 | logPrefix := "rmq.Publisher.PublishUntilAcked" 263 | nacks := 0 264 | for { 265 | defConf, err := p.PublishUntilConfirmed(ctx, confirmTimeout, pub) 266 | if err != nil { 267 | if nacks > 0 { 268 | return fmt.Errorf(logPrefix+" resent nacked Publishings %d time(s) and %w", nacks, err) 269 | } 270 | return err 271 | } 272 | 273 | if defConf.Acked() { 274 | return nil 275 | } 276 | 277 | nacks++ 278 | } 279 | } 280 | 281 | // PublishBatchUntilAcked Publishes all of your Publishings at once, and then wait's for the DeferredConfirmation to be Acked, 282 | // resending if it's been longer than confirmTimeout or if they've been nacked. 283 | // confirmTimeout defaults to 1 minute. Recommended to call with context.WithTimeout. 284 | func (p *Publisher) PublishBatchUntilAcked(ctx context.Context, confirmTimeout time.Duration, pubs ...Publishing) error { 285 | logPrefix := "rmq.Publisher.PublishBatchUntilAcked" 286 | 287 | if len(pubs) == 0 { 288 | return nil 289 | } 290 | if p.config.DontConfirm { 291 | return fmt.Errorf(logPrefix + " called on a rmq.Publisher that's not in Confirm mode") 292 | } 293 | 294 | if confirmTimeout == 0 { 295 | confirmTimeout = time.Minute 296 | } 297 | 298 | errs := make([]error, 0) 299 | pendingPubs := make([]*amqp.DeferredConfirmation, len(pubs)) 300 | ackedPubs := make([]bool, len(pubs)) 301 | 302 | for { 303 | select { 304 | case <-p.ctx.Done(): 305 | err := fmt.Errorf(logPrefix+"'s Publisher timed out because %w", context.Cause(p.ctx)) 306 | return errors.Join(append(errs, err)...) 307 | case <-ctx.Done(): 308 | err := fmt.Errorf(logPrefix+" timed out because %w", context.Cause(ctx)) 309 | return errors.Join(append(errs, err)...) 310 | default: 311 | } 312 | 313 | err := p.publishBatch(ctx, confirmTimeout, pubs, pendingPubs, ackedPubs, errs) 314 | if err == nil { 315 | return nil 316 | } 317 | clear(pendingPubs) 318 | } 319 | } 320 | 321 | // publishBatch publishes a slice of pubs once, waiting for them all to get acked. 322 | // republishes on failure, returns after they've confirmed. 323 | // blocks until context ends or confirmTimeout 324 | func (p *Publisher) publishBatch( 325 | ctx context.Context, 326 | confirmTimeout time.Duration, 327 | pubs []Publishing, 328 | pendingPubs []*amqp.DeferredConfirmation, 329 | ackedPubs []bool, 330 | errs []error, 331 | ) (err error) { 332 | logPrefix := "rmq.Publisher.publishBatch" 333 | var delay time.Duration 334 | attempt := 0 335 | published := 0 336 | remaining := 0 337 | for _, acked := range ackedPubs { 338 | if !acked { 339 | remaining++ 340 | } 341 | } 342 | 343 | for published != remaining { 344 | for i, pub := range pubs { 345 | // Skip if it's been previously acked or published 346 | if ackedPubs[i] || pendingPubs[i] != nil { 347 | continue 348 | } 349 | 350 | pendingPubs[i], err = p.Publish(ctx, pub) 351 | if err != nil { 352 | errs = append(errs, err) 353 | delay = p.config.Delay(attempt) 354 | attempt++ 355 | select { 356 | case <-p.ctx.Done(): 357 | return fmt.Errorf(logPrefix+"'s Publisher timed out because %w", context.Cause(p.ctx)) 358 | case <-ctx.Done(): 359 | return fmt.Errorf(logPrefix+" timed out because %w", context.Cause(ctx)) 360 | case <-time.After(delay): 361 | } 362 | } else { 363 | published++ 364 | attempt = 0 365 | } 366 | } 367 | } 368 | 369 | confirmTimer := time.After(confirmTimeout) 370 | confirmed := 0 371 | for confirmed != remaining { 372 | for i, pub := range pendingPubs { 373 | // Skip if it's been previously confirmed 374 | if ackedPubs[i] || pendingPubs[i] == nil { 375 | continue 376 | } 377 | 378 | select { 379 | case <-p.ctx.Done(): 380 | return fmt.Errorf(logPrefix+"'s Publisher timed out because %w", context.Cause(p.ctx)) 381 | case <-ctx.Done(): 382 | return fmt.Errorf(logPrefix+" timed out because %w", context.Cause(ctx)) 383 | case <-confirmTimer: 384 | return fmt.Errorf(logPrefix + " timed out waiting on confirms") 385 | case <-pub.Done(): 386 | if pub.Acked() { 387 | ackedPubs[i] = true 388 | } 389 | confirmed++ 390 | pendingPubs[i] = nil 391 | default: 392 | } 393 | } 394 | } 395 | 396 | return nil 397 | } 398 | --------------------------------------------------------------------------------