diff --git a/README.md b/README.md index e7ff1b1..f40ba50 100644 --- a/README.md +++ b/README.md @@ -152,7 +152,7 @@ func main() { // Collect the ids of any messages that failed. They won't be deleted, // so they return to the queue once their visibility timeout expires. - var failed []interface{} + var failed []any for _, msg := range msgs { // Assert the type of message to get the body or any other attributes diff --git a/v2/README.md b/v2/README.md new file mode 100644 index 0000000..21621df --- /dev/null +++ b/v2/README.md @@ -0,0 +1,64 @@ +# Formigo - distributed SQS worker pools. + +> This library relies heavily on [`github.com/alitto/pond/v2`](https://github.com/alitto/pond), a very flexible and well-tested worker pool library. + +Formigo is a fast, reliable worker pool consumer for SQS. + +Basic features of Pond: + +- automatic scaling to available resources/limits based on incoming queue pressure. +- fire & forget queue submission +- fire & **wait for a response** submission to a queue (e.g. for dependent jobs, or deferred behaviour). +- clean, graceful exit when shutting down. + +## Usage + +See `cmd/echo/main.go` for a complete example of how to use v2. + +### Consumers + +Your consumer should adhere to the following: + +```go +import formigo "github.com/Pod-Point/go-queue-worker/v2" + +func (context.Context, formigo.Message) error +``` + +A `Decode` method is provided on formigo.Message to help ensure proper decoding with JSON SQS messages: + +```go +type Message struct { + Foo string `json:"foo"` +} + +func (ctx context.Context, msg formigo.Message) error { + var body Message + if err := msg.Decode(&body); err != nil { + return err + } + // pass to your internal consumer, etc. +} +``` + +### Error Reporting + +If you need errors to be sent to Sentry, structured logging, etc, feel free to use `WithReporter` when setting up a manager. + +```go +formigo.WithReporter(func(err error) { + sentry.CaptureException(err) +}) +``` + +### Other Options + +All other options should be fairly self-explanatory, and have godocs and defaults. + +## Shutting Down + +Ensure the context passed to `manager.Run(ctx)` is cancelable, preferrably via `signal.NotifyContext`. + +If it isn't, the program will have no exit condition until fully terminated by the operating system. + +See the example consumer to see how this is done. \ No newline at end of file diff --git a/v2/cmd/echo/main.go b/v2/cmd/echo/main.go new file mode 100644 index 0000000..2a66bd8 --- /dev/null +++ b/v2/cmd/echo/main.go @@ -0,0 +1,82 @@ +// Package main is a test consumer for a localstack queue. +// This is not meant to be a production config, though it is set up exactly as you would use it there. +package main + +import ( + "context" + "errors" + "fmt" + "log" + "log/slog" + "os" + "os/signal" + "syscall" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/service/sqs" + + formigo "github.com/Pod-Point/go-queue-worker/v2" +) + +type Message struct { + Level string `json:"level"` + Message string `json:"message"` +} + +func main() { + ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGILL, syscall.SIGHUP) + defer cancel() + + cfg, err := config.LoadDefaultConfig(ctx) + if err != nil { + log.Printf("unable to load config: %v\n", err) + os.Exit(2) + } + + client := formigo.NewSQSClient(sqs.NewFromConfig(cfg), &sqs.ReceiveMessageInput{ + QueueUrl: aws.String("http://sqs.eu-west-1.localhost.localstack.cloud:4566/000000000000/formigo-echo-testing"), + MaxNumberOfMessages: 5, + VisibilityTimeout: 6, + WaitTimeSeconds: 10, + }) + + manager := formigo.NewManager(client, + formigo.WithDeadline(time.Second*4), + formigo.WithFetchConcurrency(2), + formigo.WithFetchDelay(time.Second*1), + formigo.WithWorkerConcurrency(3), + formigo.WithReporter(func(err error) { + slog.Error(err.Error(), slog.String("type", fmt.Sprintf("%T", err))) + }), + formigo.WithConsumer(func(ctx context.Context, msg formigo.Message) error { + var body Message + if err := msg.Decode(&body); err != nil { + return err + } + + switch body.Level { + case "error": + return errors.New(body.Message) + case "info": + slog.InfoContext(ctx, body.Message) + case "warn", "warning": + slog.WarnContext(ctx, body.Message) + case "debug": + slog.DebugContext(ctx, body.Message) + } + + return nil + }), + formigo.WithLogger(slog.New( + slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{ + Level: slog.LevelInfo, + }), + )), + ) + + if err := manager.Run(ctx); err != nil { + log.Fatal(err) + } +} diff --git a/v2/cmd/producer/main.go b/v2/cmd/producer/main.go new file mode 100644 index 0000000..0721823 --- /dev/null +++ b/v2/cmd/producer/main.go @@ -0,0 +1,57 @@ +// Package main is a test producer for a localstack queue. +// This is not meant to be a production config, though it is set up exactly as you would use it there. +package main + +import ( + "context" + "encoding/json" + "fmt" + "log" + "os" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/service/sqs" +) + +func main() { + ctx := context.Background() + + cfg, err := config.LoadDefaultConfig(ctx) + if err != nil { + log.Printf("unable to load config: %v\n", err) + os.Exit(2) + } + + if len(os.Args[1:]) != 2 { + fmt.Println("USAGE: ./echo [level] \"[message]\"") + } + + level := os.Args[1] + body := os.Args[2] + + b, err := json.Marshal(struct { + Level string `json:"level"` + Body string `json:"message"` + }{ + Level: level, + Body: body, + }) + if err != nil { + log.Fatal(err) + } + + client := sqs.NewFromConfig(cfg) + + // duplicate for testing + for range 10 { + output, err := client.SendMessage(ctx, &sqs.SendMessageInput{ + MessageBody: aws.String(string(b)), + QueueUrl: aws.String("http://sqs.eu-west-1.localhost.localstack.cloud:4566/000000000000/formigo-echo-testing"), + }) + if err != nil { + log.Fatal(err) + } + fmt.Printf("message id: %s\n", *output.MessageId) + } +} diff --git a/v2/go.mod b/v2/go.mod new file mode 100644 index 0000000..f35f577 --- /dev/null +++ b/v2/go.mod @@ -0,0 +1,29 @@ +module github.com/Pod-Point/go-queue-worker/v2 + +go 1.26 + +require ( + github.com/alitto/pond/v2 v2.7.1 + github.com/aws/aws-sdk-go-v2 v1.47.1 + github.com/aws/aws-sdk-go-v2/config v1.33.6 + github.com/aws/aws-sdk-go-v2/service/sqs v1.34.5 + github.com/stretchr/testify v1.9.0 +) + +require ( + github.com/aws/aws-sdk-go-v2/credentials v1.20.6 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.20.1 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.4 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.4 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.5.4 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.19 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.14.4 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.10.1 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.38.1 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.43.1 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.51.1 // indirect + github.com/aws/smithy-go v1.28.1 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/v2/go.sum b/v2/go.sum new file mode 100644 index 0000000..a940922 --- /dev/null +++ b/v2/go.sum @@ -0,0 +1,42 @@ +github.com/alitto/pond/v2 v2.7.1 h1:QxMbcfjcVTa0pyxX5Ib1226mM8u8D7gKUVkCUU4DYIw= +github.com/alitto/pond/v2 v2.7.1/go.mod h1:xkjYEgQ05RSpWdfSd1nM3OVv7TBhLdy7rMp3+2Nq+yE= +github.com/aws/aws-sdk-go-v2 v1.47.1 h1:uOIZnp4PK3ZhKI0dNrJrhTEsLxbpXHTAJlwoS1pvAtw= +github.com/aws/aws-sdk-go-v2 v1.47.1/go.mod h1:bttEH6JqnUL8LepvDVfdrds/fZ5bCIxzpe3abyUrhDU= +github.com/aws/aws-sdk-go-v2/config v1.33.6 h1:MBjkSTLczek/UgiK+EYPIoRTqE7gP8vtW3OFbFo7Nug= +github.com/aws/aws-sdk-go-v2/config v1.33.6/go.mod h1:grRAFzdAZJrwcbasJRg2MPvIrVjtlfXllHssN6+E1JE= +github.com/aws/aws-sdk-go-v2/credentials v1.20.6 h1:NpAFXCU7NzXNkdGK3zQTtsRJ+3v9tZQV0xcdRw8uBdw= +github.com/aws/aws-sdk-go-v2/credentials v1.20.6/go.mod h1:mcZCoiPnyMvP8VMNbygNX5lLqSlkYJIMPODylQMurOk= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.20.1 h1:8gALAAmacnIXh+z6VkdDanv4/IkG5APdg4DZLDTmLog= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.20.1/go.mod h1:Z7IJhJU+poOdJjUR2wpyY21ossQ1XS/R3Lk9Msq5kM4= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.4 h1:CLq4+8UHCI+ZZYl/EuJxXovaIVN2xeeT8JV+dsApQ5E= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.4/go.mod h1:Wv4q5sAM04xAMkoOedxLx2inVf6K5FdxYp+A61L+q/0= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.4 h1:dD4MR81I7YkpEBRk6UP9rocC2QnT3qVuXwzlYTtfGEs= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.4/go.mod h1:EcXV1kAFd5XwSkDHlj94gnF3q5CkJyYiIJfH8N0VmrE= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.5.4 h1:7Wo47d/xn/7KttCSBd8EGYeZ7ULRFRkUHr6vkZPBzVQ= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.5.4/go.mod h1:tDB2IVC1xC3vX8o+6uRlzhTxP3g1b77CZXFX/oD2FnQ= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.19 h1:bAdDl/HkGCcGPoe25ToSHEw23VIxt6CT5fLcg111BKg= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.19/go.mod h1:KaUzbLxv4CeSxh6ZCl9B4m7CuFenS8kUEaDs+f/DQr4= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.14.4 h1:29SvnfGhXjTl8ONxFwbj2rs6lbhiFXD2CgFQmbT/bXY= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.14.4/go.mod h1:wm04I5DMuNVvZHFe/dHnUxincvNbbK7AiNBbYsQivek= +github.com/aws/aws-sdk-go-v2/service/signin v1.10.1 h1:DzCCWLzcIRQ77F3DEUljud7bEjTgFOIKXP52NmVRyhU= +github.com/aws/aws-sdk-go-v2/service/signin v1.10.1/go.mod h1:xpo/geVldu8payT375WekctUzopG/hBU7miiqItMUlw= +github.com/aws/aws-sdk-go-v2/service/sqs v1.34.5 h1:HYyVDOC2/PIg+3oBX1q0wtDU5kONki6lrgIG0afrBkY= +github.com/aws/aws-sdk-go-v2/service/sqs v1.34.5/go.mod h1:7idt3XszF6sE9WPS1GqZRiDJOxw4oPtlRBXodWnCGjU= +github.com/aws/aws-sdk-go-v2/service/sso v1.38.1 h1:Umtl/0YZhng4xndfW3lKJrYYP7NLEjI6bGXVomwLcs0= +github.com/aws/aws-sdk-go-v2/service/sso v1.38.1/go.mod h1:rRD/dnm7q0HYE/I5TMaPgkWyyUGLcwuxHLABsLnQ3e0= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.43.1 h1:orIWdNiLgzrhu/11RcPPKO/SBzUUymbUQuZbSPImghg= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.43.1/go.mod h1:skwM/xsbR/1ReUTesv9BhpJp1VjajR7DWQnuVLwiXsQ= +github.com/aws/aws-sdk-go-v2/service/sts v1.51.1 h1:0HOqZXRvMytH6bFHVIc0oJX07sZjfhz0zXtjs6gdE8s= +github.com/aws/aws-sdk-go-v2/service/sts v1.51.1/go.mod h1:26zA0GhDrLo+yiLI2yXWxqB1PdsShfLikoI7GOEgugM= +github.com/aws/smithy-go v1.28.1 h1:R/nXH00c8qcfCzQVELtRw+eLQWtzv+VAIEFJ1/xxXlQ= +github.com/aws/smithy-go v1.28.1/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/v2/manager.go b/v2/manager.go new file mode 100644 index 0000000..e4a742d --- /dev/null +++ b/v2/manager.go @@ -0,0 +1,252 @@ +// Package formigo exposes an extremely simple way to run SQS messages inside a worker pool with capped concurrency. +// It also handles graceful shutdown, and ensures messages in-flight are completed on context cancellation. +package formigo + +import ( + "context" + "errors" + "log/slog" + "time" + + "github.com/alitto/pond/v2" +) + +var ( + ErrConsumerMissing = errors.New("formigo/v2: consumer is not set. Use .WithConsumer()") +) + +// Fetcher condenses RetrieveMessage into the smallest possible interface. +type Fetcher interface { + Fetch(context.Context) ([]Message, error) +} + +// Deleter condenses DeleteMessage into the smallest possible interface. +type Deleter interface { + Delete(context.Context, Message) error +} + +// Client is a thin wrapper around an SQS client that can fetch and delete messages. +type Client interface { + Fetcher + Deleter +} + +// Manager sets up a set of workers and processes an queue via them. +type Manager struct { + client Client + pool pond.Pool + deadline time.Duration + fetchDelay time.Duration + fetchConcurrency int + workerConcurrency int + consumer func(context.Context, Message) error + reporter func(error) + logger *slog.Logger +} + +// Submit will push the message onto a consumer via the worker pool. +func (m *Manager) Submit(msg Message) (err error) { + m.pool.SubmitErr(func() error { + ctx := context.Background() + + // make sure this consumer deletes the message from the queue when deferring. + // this one uses the background context with no deadline, and reports its own errors. + defer func() { + // wrapped in a func() to defer scope capture + m.Delete(ctx, msg, err) + }() + + m.Log(ctx, slog.LevelDebug, "processing message", + slog.String("messageID", msg.ID), + slog.String("receiptHandle", msg.ReceiptHandle), + ) + + if m.deadline > 0 { + m.Log(ctx, slog.LevelDebug, "applying deadline", + slog.String("messageID", msg.ID), + slog.Duration("deadline", m.deadline), + ) + + var cancel func() + ctx, cancel = context.WithDeadline(ctx, time.Now().Add(m.deadline)) + defer cancel() + } + + err = m.consumer(ctx, msg) + + if err != nil { + m.Log(ctx, slog.LevelInfo, "message: errored", + slog.String("messageID", msg.ID), + slog.String("error", err.Error()), + ) + } else { + m.Log(ctx, slog.LevelInfo, "message: processed", + slog.String("messageID", msg.ID), + ) + } + + return err + }) + + return nil +} + +// Delete will call the Client's Delete method and remove a message from the queue. +// This may succeed with no error, but still not remove the message (e.g. it was not the most recent fetch of this message). +func (m *Manager) Delete(ctx context.Context, msg Message, err error) { + m.Log(ctx, slog.LevelDebug, "deleting message", + slog.String("messageID", msg.ID), + slog.String("receiptHandle", msg.ReceiptHandle), + ) + + if err != nil { + switch { + case errors.Is(err, context.Canceled): + // in this specific case, we may have done the work but the context passed to us was cancelled. + // we may or may not *have* to delete here. + m.Log(ctx, slog.LevelDebug, "not deleting a context.Canceled message", + slog.String("messageID", msg.ID), + slog.String("error", err.Error()), + ) + return + case errors.Is(err, context.DeadlineExceeded): + // do not delete from the queue if the deadline is exceeded. + // this needs to be reported as well, or it'll cause a problem + // for end users not seeing visibility or results. + m.Log(ctx, slog.LevelDebug, "not deleting a context.DeadlineExceeded message", + slog.String("messageID", msg.ID), + slog.String("error", err.Error()), + ) + m.Report(err) + return + default: + m.Report(err) + } + } + + if err := m.client.Delete(ctx, msg); err != nil { + m.Log(ctx, slog.LevelError, "failed to delete message", + slog.String("messageID", msg.ID), + slog.String("error", err.Error()), + ) + m.Report(err) + } else { + m.Log(ctx, slog.LevelInfo, "message deleted", + slog.String("messageID", msg.ID), + ) + } +} + +// Fetch will run a continuous loop of querying SQS for messages. +func (m *Manager) Fetch(ctx context.Context) error { + for { + select { + case <-ctx.Done(): + m.Log(ctx, slog.LevelDebug, "exiting fetch instance") + return ctx.Err() + default: + m.Log(ctx, slog.LevelDebug, "fetching messages") + messages, err := m.client.Fetch(ctx) + if err != nil { + m.Report(err) // report errors with the initial fetch. + continue + } + + m.Log(ctx, slog.LevelDebug, "fetched messages", slog.Int("count", len(messages))) + + for _, msg := range messages { + if err := m.Submit(msg); err != nil { + m.Report(err) + } + } + + m.Log(ctx, slog.LevelDebug, "waiting before polling", slog.Duration("delay", m.fetchDelay)) + + // allow configuring a delay after fetching messages + // this may help prevent excessive pressure on AWS. + if m.fetchDelay > 0 { + time.Sleep(m.fetchDelay) + } + } + } +} + +// Report will report an error whenever a func for it is set. +func (m *Manager) Report(err error) { + if m.reporter != nil { + m.reporter(err) + } +} + +// Run will start up a worker pool and start processing a queue. +// It must exit via the context, e.g. via signal.NotifyContext or cancel(). +// Using context.Background will never safely exit. +// +// This will error if a consumer is not set. +func (m *Manager) Run(ctx context.Context) error { + if m.consumer == nil { + return ErrConsumerMissing + } + + m.pool = pond.NewPool(m.workerConcurrency, pond.WithContext(ctx)) + + // set up the fetchers to feed the workers and start them + fetcher := pond.NewPool(m.fetchConcurrency, pond.WithContext(ctx)) + for i := range m.fetchConcurrency { + m.Log(ctx, slog.LevelInfo, "starting fetcher", slog.Int("id", i)) + + fetcher.SubmitErr(func() error { + return m.Fetch(ctx) + }) + } + + select { + case <-ctx.Done(): + m.Log(ctx, slog.LevelInfo, "stopping fetchers") + + // stop the fetcher first to be sure it has drained + fetcher.StopAndWait() + + m.Log(ctx, slog.LevelInfo, "stopping workers") + + // stop and drain the worker pool + m.pool.StopAndWait() + return ctx.Err() + } +} + +// Log will send a log record to a configured slog.Logger, if set. +func (m *Manager) Log(ctx context.Context, level slog.Level, msg string, args ...any) { + if m.logger != nil { + m.logger.Log(ctx, level, msg, args...) + } +} + +// NewManager will create a new *Manager with the given client and options. +func NewManager(client Client, opts ...Option) *Manager { + manager := NewDefaultManager(client) + + for _, opt := range opts { + opt.Apply(manager) + } + + manager.Log(context.Background(), slog.LevelInfo, "applied options", + slog.Int("workerConcurrency", manager.workerConcurrency), + slog.Int("fetchConcurrency", manager.fetchConcurrency), + slog.Duration("fetchDelay", manager.fetchDelay), + slog.Duration("deadline", manager.deadline), + ) + + return manager +} + +// NewDefaultManager returns a Manager with all defaults set in the options. +func NewDefaultManager(client Client) *Manager { + return &Manager{ + client: client, + fetchDelay: 0, + fetchConcurrency: 2, + workerConcurrency: 20, + deadline: time.Second * 30, + } +} diff --git a/v2/manager_test.go b/v2/manager_test.go new file mode 100644 index 0000000..aafbf29 --- /dev/null +++ b/v2/manager_test.go @@ -0,0 +1,48 @@ +package formigo + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +type TestClient struct { + // +} + +func (t TestClient) Fetch(ctx context.Context) ([]Message, error) { + return nil, nil +} + +func (t TestClient) Delete(ctx context.Context, message Message) error { + return nil +} + +var _ Client = (*TestClient)(nil) + +// testContext is a safety measure for testing Run in case there is never an exit condition. +func testContext() (context.Context, context.CancelFunc) { + return context.WithDeadline(context.Background(), time.Now().Add(time.Second*4)) +} + +func TestNewDefaultManager_VerifiesDefaults(t *testing.T) { + manager := NewDefaultManager(TestClient{}) + + assert.Nil(t, manager.consumer) + assert.Nil(t, manager.reporter) + + assert.Equal(t, 2, manager.fetchConcurrency) + assert.Equal(t, 20, manager.workerConcurrency) + assert.Equal(t, time.Millisecond*500, manager.fetchDelay) + assert.Equal(t, time.Second*30, manager.deadline) +} + +func TestManager_Run_FailsWithoutConsumer(t *testing.T) { + manager := NewDefaultManager(TestClient{}) + ctx, cancel := testContext() + defer cancel() + + assert.ErrorIs(t, manager.Run(ctx), ErrConsumerMissing) +} diff --git a/v2/message.go b/v2/message.go new file mode 100644 index 0000000..e38d92e --- /dev/null +++ b/v2/message.go @@ -0,0 +1,18 @@ +package formigo + +import ( + "encoding/json" +) + +type Message struct { + ID string + Attributes map[string]string + Body string + MessageAttributes map[string]any + ReceiptHandle string +} + +// Decode attempts to decode a message's Body using json.Unmarshal. +func (m Message) Decode(value any) error { + return json.Unmarshal([]byte(m.Body), &value) +} diff --git a/v2/options.go b/v2/options.go new file mode 100644 index 0000000..7f890c1 --- /dev/null +++ b/v2/options.go @@ -0,0 +1,77 @@ +package formigo + +import ( + "context" + "log/slog" + "time" +) + +type Option interface { + Apply(*Manager) +} + +type OptionFunc func(*Manager) + +func (f OptionFunc) Apply(manager *Manager) { + f(manager) +} + +// WithDeadline sets the context deadline used for consumer processes when they are given a message. +// Ensure this is lower than VisibilityTimeout in SQS, or you may process the same message twice. +// +// Default: 30s +func WithDeadline(deadline time.Duration) OptionFunc { + return func(manager *Manager) { + manager.deadline = deadline + } +} + +// WithWorkerConcurrency will set the number of workers used to run a consumer process. +// Set this high enough so that you can process a reasonably high number of messages, +// but not so high that you throttle/cpu-limit the other workers. +// +// Default: 20 +func WithWorkerConcurrency(concurrency int) OptionFunc { + return func(manager *Manager) { + manager.workerConcurrency = concurrency + } +} + +// WithFetchConcurrency will set the number of workers that will be continuously fetching messages from the given queue. +// Do not set this significantly higher than needed (3-4 is sufficient) or you may grow the queue backlog. +// +// Default: 2 +func WithFetchConcurrency(concurrency int) OptionFunc { + return func(manager *Manager) { + manager.fetchConcurrency = concurrency + } +} + +// WithFetchDelay will set the amount of time to sleep between fetching new messages. +// Default: 500ms +func WithFetchDelay(delay time.Duration) OptionFunc { + return func(manager *Manager) { + manager.fetchDelay = delay + } +} + +// WithReporter will set the error reporting function (e.g. log +func WithReporter(f func(error)) OptionFunc { + return func(manager *Manager) { + manager.reporter = f + } +} + +// WithConsumer sets the consumer to use when processing messages. +// Required; panics if not set since this is a dev issue. +func WithConsumer(f func(context.Context, Message) error) OptionFunc { + return func(manager *Manager) { + manager.consumer = f + } +} + +func WithLogger(logger *slog.Logger) OptionFunc { + return func(manager *Manager) { + manager.logger = logger + } +} diff --git a/v2/sqs.go b/v2/sqs.go new file mode 100644 index 0000000..39c417c --- /dev/null +++ b/v2/sqs.go @@ -0,0 +1,111 @@ +package formigo + +import ( + "context" + "errors" + + "github.com/aws/aws-sdk-go-v2/service/sqs" + "github.com/aws/aws-sdk-go-v2/service/sqs/types" +) + +var ( + ErrClientNil = errors.New("formigo/v2: sqs client is nil") + ErrInputNil = errors.New("formigo/v2: sqs ReceiveMessage input is nil") + ErrReceiptHandleEmpty = errors.New("formigo/v2: message ReceiptHandle is empty when deleting") +) + +// SQSClient is an implementation of Client, and is a thin wrapper around AWS SQS. +type SQSClient struct { + client *sqs.Client + input *sqs.ReceiveMessageInput +} + +// Fetch will attempt to fetch SQS messages from the configured queue, and convert them to Message. +func (s SQSClient) Fetch(ctx context.Context) ([]Message, error) { + if s.client == nil { + return nil, ErrClientNil + } + + if s.input == nil { + return nil, ErrInputNil + } + + output, err := s.client.ReceiveMessage(ctx, s.input) + if err != nil { + return nil, err + } + + return convert(output), nil +} + +// Delete takes a message and deletes it from the queue. +func (s SQSClient) Delete(ctx context.Context, msg Message) error { + if msg.ReceiptHandle == "" { + return ErrReceiptHandleEmpty + } + + _, err := s.client.DeleteMessage(ctx, &sqs.DeleteMessageInput{ + QueueUrl: s.input.QueueUrl, + ReceiptHandle: &msg.ReceiptHandle, + }) + if err != nil { + return err + } + + return nil +} + +// NewSQSClient makes a new SQSClient from an AWS SQS client and the parameters used for ReceiveMessage. +func NewSQSClient(client *sqs.Client, input *sqs.ReceiveMessageInput) *SQSClient { + return &SQSClient{ + client: client, + input: input, + } +} + +// convert will take sqs output and convert it into a slice of Message. +func convert(output *sqs.ReceiveMessageOutput) []Message { + messages := make([]Message, 0, len(output.Messages)) + + for _, msg := range output.Messages { + if msg.MessageId == nil { + // something horrible has happened + continue + } + + if msg.Body == nil { + var empty string + msg.Body = &empty + } + + messages = append(messages, Message{ + ID: *msg.MessageId, + Body: *msg.Body, + ReceiptHandle: *msg.ReceiptHandle, + Attributes: msg.Attributes, + MessageAttributes: convertAttribute(msg.MessageAttributes), + }) + } + + return messages +} + +// convertAttribute turns an AWS type into `any` instead. +func convertAttribute(attributes map[string]types.MessageAttributeValue) map[string]any { + attrs := make(map[string]any) + + for k, v := range attributes { + if v.DataType == nil { + continue + } + + switch *v.DataType { + case "String", "Number": + attrs[k] = *v.StringValue + case "Binary": + attrs[k] = v.BinaryValue + } + } + + return attrs +}