diff --git a/client_test.go b/client_test.go index 908ee18..3dfa92a 100644 --- a/client_test.go +++ b/client_test.go @@ -2,6 +2,7 @@ package kq import ( "context" + "sync" "testing" "time" @@ -9,12 +10,17 @@ import ( ) type fakeProducer struct { - record *kgo.Record - err error + mu sync.Mutex + record *kgo.Record + records []*kgo.Record + err error } func (p *fakeProducer) Produce(_ context.Context, record *kgo.Record) error { + p.mu.Lock() + defer p.mu.Unlock() p.record = record + p.records = append(p.records, record) return p.err } diff --git a/share_consumer.go b/share_consumer.go index 870959a..09c3ed4 100644 --- a/share_consumer.go +++ b/share_consumer.go @@ -2,7 +2,9 @@ package kq import ( "context" + "errors" "fmt" + "sync" "github.com/twmb/franz-go/pkg/kgo" ) @@ -22,14 +24,18 @@ type shareGroupConsumer interface { type kafkaShareGroupConsumer struct { client *kgo.Client + mu sync.Mutex + ackErr error } func newKafkaShareGroupConsumer(config WorkerConfig) (*kafkaShareGroupConsumer, error) { + consumer := new(kafkaShareGroupConsumer) options, err := config.Config.kafkaOptions( kgo.ConsumeTopics(config.readyTopic()), kgo.ShareGroup(config.workerGroup()), kgo.ShareMaxRecords(int32(config.Concurrency)), kgo.ShareMaxRecordsStrict(), + kgo.ShareAckCallback(consumer.recordAckResult), ) if err != nil { return nil, err @@ -40,7 +46,8 @@ func newKafkaShareGroupConsumer(config WorkerConfig) (*kafkaShareGroupConsumer, return nil, fmt.Errorf("kq: create share consumer: %w", err) } - return &kafkaShareGroupConsumer{client: client}, nil + consumer.client = client + return consumer, nil } func (c *kafkaShareGroupConsumer) Poll(ctx context.Context, limit int) sharePollResult { @@ -57,9 +64,26 @@ func (*kafkaShareGroupConsumer) Ack(record *kgo.Record, status kgo.AckStatus) { } func (c *kafkaShareGroupConsumer) FlushAcks(ctx context.Context) error { - return c.client.FlushAcks(ctx) + flushErr := c.client.FlushAcks(ctx) + return errors.Join(flushErr, c.takeAckError()) +} + +func (c *kafkaShareGroupConsumer) takeAckError() error { + c.mu.Lock() + defer c.mu.Unlock() + ackErr := c.ackErr + c.ackErr = nil + return ackErr } func (c *kafkaShareGroupConsumer) Close() { c.client.Close() } + +func (c *kafkaShareGroupConsumer) recordAckResult(_ *kgo.Client, results kgo.ShareAckResults) { + if err := results.Error(); err != nil { + c.mu.Lock() + c.ackErr = errors.Join(c.ackErr, err) + c.mu.Unlock() + } +} diff --git a/share_consumer_test.go b/share_consumer_test.go new file mode 100644 index 0000000..64dabd7 --- /dev/null +++ b/share_consumer_test.go @@ -0,0 +1,25 @@ +package kq + +import ( + "errors" + "testing" + + "github.com/twmb/franz-go/pkg/kgo" +) + +func TestKafkaShareGroupConsumerTracksAckErrors(t *testing.T) { + first := errors.New("first ack failed") + second := errors.New("second ack failed") + consumer := new(kafkaShareGroupConsumer) + + consumer.recordAckResult(nil, kgo.ShareAckResults{{Err: first}}) + consumer.recordAckResult(nil, kgo.ShareAckResults{{Err: second}}) + + err := consumer.takeAckError() + if !errors.Is(err, first) || !errors.Is(err, second) { + t.Fatalf("error = %v, want both ack errors", err) + } + if err := consumer.takeAckError(); err != nil { + t.Fatalf("second take = %v, want nil", err) + } +} diff --git a/worker.go b/worker.go index 4b49479..e9436ae 100644 --- a/worker.go +++ b/worker.go @@ -121,6 +121,9 @@ func (w *Worker) Run(ctx context.Context) error { if ackErr != nil { ackErr = fmt.Errorf("kq: ack tasks: %w", ackErr) } + if ctx.Err() != nil { + return ackErr + } if err := errors.Join(taskErr, ackErr, pollErr); err != nil { return err } @@ -140,6 +143,9 @@ func (w *Worker) handle(ctx context.Context, value []byte) error { if handlerErr == nil { return nil } + if ctx.Err() != nil { + return errors.Join(handlerErr, ctx.Err()) + } if envelope.Retried >= envelope.Retries { if err := w.deadLetter(ctx, envelope, handlerErr); err != nil { diff --git a/worker_test.go b/worker_test.go index 5357583..c8a70d7 100644 --- a/worker_test.go +++ b/worker_test.go @@ -14,12 +14,16 @@ import ( ) type fakeShareGroupConsumer struct { - mu sync.Mutex - results []sharePollResult - pollLimits []int - acks map[*kgo.Record][]kgo.AckStatus - flushes int - flushErr error + mu sync.Mutex + results []sharePollResult + pollLimits []int + acks map[*kgo.Record][]kgo.AckStatus + ackOrder []*kgo.Record + acked chan *kgo.Record + flushes int + flushErr error + flushStarted chan struct{} + flushRelease chan struct{} } func (c *fakeShareGroupConsumer) Poll(ctx context.Context, limit int) sharePollResult { @@ -39,22 +43,54 @@ func (c *fakeShareGroupConsumer) Poll(ctx context.Context, limit int) sharePollR func (c *fakeShareGroupConsumer) Ack(record *kgo.Record, status kgo.AckStatus) { c.mu.Lock() - defer c.mu.Unlock() if c.acks == nil { c.acks = make(map[*kgo.Record][]kgo.AckStatus) } c.acks[record] = append(c.acks[record], status) + c.ackOrder = append(c.ackOrder, record) + acked := c.acked + c.mu.Unlock() + if acked != nil { + acked <- record + } } func (c *fakeShareGroupConsumer) FlushAcks(context.Context) error { c.mu.Lock() - defer c.mu.Unlock() c.flushes++ - return c.flushErr + err := c.flushErr + started := c.flushStarted + release := c.flushRelease + c.mu.Unlock() + if started != nil { + started <- struct{}{} + } + if release != nil { + <-release + } + return err } func (*fakeShareGroupConsumer) Close() {} +func (c *fakeShareGroupConsumer) pollCount() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.pollLimits) +} + +func (c *fakeShareGroupConsumer) statuses(record *kgo.Record) []kgo.AckStatus { + c.mu.Lock() + defer c.mu.Unlock() + return append([]kgo.AckStatus(nil), c.acks[record]...) +} + +func (c *fakeShareGroupConsumer) flushCount() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.flushes +} + func TestWorkerProcessesBoundedPollGeneration(t *testing.T) { records := []*kgo.Record{ {Value: encodeWorkerTestTask(t, Task{Type: "first"}, 0)}, @@ -152,6 +188,169 @@ func TestWorkerExecutesGenerationConcurrently(t *testing.T) { } } +func TestWorkerFlushesGenerationBeforeNextPoll(t *testing.T) { + records := []*kgo.Record{ + {Value: encodeWorkerTestTask(t, Task{Type: "first"}, 0)}, + {Value: encodeWorkerTestTask(t, Task{Type: "second"}, 0)}, + } + consumer := &fakeShareGroupConsumer{ + results: []sharePollResult{ + {records: records[:1]}, + {records: records[1:]}, + {closed: true}, + }, + flushStarted: make(chan struct{}, 2), + flushRelease: make(chan struct{}), + } + worker := &Worker{ + consumer: consumer, + producer: new(fakeProducer), + config: WorkerConfig{Concurrency: 1}, + handler: func(context.Context, Task) error { return nil }, + } + done := make(chan error, 1) + go func() { done <- worker.Run(context.Background()) }() + + <-consumer.flushStarted + if got := consumer.pollCount(); got != 1 { + t.Fatalf("polls during first flush = %d, want 1", got) + } + consumer.flushRelease <- struct{}{} + <-consumer.flushStarted + if got := consumer.pollCount(); got != 2 { + t.Fatalf("polls during second flush = %d, want 2", got) + } + consumer.flushRelease <- struct{}{} + if err := <-done; err != nil { + t.Fatal(err) + } + if got := consumer.pollCount(); got != 3 { + t.Fatalf("total polls = %d, want 3", got) + } +} + +func TestWorkerAcknowledgesTasksAsTheyFinish(t *testing.T) { + slow := &kgo.Record{Value: encodeWorkerTestTask(t, Task{Type: "slow"}, 0)} + fast := &kgo.Record{Value: encodeWorkerTestTask(t, Task{Type: "fast"}, 0)} + consumer := &fakeShareGroupConsumer{ + results: []sharePollResult{{records: []*kgo.Record{slow, fast}}, {closed: true}}, + acked: make(chan *kgo.Record, 2), + } + slowStarted := make(chan struct{}) + releaseSlow := make(chan struct{}) + worker := &Worker{ + consumer: consumer, + producer: new(fakeProducer), + config: WorkerConfig{Concurrency: 2}, + handler: func(_ context.Context, task Task) error { + if task.Type == "slow" { + close(slowStarted) + <-releaseSlow + } + return nil + }, + } + done := make(chan error, 1) + go func() { done <- worker.Run(context.Background()) }() + + <-slowStarted + if first := <-consumer.acked; first != fast { + t.Fatalf("first acknowledged record = %p, want fast record %p", first, fast) + } + close(releaseSlow) + if err := <-done; err != nil { + t.Fatal(err) + } +} + +func TestWorkerJoinsConcurrentTaskAndFlushErrors(t *testing.T) { + first := errors.New("first handler failed") + second := errors.New("second handler failed") + writeErr := errors.New("retry write failed") + flushErr := errors.New("ack flush failed") + records := []*kgo.Record{ + {Value: encodeWorkerTestTask(t, Task{Type: "first"}, 1)}, + {Value: encodeWorkerTestTask(t, Task{Type: "second"}, 1)}, + } + consumer := &fakeShareGroupConsumer{ + results: []sharePollResult{{records: records}}, + flushErr: flushErr, + } + config := workerRetryConfig(t, time.Minute) + config.Concurrency = 2 + worker := &Worker{ + consumer: consumer, + producer: &fakeProducer{err: writeErr}, + config: config, + handler: func(_ context.Context, task Task) error { + if task.Type == "first" { + return first + } + return second + }, + } + + err := worker.Run(context.Background()) + for _, want := range []error{first, second, writeErr, flushErr} { + if !errors.Is(err, want) { + t.Fatalf("error = %v, want cause %v", err, want) + } + } + for _, record := range records { + statuses := consumer.statuses(record) + if len(statuses) != 1 || statuses[0] != kgo.AckRelease { + t.Fatalf("statuses = %v, want [release]", statuses) + } + } +} + +func TestWorkerCancellationDrainsPool(t *testing.T) { + const concurrency = 4 + records := make([]*kgo.Record, 0, concurrency) + for i := range concurrency { + records = append(records, &kgo.Record{ + Value: encodeWorkerTestTask(t, Task{Type: fmt.Sprintf("task-%d", i)}, 0), + }) + } + consumer := &fakeShareGroupConsumer{results: []sharePollResult{{records: records}}} + started := make(chan struct{}, concurrency) + worker := &Worker{ + consumer: consumer, + producer: new(fakeProducer), + config: WorkerConfig{Concurrency: concurrency}, + handler: func(ctx context.Context, _ Task) error { + started <- struct{}{} + <-ctx.Done() + return ctx.Err() + }, + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- worker.Run(ctx) }() + + for range concurrency { + <-started + } + cancel() + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-time.After(time.Second): + t.Fatal("worker did not drain canceled pool") + } + for _, record := range records { + statuses := consumer.statuses(record) + if len(statuses) != 1 || statuses[0] != kgo.AckRelease { + t.Fatalf("statuses = %v, want [release]", statuses) + } + } + if got := consumer.flushCount(); got != 1 { + t.Fatalf("flushes = %d, want 1", got) + } +} + func TestWorkerHandlesTask(t *testing.T) { producer := new(fakeProducer) value := encodeWorkerTestTask(t, Task{