feat: reconcile mailbox sync recovery with merged worker Kafka retry and broker-confirmed delivery safeguards

This commit is contained in:
Devin AI
2026-10-09 18:18:26 +00:00
25 changed files with 1221 additions and 91 deletions
+1 -1
View File
@@ -202,7 +202,7 @@ jobs:
env:
WARMBLY_TEST_DB: postgres://warmbly:warmbly@127.0.0.1:5432/warmbly_test?sslmode=disable
run: |
go test -race ./internal/repository -run '^TestLive(Monitoring|LastWarmupSendFailure|PersistentMailboxLoadingFailure|DispatchRetry|LineageQueueRevisions|AdminPlatformOverview|WarmupRecoversOldPreparation|WarmupPreparationRecoveryRetains)' -count=1
go test -race ./internal/repository -run '^TestLive(Monitoring|LastWarmupSendFailure|PersistentMailboxLoadingFailure|DispatchRetry|LineageQueueRevisions|AdminPlatformOverview|WarmupRecoversOldPreparation|WarmupPreparationRecoveryRetains|WarmupFilingRequiresAcknowledgement)' -count=1
go test -race ./internal/tasks -run '^Test(LiveWarmup(SafetyHoldReschedules|Recovery|TaskDoesNotApplyResults)|WarmupSenderDistinguishes)' -count=1
- name: Mailbox sync admission, replay, truthful state and upgrade regressions (PostgreSQL)
+1 -1
View File
@@ -1943,7 +1943,7 @@ func main() {
// warming or backing a live campaign (the health-check lane). This is
// the bootstrap — enabling warmup or starting a campaign doesn't itself
// enqueue the first warmup task.
go tasksService.StartWarmupReconciler(ctx, 10*time.Minute)
go tasksService.StartWarmupReconciler(ctx, time.Minute)
go jobs.NewWarmupOperationsJob(warmupRepository, taskRepository, cache, opsNotifier).Start(ctx)
// Tracking-domain sweep: re-resolve every custom tracking domain hourly,
+6 -1
View File
@@ -218,7 +218,12 @@ func main() {
}()
log.Printf("Worker %s started, listening on topic %s", workerID, workerTopic)
if err := bus.Subscribe(ctx, []string{workerTopic}, "worker-"+workerID.String(), workerService.Receive); err != nil {
lanes := worker.CommandConcurrency()
if _, ok := bus.(eventbus.KeyedSubscriber); !ok {
lanes = 1
}
log.Printf("Worker command concurrency: %d ordered mailbox lanes", lanes)
if err := eventbus.SubscribeKeyed(ctx, bus, []string{workerTopic}, "worker-"+workerID.String(), lanes, workerService.ResolveCommandKey, workerService.Receive); err != nil {
log.Println("event bus subscribe ended:", err)
}
@@ -311,6 +311,8 @@ On `filesystem`, a remote worker writes blobs to its own disk rather than a volu
| `PUBSUB_ENABLED` | `false` uses the Redis bridge for realtime fanout, `true` uses Google Pub/Sub | `false` | yes |
| `GCP_PROJECT_ID` | Project when `PUBSUB_ENABLED=true` | unset | yes |
Kafka event publication waits for the broker's delivery report with all in-sync replica acknowledgements, bounded by the caller's deadline and `EVENTBUS_HANDLER_TIMEOUT`. Local queue admission is not evidence of publication. A rejected or timed-out result keeps the original command unacknowledged for replay; a timeout does not prove the event was never delivered, so consumers still deduplicate replayed outcomes. Generic producer events are drained separately through shutdown, without consuming a publication's independent delivery report.
<Callout title="Topics are created by the bus, not by hand">
A worker's command topic is named after the node id issued when it joined, so the set of topics is not knowable in advance. The Kafka backend creates what it uses, once per topic per process, rather than depending on the broker's `auto.create.topics.enable`, which is off by default on Confluent Cloud and configurable only on Dedicated clusters, not on Basic, Standard, Enterprise or Freight. The credential therefore needs `CREATE` on topics as well as read and write.
</Callout>
@@ -436,6 +438,7 @@ Keep the copies private. GeoLite2 may not be redistributed, which is also why no
| `WORKER_BIND_IP` | Source address to bind outbound connections to, and the seed for a derived `WORKER_ID` | unset | yes |
| `WORKER_PUBLIC_IP` | Public IPv4 reported by the worker, never replaced by the backend-observed address, which may be a reverse proxy. Otherwise the node discovers HTTPS egress IPv4 using `api.ipify.org` every ten minutes. Set this explicitly when mail and HTTPS have different egress routes | detected | yes |
| `WARMBLY_WORKER_CAPACITY` | Assigned-mailbox planning target for this worker. This guides placement and is not a mailbox-provider send limit | `100` | yes |
| `WORKER_COMMAND_CONCURRENCY` | Kafka worker command lanes, from `1` to `64`. Commands are grouped by the mailbox in the payload, including older task-keyed sends. Each mailbox stays in order; independent lanes run concurrently. NATS remains serial. Invalid values use the default | `8` | yes |
| `WARMBLY_NODE_REGION` | Free-form label for where this node egresses from, e.g. `eu-central`. Placement prefers a worker near where a mailbox's provider expects sign-ins; unset scores neutral. Written by the join script from `--region` | unset | yes |
| `ENCRYPTED_KEYS_BACKEND_URL` | Backend base the worker fetches organization keys from | unset | yes |
| `ENCRYPTED_KEYS_WORKER_TOKEN` | The worker's copy of `INTERNAL_API_TOKEN` | unset | yes |
+6
View File
@@ -37,6 +37,12 @@ Persistent loading failures and pending warmup tasks at least an hour past their
A missing mailbox also requests another reload through the worker heartbeat. Existing warmup sends stuck before execution recover automatically: after ten minutes, the scheduler retires a nonce-protected command only when neither worker execution nor a delivery receipt exists, refunds its reserved count, and resumes the mailbox's normal schedule. Delayed copies of the retired command cannot send. Customers do not need to disconnect or re-add their mailboxes. Actually started sends with missing outcomes, conflicting outcomes, provider cooldowns, authentication failures and other safety holds are not cleared by this repair.
Recovery runs at backend startup and every minute in bounded batches. The ten-minute cutoff is unchanged: it protects recent authorizations, not a waiting period added after each release. Existing older commands can be retired on the first pass; large fleets can need subsequent passes to reseed every schedule. Sending hours, quotas and actual provider failures still apply.
Kafka workers process independent mailbox lanes concurrently instead of making the whole worker wait for one command. Older queued commands are grouped by their mailbox payload too. The broker checkpoint never passes unfinished or failed work, and failed subscriptions reopen for replay. The default is eight lanes; see [worker configuration](/development/configuration/#workers).
Unacknowledged warmup filing actions remain durable. They retry every five minutes initially, then every fifteen minutes after twenty minutes of pending time, capped at every thirty minutes after an hour. Repeated arrivals do not reset that age. This reduces repeat commands during a worker backlog without discarding pending filing or counting it as a failed send.
If outbound admission refuses a warmup command before publication, a successful deferral returns it to the schedule without recording a failed send. A deferral whose database outcome is uncertain is not treated as proof of failure. Existing records affected by an older preparation-result race also recover automatically, but only when the recorded internal preparation failure predates the new authorization and there is no execution, outbound reservation, provider-confirmed send or receipt. Only that provably unstarted command's false conflict is released; genuine conflicting outcomes remain held.
Follow a provider-specific error when one is shown. Google and Microsoft sign-ins use their native APIs, not SMTP settings. Reconnect those mailboxes only for a reported sign-in or permission error. SMTP host, port, password and IP allowlist checks apply to SMTP mailboxes when the error reports a mail-server connection problem. Warmbly Cloud makes those connections from its own network, not your self-hosted instance's address.
+51
View File
@@ -0,0 +1,51 @@
package worker
import (
"context"
"encoding/json"
"fmt"
"os"
"strconv"
"github.com/google/uuid"
"github.com/warmbly/warmbly/internal/infrastructure/eventbus"
"github.com/warmbly/warmbly/internal/models"
)
func CommandConcurrency() int {
if n, err := strconv.Atoi(os.Getenv("WORKER_COMMAND_CONCURRENCY")); err == nil && n >= 1 && n <= 64 {
return n
}
return 8
}
// ResolveCommandKey uses the payload so old task-keyed sends share their mailbox's lane too.
func (w *WorkerService) ResolveCommandKey(ctx context.Context, msg eventbus.Message) (string, error) {
var event models.WorkerEvent
if err := w.Codec.Deserialize(ctx, msg.Topic, msg.Payload, &event); err != nil {
return "", err
}
raw, err := json.Marshal(event.Body)
if err != nil {
return "", err
}
var body struct {
EmailID uuid.UUID `json:"email_id"`
ID uuid.UUID `json:"id"`
ProcessID uuid.UUID `json:"process_id"`
}
if err := json.Unmarshal(raw, &body); err != nil {
return "", err
}
id := body.EmailID
switch event.Type {
case models.WorkerEventTypeAddEmail:
id = body.ID
case models.WorkerEventTypeEmailValidation:
id = body.ProcessID
}
if id == uuid.Nil {
return "", fmt.Errorf("worker command %s has no mailbox or validation identity", event.Type)
}
return id.String(), nil
}
@@ -0,0 +1,17 @@
//go:build kafka
package worker
import (
"testing"
"github.com/warmbly/warmbly/internal/infrastructure/codec"
)
func TestAvroCommandKeysGroupLegacySendKeysWithMailboxMutations(t *testing.T) {
c, err := codec.NewAvro("mock://command-key", "", "")
if err != nil {
t.Fatal(err)
}
checkCommandKeys(t, &WorkerService{Codec: c})
}
+62
View File
@@ -0,0 +1,62 @@
package worker
import (
"testing"
"github.com/google/uuid"
"github.com/warmbly/warmbly/internal/infrastructure/codec"
"github.com/warmbly/warmbly/internal/infrastructure/eventbus"
"github.com/warmbly/warmbly/internal/models"
)
func TestCommandKeysGroupOldSendKeysWithMailboxMutations(t *testing.T) {
w := &WorkerService{Codec: codec.NewJSON()}
checkCommandKeys(t, w)
for _, body := range []string{`{"type":"SEND_EMAIL","body":{}}`, `{"type":"REMOVE_EMAIL","body":{"email_id":"invalid"}}`, "invalid"} {
if _, err := w.ResolveCommandKey(t.Context(), eventbus.Message{Payload: []byte(body)}); err == nil {
t.Fatal("malformed command was assigned an unsafe fallback key")
}
}
}
func checkCommandKeys(t *testing.T, w *WorkerService) {
t.Helper()
mailbox, task, process := uuid.New(), uuid.New(), uuid.New()
for _, test := range []struct {
kind models.WorkerEventType
body any
want uuid.UUID
}{
{models.WorkerEventTypeSendEmail, &models.SendEmail{EmailID: mailbox, TaskID: task}, mailbox},
{models.WorkerEventTypeWarmupAction, &models.WarmupEmailAction{EmailID: mailbox}, mailbox},
{models.WorkerEventTypeMessageSeen, &models.MessageSeenAction{EmailID: mailbox}, mailbox},
{models.WorkerEventTypeMessageFolder, &models.MessageFolderAction{EmailID: mailbox}, mailbox},
{models.WorkerEventTypeMailboxIdentity, models.EventWorkerMailboxIdentity{EmailID: mailbox}, mailbox},
{models.WorkerEventTypeAddEmail, &models.AddWorkerEmail{ID: mailbox}, mailbox},
{models.WorkerEventTypeRemoveEmail, &models.RemoveWorkerEmail{EmailID: mailbox.String()}, mailbox},
{models.WorkerEventTypeEmailValidation, models.EventWorkerEmailValidation{OrgID: mailbox, ProcessID: process}, process},
} {
t.Run(string(test.kind), func(t *testing.T) {
payload, err := w.Codec.Serialize(t.Context(), "commands", models.WorkerEvent{Type: test.kind, Body: test.body})
if err != nil {
t.Fatal(err)
}
key, err := w.ResolveCommandKey(t.Context(), eventbus.Message{Topic: "commands", Key: task.String(), Payload: payload})
if err != nil || key != test.want.String() {
t.Fatalf("key=%q want=%q err=%v", key, test.want, err)
}
})
}
}
func TestCommandConcurrencyDefaultsAndBoundedOverride(t *testing.T) {
for _, test := range []struct {
env string
want int
}{{"", 8}, {"invalid", 8}, {"0", 8}, {"65", 8}, {"1", 1}, {"16", 16}, {"64", 64}} {
t.Setenv("WORKER_COMMAND_CONCURRENCY", test.env)
if got := CommandConcurrency(); got != test.want {
t.Fatalf("env=%q got=%d want=%d", test.env, got, test.want)
}
}
}
+5 -7
View File
@@ -344,10 +344,7 @@ func (w *WorkerService) sendEmailSuccess(taskID uuid.UUID, messageID, providerMs
}
}
// sendNotLoadedRedeliveries is how many bus deliveries a send may burn waiting
// for a transient condition (mailbox not loaded yet, object storage blip)
// before the worker reports it failed. Each redelivery is a second apart, and
// the bus stops redelivering at ten, so this stays well inside that.
// Report unattempted sends before either broker's retry budget expires.
const sendNotLoadedRedeliveries = 5
// failSend reports a send the worker could not attempt. A retryable condition
@@ -360,8 +357,7 @@ func (w *WorkerService) failSend(ctx context.Context, sendEmail models.SendEmail
if d := deliveryOf(ctx); retryable && d.redelivers && d.attempt < sendNotLoadedRedeliveries {
return errors.New(reason)
}
w.sendEmailFailure(sendEmail.TaskID, sendEmail.EmailID, nil, reason)
return nil
return w.sendEmailFailure(sendEmail.TaskID, sendEmail.EmailID, nil, reason)
}
// sendEmailError reports a failed send attempt. The per-task result is always
@@ -416,7 +412,7 @@ func (w *WorkerService) sendEmailError(taskID uuid.UUID, emailID uuid.UUID, mail
}
// sendEmailFailure sends a generic failure result (for non-MailError cases)
func (w *WorkerService) sendEmailFailure(taskID uuid.UUID, emailID uuid.UUID, mail *wmail.WMail, errorMsg string) {
func (w *WorkerService) sendEmailFailure(taskID uuid.UUID, emailID uuid.UUID, mail *wmail.WMail, errorMsg string) error {
now := time.Now().UTC()
result := models.SendEmailResult{
TaskID: taskID,
@@ -430,5 +426,7 @@ func (w *WorkerService) sendEmailFailure(taskID uuid.UUID, emailID uuid.UUID, ma
if err := w.Produce(models.JobEventTypeEmailFailed, taskID.String(), result); err != nil {
log.Error().Err(err).Str("task_id", taskID.String()).Msg("Failed to produce email failure event")
return err
}
return nil
}
@@ -2,9 +2,12 @@ package worker
import (
"context"
"errors"
"testing"
"github.com/google/uuid"
"github.com/warmbly/warmbly/internal/infrastructure/codec"
"github.com/warmbly/warmbly/internal/infrastructure/eventbus"
"github.com/warmbly/warmbly/internal/models"
"github.com/warmbly/warmbly/internal/pkg/emailverify"
@@ -29,3 +32,34 @@ func TestMissingMailboxRequestsReloadBeforeRedelivery(t *testing.T) {
t.Fatal("missing mailbox did not request a reload")
}
}
func TestUnloadedMailboxPreparationResultMustPublishBeforeAcknowledgement(t *testing.T) {
w := newLoadingWorker(&capturedEvents{})
bus := &filingBus{err: errors.New("broker unavailable")}
w.Bus, w.Codec = bus, codec.NewJSON()
w.InitEvents()
event := &models.WorkerEvent{Type: models.WorkerEventTypeSendEmail, Body: &models.SendEmail{TaskID: uuid.New(), EmailID: uuid.New()}}
payload, err := w.Codec.Serialize(t.Context(), "commands", event)
if err != nil {
t.Fatal(err)
}
msg := eventbus.Message{Payload: payload, Attempt: 5, Redelivers: true}
if err := w.Receive(t.Context(), msg); !errors.Is(err, bus.err) || len(bus.events) != 0 {
t.Fatalf("unpublished failure was acknowledged: events=%v err=%v", bus.events, err)
}
bus.err = nil
if err := w.Receive(t.Context(), msg); err != nil || len(bus.events) != 1 {
t.Fatalf("preparation did not resolve: events=%v err=%v", bus.events, err)
}
raw, err := w.Codec.Serialize(t.Context(), "result", bus.events[0].Body)
if err != nil {
t.Fatal(err)
}
var result models.SendEmailResult
if err := w.Codec.Deserialize(t.Context(), "result", raw, &result); err != nil {
t.Fatal(err)
}
if result.Success || result.Error == nil || result.Error.Failure == nil || result.Error.Failure.Stage != "prepare" || result.Error.Failure.Protocol != "internal" {
t.Fatalf("unattempted send misrepresented as provider evidence: %+v", result)
}
}
+1 -1
View File
@@ -203,7 +203,7 @@ func (p *publisher) PublishSendEmail(ctx context.Context, workerID uuid.UUID, pa
}
workerTopic := kafka.GetWorkerTopic(workerID.String())
return p.publish(workerTopic, params.TaskID.String(), workerEvent)
return p.publish(workerTopic, params.EmailID.String(), workerEvent)
}
// StoreEmailBody stores email body in S3 and returns the S3 key. It is the
+31 -2
View File
@@ -40,6 +40,8 @@ import (
// distinguish it from a broker fault so they do not report a clean stop.
var ErrBusClosed = errors.New("eventbus: bus closed")
var ErrSubscriptionRebalanced = errors.New("eventbus: subscription rebalanced")
// EventBus is the transport-level interface. Implementations must be safe for
// concurrent use by multiple goroutines.
type EventBus interface {
@@ -72,6 +74,29 @@ type EventBus interface {
// the message; returning an error leaves it for redelivery.
type Handler func(ctx context.Context, msg Message) error
type KeyResolver func(ctx context.Context, msg Message) (string, error)
func resolveKey(ctx context.Context, key KeyResolver, msg Message) (value string, err error) {
defer func() {
if r := recover(); r != nil {
err = fmt.Errorf("eventbus: key resolver panic on topic %s: %v", msg.Topic, r)
}
}()
return key(ctx, msg)
}
// KeyedSubscriber preserves resolved-key order while handling independent keys concurrently.
type KeyedSubscriber interface {
SubscribeKeyed(context.Context, []string, string, int, KeyResolver, Handler) error
}
func SubscribeKeyed(ctx context.Context, bus EventBus, topics []string, group string, lanes int, key KeyResolver, handler Handler) error {
if keyed, ok := bus.(KeyedSubscriber); ok {
return keyed.SubscribeKeyed(ctx, topics, group, lanes, key, handler)
}
return bus.Subscribe(ctx, topics, group, handler)
}
func retrySubscription(ctx context.Context, subscribe func(context.Context) error) error {
for {
if err := ctx.Err(); err != nil {
@@ -84,7 +109,11 @@ func retrySubscription(ctx context.Context, subscribe func(context.Context) erro
if errors.Is(err, ErrBusClosed) {
return err
}
log.Warn().Msg("eventbus subscription interrupted; reopening for replay")
if errors.Is(err, ErrSubscriptionRebalanced) {
log.Debug().Msg("eventbus subscription rebalanced; reopening for replay")
} else {
log.Warn().Err(err).Msg("eventbus subscription interrupted; reopening for replay")
}
timer := time.NewTimer(time.Second)
select {
case <-ctx.Done():
@@ -106,7 +135,7 @@ type Message struct {
Key string
Payload []byte
// Attempt is the 1-based delivery count of this message (NATS reports it
// from the consumer's redelivery metadata; Kafka always reports 1). A
// from redelivery metadata; Kafka counts retries within a subscription). A
// handler that retries by returning an error can read it to know when the
// broker is about to stop redelivering and give up cleanly instead.
Attempt int
+53 -27
View File
@@ -78,10 +78,7 @@ func NewKafkaFromProducer(p *kafka.Producer, cfg KafkaConfig) *KafkaBus {
func (b *KafkaBus) Name() string { return "kafka" }
// Publish writes to the underlying producer. The context is checked for
// cancellation before the call but not used to bound the produce itself: the
// confluent-kafka-go producer is asynchronous and Produce returns immediately
// after enqueuing.
// Publish succeeds only after a native broker delivery report.
func (b *KafkaBus) Publish(ctx context.Context, topic, key string, payload []byte) error {
if err := ctx.Err(); err != nil {
return err
@@ -97,7 +94,9 @@ func (b *KafkaBus) Publish(ctx context.Context, topic, key string, payload []byt
if err := b.ensureTopics(ctx, topic); err != nil {
return err
}
err := b.producer.Produce(topic, []byte(key), payload)
pctx, cancel := context.WithTimeout(ctx, handlerTimeout())
defer cancel()
err := b.producer.ProduceConfirmed(pctx, topic, []byte(key), payload)
if errors.Is(err, kafka.ErrClientClosed) {
return ErrBusClosed
}
@@ -108,6 +107,10 @@ func (b *KafkaBus) Publish(ctx context.Context, topic, key string, payload []byt
// topics, and blocks reading messages until ctx is cancelled or a fatal error
// occurs. Failed handlers never store offsets; exhausted retries reopen for replay.
func (b *KafkaBus) Subscribe(ctx context.Context, topics []string, group string, handler Handler) error {
return b.SubscribeKeyed(ctx, topics, group, 1, nil, handler)
}
func (b *KafkaBus) SubscribeKeyed(ctx context.Context, topics []string, group string, lanes int, key KeyResolver, handler Handler) error {
if len(topics) == 0 {
return errors.New("eventbus kafka: at least one topic required")
}
@@ -117,12 +120,15 @@ func (b *KafkaBus) Subscribe(ctx context.Context, topics []string, group string,
if handler == nil {
return errors.New("eventbus kafka: handler required")
}
if lanes < 1 || lanes > 64 || (lanes > 1 && key == nil) {
return errors.New("eventbus kafka: keyed subscription requires 1-64 lanes and a key resolver")
}
return retrySubscription(ctx, func(ctx context.Context) error {
return b.subscribeOnce(ctx, topics, group, handler)
return b.subscribeOnce(ctx, topics, group, lanes, key, handler)
})
}
func (b *KafkaBus) subscribeOnce(ctx context.Context, topics []string, group string, handler Handler) error {
func (b *KafkaBus) subscribeOnce(ctx context.Context, topics []string, group string, lanes int, key KeyResolver, handler Handler) error {
// Subscribing to a topic that does not exist yet returns no messages and
// no error, so a worker would sit silent rather than fail.
@@ -170,32 +176,52 @@ func (b *KafkaBus) subscribeOnce(ctx context.Context, topics []string, group str
cons.Close()
}()
err = cons.Consume(ctx, func(msg *ckf.Message) error {
topic := ""
if msg.TopicPartition.Topic != nil {
topic = *msg.TopicPartition.Topic
}
// A message already being handled finishes after a shutdown signal; only reading the next one stops.
hctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), handlerTimeout())
defer cancel()
if err := invokeHandler(hctx, handler, Message{
Topic: topic,
Key: string(msg.Key),
Payload: msg.Value,
Attempt: 1,
Redelivers: true,
}); err != nil {
log.Error().Err(err).Str("topic", topic).Msg("eventbus kafka handler error")
return err
}
return nil
})
resolve := func(msg *ckf.Message) (string, error) {
return resolveKey(ctx, key, kafkaEnvelope(msg))
}
deliver := kafkaHandler(ctx, handler)
err = cons.ConsumeConcurrent(ctx, lanes, resolve, deliver)
if errors.Is(err, kafka.ErrClientClosed) {
return ErrBusClosed
}
if errors.Is(err, kafka.ErrAssignmentLost) {
return errors.Join(ErrSubscriptionRebalanced, err)
}
return err
}
func kafkaHandler(ctx context.Context, handler Handler) func(*ckf.Message) error {
var mu sync.Mutex
attempts := map[*ckf.Message]int{}
return func(msg *ckf.Message) error {
mu.Lock()
attempts[msg]++
attempt := attempts[msg]
mu.Unlock()
envelope := kafkaEnvelope(msg)
envelope.Attempt = attempt
// A message already being handled finishes after a shutdown signal; only reading the next one stops.
hctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), handlerTimeout())
defer cancel()
if err := invokeHandler(hctx, handler, envelope); err != nil {
log.Error().Err(err).Str("topic", envelope.Topic).Msg("eventbus kafka handler error")
return err
}
mu.Lock()
delete(attempts, msg)
mu.Unlock()
return nil
}
}
func kafkaEnvelope(msg *ckf.Message) Message {
topic := ""
if msg.TopicPartition.Topic != nil {
topic = *msg.TopicPartition.Topic
}
return Message{Topic: topic, Key: string(msg.Key), Payload: msg.Value, Attempt: 1, Redelivers: true}
}
// Close flushes the producer and closes every consumer that was opened via
// Subscribe.
func (b *KafkaBus) Close() error {
+31 -1
View File
@@ -7,9 +7,39 @@ import (
"errors"
"testing"
ckf "github.com/confluentinc/confluent-kafka-go/v2/kafka"
"github.com/warmbly/warmbly/internal/infrastructure/kafka"
)
func TestKafkaHandlerCountsIndependentRetriesUntilPreparationCanResolve(t *testing.T) {
first, other := &ckf.Message{}, &ckf.Message{}
var seen []int
deliver := kafkaHandler(t.Context(), func(_ context.Context, msg Message) error {
seen = append(seen, msg.Attempt)
if !msg.Redelivers {
t.Fatal("retry delivery was not marked")
}
if msg.Attempt < 5 {
return errors.New("mailbox not loaded")
}
return nil
})
for attempt := 1; attempt <= 5; attempt++ {
err := deliver(first)
if (err == nil) != (attempt == 5) || seen[len(seen)-1] != attempt {
t.Fatalf("attempt=%d seen=%v err=%v", attempt, seen, err)
}
if attempt == 3 {
if err := deliver(other); err == nil || seen[len(seen)-1] != 1 {
t.Fatal("another mailbox inherited retry state")
}
}
}
if err := deliver(first); err == nil || seen[len(seen)-1] != 1 {
t.Fatal("successful delivery retained retry state")
}
}
func TestKafkaSubscriptionReopensAfterFailureOrUnexpectedExit(t *testing.T) {
for _, failure := range []error{errors.New("offset assignment lost"), nil} {
ctx, cancel := context.WithCancel(t.Context())
@@ -38,7 +68,7 @@ func TestKafkaSubscriptionRetiresClientsBeforeReopening(t *testing.T) {
ctx, cancel := context.WithCancel(t.Context())
cancel()
for i := 0; i < 3; i++ {
if err := b.subscribeOnce(ctx, []string{"sync-test"}, "sync-test", func(context.Context, Message) error { return nil }); !errors.Is(err, context.Canceled) {
if err := b.subscribeOnce(ctx, []string{"sync-test"}, "sync-test", 1, nil, func(context.Context, Message) error { return nil }); !errors.Is(err, context.Canceled) {
t.Fatal(err)
}
b.mu.Lock()
+56 -31
View File
@@ -7,6 +7,7 @@ import (
"errors"
"fmt"
"sync"
"sync/atomic"
"time"
ckf "github.com/confluentinc/confluent-kafka-go/v2/kafka"
@@ -24,11 +25,12 @@ type consumerClient interface {
}
type Consumer struct {
c consumerClient
Avrov2 *Avrov2
mu sync.Mutex
closed bool
retryLimit int
c consumerClient
Avrov2 *Avrov2
mu sync.Mutex
closed bool
retryLimit int
assignmentEpoch atomic.Uint64
}
type ConsumerConfig struct {
@@ -93,7 +95,13 @@ func (cons *Consumer) SubscribeTopics(topics []string) error {
if cons.closed {
return ErrClientClosed
}
return cons.c.SubscribeTopics(topics, nil)
return cons.c.SubscribeTopics(topics, func(_ *ckf.Consumer, event ckf.Event) error {
if _, revoked := event.(ckf.RevokedPartitions); revoked {
cons.assignmentEpoch.Add(1)
}
// Let the native client perform its usual assignment protocol.
return nil
})
}
// Polling and offset storage must finish before the native client is destroyed.
@@ -107,11 +115,18 @@ func (cons *Consumer) readMessage() (*ckf.Message, error) {
}
func (cons *Consumer) storeMessage(msg *ckf.Message) error {
return cons.storeMessageAt(msg, cons.assignmentEpoch.Load())
}
func (cons *Consumer) storeMessageAt(msg *ckf.Message, epoch uint64) error {
cons.mu.Lock()
defer cons.mu.Unlock()
if cons.closed {
return ErrClientClosed
}
if cons.assignmentEpoch.Load() != epoch {
return ErrAssignmentLost
}
if native, ok := cons.c.(interface{ AssignmentLost() bool }); ok && native.AssignmentLost() {
return ErrAssignmentLost
}
@@ -120,12 +135,16 @@ func (cons *Consumer) storeMessage(msg *ckf.Message) error {
}
func (cons *Consumer) Consume(ctx context.Context, handler func(msg *ckf.Message) error) error {
epoch := cons.assignmentEpoch.Load()
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
msg, err := cons.readMessage()
if cons.assignmentEpoch.Load() != epoch {
return ErrAssignmentLost
}
if err != nil {
if kafkaErr, ok := err.(ckf.Error); ok {
if kafkaErr.Code() == ckf.ErrTimedOut {
@@ -144,34 +163,12 @@ func (cons *Consumer) Consume(ctx context.Context, handler func(msg *ckf.Message
return fmt.Errorf("error reading message: %w", err)
}
backoff := 100 * time.Millisecond
limit := cons.retryLimit
if limit <= 0 {
limit = 4
}
for attempt := 1; ; attempt++ {
if err := handler(msg); err == nil {
break
} else {
log.Error().Err(err).Msg("kafka message handler error; retrying before advancing offsets")
if attempt >= limit {
return fmt.Errorf("kafka: handler retry budget exhausted before offset storage: %w", err)
}
}
if err := consumerBackoff(ctx, backoff); err != nil {
return err
}
cons.mu.Lock()
closed := cons.closed
cons.mu.Unlock()
if closed {
return ErrClientClosed
}
backoff = min(2*backoff, 5*time.Second)
if err := cons.handleWithRetry(ctx, msg, handler); err != nil {
return err
}
// Needs enable.auto.offset.store=false; the background commit sends it.
if err := cons.storeMessage(msg); err != nil {
if err := cons.storeMessageAt(msg, epoch); err != nil {
if errors.Is(err, ErrClientClosed) {
return err
}
@@ -191,3 +188,31 @@ func consumerBackoff(ctx context.Context, delay time.Duration) error {
return nil
}
}
func (cons *Consumer) handleWithRetry(ctx context.Context, msg *ckf.Message, handler func(msg *ckf.Message) error) error {
backoff := 100 * time.Millisecond
limit := cons.retryLimit
if limit <= 0 {
limit = 6
}
for attempt := 1; ; attempt++ {
err := handler(msg)
if err == nil {
return nil
}
log.Error().Err(err).Msg("kafka message handler error; retrying before advancing offsets")
if attempt >= limit {
return fmt.Errorf("kafka: handler retry budget exhausted before offset storage: %w", err)
}
if err := consumerBackoff(ctx, backoff); err != nil {
return err
}
cons.mu.Lock()
closed := cons.closed
cons.mu.Unlock()
if closed {
return ErrClientClosed
}
backoff = min(2*backoff, 5*time.Second)
}
}
@@ -0,0 +1,221 @@
//go:build kafka
package kafka
import (
"context"
"errors"
"fmt"
"hash/fnv"
"sync"
"time"
ckf "github.com/confluentinc/confluent-kafka-go/v2/kafka"
"github.com/rs/zerolog/log"
)
const laneBuffer = 32
type partitionKey struct {
topic string
partition int32
}
type laneResult struct {
msg *ckf.Message
err error
}
// partitionOffsets stores an offset only once every earlier read offset in its partition succeeded.
type partitionOffsets struct {
pending []*ckf.Message
done map[ckf.Offset]bool
}
func (p *partitionOffsets) complete(offset ckf.Offset) *ckf.Message {
p.done[offset] = true
var last *ckf.Message
for len(p.pending) > 0 && p.done[p.pending[0].TopicPartition.Offset] {
last = p.pending[0]
delete(p.done, last.TopicPartition.Offset)
p.pending[0] = nil
p.pending = p.pending[1:]
}
return last
}
func messagePartition(msg *ckf.Message) partitionKey {
k := partitionKey{partition: msg.TopicPartition.Partition}
if msg.TopicPartition.Topic != nil {
k.topic = *msg.TopicPartition.Topic
}
return k
}
func messageLane(key string, lanes int) int {
h := fnv.New32a()
_, _ = h.Write([]byte(key))
return int(h.Sum32() % uint32(lanes))
}
// ConsumeConcurrent preserves key order and stores only contiguous successful offsets.
func (cons *Consumer) ConsumeConcurrent(ctx context.Context, lanes int, key func(*ckf.Message) (string, error), handler func(msg *ckf.Message) error) error {
if lanes < 1 || lanes > 64 {
return errors.New("kafka: concurrency must be between 1 and 64")
}
if lanes <= 1 {
return cons.Consume(ctx, handler)
}
results := make(chan laneResult, lanes*laneBuffer)
runCtx, cancel := context.WithCancel(ctx)
defer cancel()
epoch := cons.assignmentEpoch.Load()
queues := make([]chan *ckf.Message, lanes)
var wg sync.WaitGroup
for i := range queues {
queues[i] = make(chan *ckf.Message, laneBuffer)
wg.Add(1)
go func(queue <-chan *ckf.Message) {
defer wg.Done()
var failed bool
for msg := range queue {
if failed || runCtx.Err() != nil {
results <- laneResult{msg: msg, err: errLaneStopped}
continue
}
err := cons.handleWithRetry(runCtx, msg, handler)
failed = err != nil
results <- laneResult{msg: msg, err: err}
}
}(queues[i])
}
partitions := map[partitionKey]*partitionOffsets{}
pending := 0
var stopErr error
settle := func(r laneResult) {
if r.err != nil {
if stopErr == nil && !errors.Is(r.err, errLaneStopped) {
stopErr = r.err
}
return
}
if errors.Is(stopErr, ErrClientClosed) || errors.Is(stopErr, ErrAssignmentLost) {
return
}
// Contiguous successes are stored even after another lane failed.
p := partitions[messagePartition(r.msg)]
before := len(p.pending)
cons.storeSettled(partitions, r.msg, epoch, &stopErr)
pending -= before - len(p.pending)
}
shutdown := func() error {
cancel()
for _, q := range queues {
close(q)
}
go func() { wg.Wait(); close(results) }()
for r := range results {
settle(r)
}
return stopErr
}
for stopErr == nil {
select {
case r := <-results:
settle(r)
continue
case <-ctx.Done():
stopErr = ctx.Err()
continue
default:
}
if pending >= lanes*(laneBuffer+1) {
select {
case r := <-results:
settle(r)
case <-ctx.Done():
stopErr = ctx.Err()
}
continue
}
msg, err := cons.readMessage()
if cons.assignmentEpoch.Load() != epoch {
stopErr = ErrAssignmentLost
continue
}
if err != nil {
var kafkaErr ckf.Error
if errors.As(err, &kafkaErr) {
if kafkaErr.Code() == ckf.ErrTimedOut {
continue
}
if kafkaErr.IsFatal() || kafkaErr.Code() == ckf.ErrMaxPollExceeded {
stopErr = fmt.Errorf("kafka: consumer must reopen: %w", err)
continue
}
log.Warn().Str("code", fmt.Sprintf("%d", kafkaErr.Code())).Err(kafkaErr).Msg("kafka consumer error")
if err := consumerBackoff(ctx, time.Second); err != nil {
stopErr = err
}
continue
}
stopErr = fmt.Errorf("error reading message: %w", err)
continue
}
laneKey := string(msg.Key)
if key != nil {
laneKey, err = key(msg)
if err != nil {
stopErr = fmt.Errorf("kafka: resolving handler key: %w", err)
continue
}
}
pk := messagePartition(msg)
p := partitions[pk]
if p == nil {
p = &partitionOffsets{done: map[ckf.Offset]bool{}}
partitions[pk] = p
}
p.pending = append(p.pending, msg)
pending++
queue := queues[messageLane(laneKey, lanes)]
dispatch:
for {
select {
case queue <- msg:
break dispatch
case r := <-results:
settle(r)
if stopErr != nil {
break dispatch
}
case <-ctx.Done():
stopErr = ctx.Err()
break dispatch
}
}
}
return shutdown()
}
var errLaneStopped = errors.New("kafka: lane stopped after an earlier failure")
func (cons *Consumer) storeSettled(partitions map[partitionKey]*partitionOffsets, msg *ckf.Message, epoch uint64, stopErr *error) {
p := partitions[messagePartition(msg)]
if p == nil {
return
}
last := p.complete(msg.TopicPartition.Offset)
if last == nil {
return
}
if err := cons.storeMessageAt(last, epoch); err != nil && *stopErr == nil {
if errors.Is(err, ErrClientClosed) || errors.Is(err, ErrAssignmentLost) {
*stopErr = err
return
}
*stopErr = fmt.Errorf("kafka: storing a handled offset failed: %w", err)
}
}
@@ -0,0 +1,234 @@
//go:build kafka
package kafka
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
ckf "github.com/confluentinc/confluent-kafka-go/v2/kafka"
)
type concurrentConsumer struct {
shutdownConsumer
mu sync.Mutex
queue []*ckf.Message
stored []*ckf.Message
reads atomic.Int64
acked chan struct{}
onRead func()
}
func (c *concurrentConsumer) ReadMessage(time.Duration) (*ckf.Message, error) {
if c.onRead != nil {
c.onRead()
}
c.mu.Lock()
if len(c.queue) > 0 {
msg := c.queue[0]
c.queue = c.queue[1:]
c.reads.Add(1)
c.mu.Unlock()
return msg, nil
}
c.mu.Unlock()
time.Sleep(time.Millisecond)
return nil, ckf.NewError(ckf.ErrTimedOut, "idle", false)
}
func (c *concurrentConsumer) StoreMessage(msg *ckf.Message) ([]ckf.TopicPartition, error) {
c.mu.Lock()
defer c.mu.Unlock()
c.stored = append(c.stored, msg)
if c.acked != nil {
select {
case c.acked <- struct{}{}:
default:
}
}
return nil, nil
}
func concurrentMessage(offset int, partition int32, mailbox string) *ckf.Message {
topic := "commands"
return &ckf.Message{
TopicPartition: ckf.TopicPartition{Topic: &topic, Partition: partition, Offset: ckf.Offset(offset)},
Key: []byte("legacy-task-key"),
Value: []byte(mailbox),
}
}
func mailboxKey(msg *ckf.Message) (string, error) { return string(msg.Value), nil }
func awaitConcurrent(t *testing.T, ch <-chan struct{}) {
t.Helper()
select {
case <-ch:
case <-time.After(3 * time.Second):
t.Fatal("timed out waiting for concurrent handler")
}
}
func TestConcurrentConsumerOrdersMailboxAndFencesOffsets(t *testing.T) {
if messageLane("a", 2) == messageLane("b", 2) {
t.Fatal("test requires separate lanes")
}
c := &concurrentConsumer{queue: []*ckf.Message{
concurrentMessage(10, 0, "a"), concurrentMessage(12, 0, "b"), concurrentMessage(18, 0, "a"),
}, acked: make(chan struct{}, 5)}
cons := &Consumer{c: c}
ctx, cancel := context.WithCancel(t.Context())
defer cancel()
release := make(chan struct{})
started, independent, later := make(chan struct{}), make(chan struct{}), make(chan struct{})
done := make(chan error, 1)
go func() {
done <- cons.ConsumeConcurrent(ctx, 2, mailboxKey, func(msg *ckf.Message) error {
switch msg.TopicPartition.Offset {
case 10:
close(started)
<-release
case 12:
close(independent)
case 18:
close(later)
}
return nil
})
}()
awaitConcurrent(t, started)
awaitConcurrent(t, independent)
select {
case <-later:
t.Fatal("same mailbox ran before earlier handler completed")
default:
}
c.mu.Lock()
count := len(c.stored)
c.mu.Unlock()
if count != 0 {
t.Fatal("stored an offset beyond an unfinished command")
}
close(release)
awaitConcurrent(t, later)
for {
awaitConcurrent(t, c.acked)
c.mu.Lock()
last := c.stored[len(c.stored)-1].TopicPartition.Offset
c.mu.Unlock()
if last == 18 {
break
}
}
cancel()
if err := <-done; !errors.Is(err, context.Canceled) {
t.Fatalf("shutdown: %v", err)
}
}
func TestConcurrentConsumerFailureKeepsPartitionReplayAndStopsLane(t *testing.T) {
c := &concurrentConsumer{queue: []*ckf.Message{
concurrentMessage(0, 0, "a"), concurrentMessage(1, 0, "b"),
concurrentMessage(2, 0, "a"), concurrentMessage(3, 1, "b"),
}}
cons := &Consumer{c: c, retryLimit: 1}
independent := make(chan struct{})
failure := errors.New("provider command failed")
var later atomic.Bool
err := cons.ConsumeConcurrent(t.Context(), 2, mailboxKey, func(msg *ckf.Message) error {
if msg.TopicPartition.Offset == 0 {
<-independent
return failure
}
if msg.TopicPartition.Offset == 1 {
close(independent)
}
if msg.TopicPartition.Offset == 2 {
later.Store(true)
}
return nil
})
if !errors.Is(err, failure) || later.Load() {
t.Fatalf("err=%v later=%v", err, later.Load())
}
for _, msg := range c.stored {
if msg.TopicPartition.Partition == 0 {
t.Fatal("advanced past failed partition offset")
}
}
}
func TestConcurrentConsumerBoundsReadAheadBehindSlowMailbox(t *testing.T) {
c := &concurrentConsumer{acked: make(chan struct{}, 100)}
c.queue = append(c.queue, concurrentMessage(0, 0, "a"))
for i := 1; i < 200; i++ {
c.queue = append(c.queue, concurrentMessage(i, 0, "b"))
}
ctx, cancel := context.WithCancel(t.Context())
defer cancel()
release := make(chan struct{})
done := make(chan error, 1)
var handled atomic.Int64
go func() {
done <- (&Consumer{c: c}).ConsumeConcurrent(ctx, 2, mailboxKey, func(msg *ckf.Message) error {
if msg.TopicPartition.Offset == 0 {
<-release
}
handled.Add(1)
return nil
})
}()
deadline := time.After(3 * time.Second)
for handled.Load() < 2*(laneBuffer+1)-1 {
select {
case <-deadline:
t.Fatal("did not process independent messages")
default:
time.Sleep(time.Millisecond)
}
}
if read := c.reads.Load(); read > 2*(laneBuffer+1) {
t.Fatalf("unbounded read-ahead: %d", read)
}
cancel()
close(release)
if err := <-done; !errors.Is(err, context.Canceled) {
t.Fatalf("shutdown: %v", err)
}
}
func TestConcurrentConsumerRejectsOffsetsFromRevokedAssignment(t *testing.T) {
c := &concurrentConsumer{queue: []*ckf.Message{concurrentMessage(0, 0, "a")}}
cons := &Consumer{c: c}
released := make(chan struct{})
var revoked atomic.Bool
c.onRead = func() {
if revoked.Load() {
cons.assignmentEpoch.Add(1)
close(released)
c.onRead = nil
}
}
err := cons.ConsumeConcurrent(t.Context(), 2, mailboxKey, func(*ckf.Message) error {
revoked.Store(true)
<-released
return nil
})
if !errors.Is(err, ErrAssignmentLost) || len(c.stored) != 0 {
t.Fatalf("err=%v stored=%v", err, c.stored)
}
}
func TestPartitionOffsetsTracksReadOrderNotNumericAdjacency(t *testing.T) {
p := &partitionOffsets{pending: []*ckf.Message{concurrentMessage(4, 0, "a"), concurrentMessage(9, 0, "b")}, done: map[ckf.Offset]bool{}}
if msg := p.complete(9); msg != nil {
t.Fatal("completed past unfinished predecessor")
}
if msg := p.complete(4); msg == nil || msg.TopicPartition.Offset != 9 || len(p.pending) != 0 {
t.Fatalf("last=%v pending=%v", msg, p.pending)
}
}
@@ -80,6 +80,32 @@ func TestConsumerRetriesFailedRecordBeforeLaterOffsets(t *testing.T) {
}
}
func TestConsumerRetryBudgetAllowsUnattemptedSendResultBeforeLaterReload(t *testing.T) {
first := &ckf.Message{TopicPartition: ckf.TopicPartition{Partition: 0, Offset: 10}}
next := &ckf.Message{TopicPartition: ckf.TopicPartition{Partition: 0, Offset: 11}}
c := &orderedConsumer{queue: []*ckf.Message{first, next}}
ctx, cancel := context.WithCancel(t.Context())
defer cancel()
attempts := 0
err := (&Consumer{c: c}).Consume(ctx, func(msg *ckf.Message) error {
if msg == first {
attempts++
if attempts < 5 {
return errors.New("mailbox not loaded")
}
return nil
}
if len(c.stored) != 1 || c.stored[0] != 10 || attempts != 5 {
t.Fatal("reload ran before resolving the unattempted send")
}
cancel()
return nil
})
if !errors.Is(err, context.Canceled) || attempts != 5 || len(c.stored) != 2 {
t.Fatalf("attempts=%d stored=%v err=%v", attempts, c.stored, err)
}
}
func TestConsumerFailureCancellationNeverStoresOffset(t *testing.T) {
c := &orderedConsumer{queue: []*ckf.Message{{}}}
ctx, cancel := context.WithCancel(context.Background())
+64 -8
View File
@@ -3,6 +3,8 @@
package kafka
import (
"context"
"fmt"
"sync"
"time"
@@ -17,10 +19,11 @@ type producerClient interface {
}
type Producer struct {
p producerClient
Avrov2 *Avrov2
mu sync.Mutex
closed bool
p producerClient
Avrov2 *Avrov2
mu sync.Mutex
closed bool
eventsDrained <-chan struct{}
}
type ProducerConfig struct {
@@ -31,6 +34,7 @@ func NewProducer(servers string) *ProducerConfig {
return &ProducerConfig{
config: map[string]ckf.ConfigValue{
"bootstrap.servers": servers,
"acks": "all",
},
}
}
@@ -57,8 +61,24 @@ func (conf *ProducerConfig) Connect() (*Producer, error) {
return nil, err
}
eventsDrained := make(chan struct{})
go func() {
defer close(eventsDrained)
for event := range p.Events() {
switch event := event.(type) {
case *ckf.Message:
if event != nil && event.TopicPartition.Error != nil {
log.Warn().Err(event.TopicPartition.Error).Msg("kafka background delivery failed")
}
case ckf.Error:
log.Debug().Str("code", event.Code().String()).Msg("kafka producer event")
}
}
}()
return &Producer{
p: p,
p: p,
eventsDrained: eventsDrained,
}, nil
}
@@ -74,11 +94,14 @@ func (pr *Producer) Close() {
}
pr.closed = true
pr.mu.Unlock()
undelivered := pr.p.Flush(30_000) // 30 seconds
if undelivered > 0 {
log.Warn().Int("count", undelivered).Msg("messages not delivered during shutdown")
outstanding := pr.p.Flush(30_000) // 30 seconds
if outstanding > 0 {
log.Warn().Int("count", outstanding).Msg("kafka producer has unflushed events during shutdown")
}
pr.p.Close()
if pr.eventsDrained != nil {
<-pr.eventsDrained
}
}
func (pr *Producer) Produce(topic string, key, value []byte) error {
@@ -94,3 +117,36 @@ func (pr *Producer) Produce(topic string, key, value []byte) error {
Timestamp: time.Now(),
}, nil)
}
// ProduceConfirmed waits for broker delivery, not merely local queue admission.
func (pr *Producer) ProduceConfirmed(ctx context.Context, topic string, key, value []byte) error {
if err := ctx.Err(); err != nil {
return err
}
report := make(chan ckf.Event, 1)
pr.mu.Lock()
if pr.closed {
pr.mu.Unlock()
return ErrClientClosed
}
err := pr.p.Produce(&ckf.Message{
TopicPartition: ckf.TopicPartition{Topic: &topic, Partition: ckf.PartitionAny},
Key: key,
Value: value,
Timestamp: time.Now(),
}, report)
pr.mu.Unlock()
if err != nil {
return err
}
select {
case <-ctx.Done():
return ctx.Err()
case event := <-report:
msg, ok := event.(*ckf.Message)
if !ok || msg == nil {
return fmt.Errorf("kafka: unexpected delivery report %T", event)
}
return msg.TopicPartition.Error
}
}
@@ -0,0 +1,160 @@
//go:build kafka
package kafka
import (
"context"
"errors"
"testing"
"time"
ckf "github.com/confluentinc/confluent-kafka-go/v2/kafka"
)
type queuedPublication struct {
msg *ckf.Message
report chan ckf.Event
}
type deliveryProducer struct {
queued chan queuedPublication
err error
}
func (p *deliveryProducer) Produce(msg *ckf.Message, report chan ckf.Event) error {
if p.err != nil {
return p.err
}
p.queued <- queuedPublication{msg, report}
return nil
}
func (*deliveryProducer) Flush(int) int { return 0 }
func (*deliveryProducer) Close() {}
func confirmedPublication(t *testing.T, p *Producer, ctx context.Context) <-chan error {
t.Helper()
done := make(chan error, 1)
go func() { done <- p.ProduceConfirmed(ctx, "jobs", []byte("mailbox"), []byte("result")) }()
return done
}
func waitDelivery(t *testing.T, done <-chan error) error {
t.Helper()
select {
case err := <-done:
return err
case <-time.After(2 * time.Second):
t.Fatal("publication did not settle")
return nil
}
}
func waitEnqueue(t *testing.T, p *deliveryProducer, done <-chan error) queuedPublication {
t.Helper()
select {
case queued := <-p.queued:
return queued
case err := <-done:
t.Fatalf("publication returned before enqueue: %v", err)
case <-time.After(2 * time.Second):
t.Fatal("publication did not enter the native queue")
}
return queuedPublication{}
}
func TestProducerConfirmedPublicationWaitsForIndependentDeliveryReports(t *testing.T) {
p := &deliveryProducer{queued: make(chan queuedPublication, 2)}
pr := &Producer{p: p}
first := confirmedPublication(t, pr, t.Context())
a := waitEnqueue(t, p, first)
second := confirmedPublication(t, pr, t.Context())
b := waitEnqueue(t, p, second)
select {
case err := <-first:
t.Fatalf("local admission was treated as delivered: %v", err)
default:
}
b.report <- b.msg
if err := waitDelivery(t, second); err != nil {
t.Fatal(err)
}
failed := errors.New("broker rejected result")
a.msg.TopicPartition.Error = failed
a.report <- a.msg
if err := waitDelivery(t, first); !errors.Is(err, failed) {
t.Fatalf("delivery failure was lost: %v", err)
}
if NewProducer("broker").config["acks"] != "all" {
t.Fatal("producer did not request broker acknowledgement")
}
}
func TestProducerConfirmedPublicationCancellationAllowsLateReport(t *testing.T) {
p := &deliveryProducer{queued: make(chan queuedPublication, 1)}
ctx, cancel := context.WithCancel(t.Context())
done := confirmedPublication(t, &Producer{p: p}, ctx)
a := waitEnqueue(t, p, done)
cancel()
if err := waitDelivery(t, done); !errors.Is(err, context.Canceled) {
t.Fatal("canceled publication was acknowledged", err)
}
select {
case a.report <- a.msg:
default:
t.Fatal("late delivery report blocks native producer")
}
if err := (&Producer{p: p}).ProduceConfirmed(ctx, "jobs", nil, nil); !errors.Is(err, context.Canceled) {
t.Fatal("already canceled context was published", err)
}
if len(p.queued) != 0 {
t.Fatal("canceled publication entered native queue")
}
}
func TestProducerConfirmedPublicationRejectsEnqueueFailureOrClosedClient(t *testing.T) {
failed := errors.New("local queue full")
for _, pr := range []*Producer{{p: &deliveryProducer{err: failed}}, {closed: true}} {
want := failed
if pr.closed {
want = ErrClientClosed
}
if err := pr.ProduceConfirmed(t.Context(), "jobs", nil, nil); !errors.Is(err, want) {
t.Fatalf("want %v, got %v", want, err)
}
}
}
func TestProducerConfirmedPublicationReportsNativeTimeout(t *testing.T) {
conf := NewProducer("127.0.0.1:1")
conf.Set("message.timeout.ms", 100)
conf.Set("socket.timeout.ms", 100)
conf.Set("log_level", 0)
pr, err := conf.Connect()
if err != nil {
t.Fatal(err)
}
t.Cleanup(pr.Close)
ctx, cancel := context.WithTimeout(t.Context(), 3*time.Second)
defer cancel()
err = pr.ProduceConfirmed(ctx, "jobs", nil, []byte("result"))
var native ckf.Error
if !errors.As(err, &native) || native.Code() != ckf.ErrMsgTimedOut {
t.Fatalf("native undelivered result treated as success: %v", err)
}
closed := make(chan struct{})
go func() {
pr.Close()
close(closed)
}()
select {
case <-closed:
case <-time.After(5 * time.Second):
t.Fatal("generic producer events prevented shutdown after delivery settled")
}
select {
case <-pr.eventsDrained:
default:
t.Fatal("producer event drain outlived shutdown")
}
}
+4 -1
View File
@@ -107,7 +107,10 @@ func (r *warmupRecoveryRepository) EnqueueFiling(ctx context.Context, action mod
}
func (r *warmupRecoveryRepository) ClaimFilings(ctx context.Context, limit int) ([]models.WarmupEmailAction, error) {
rows, err := r.db.Query(ctx, `UPDATE warmup_pending_filings SET next_attempt_at = NOW() + INTERVAL '5 minutes'
rows, err := r.db.Query(ctx, `UPDATE warmup_pending_filings SET next_attempt_at = NOW() + CASE
WHEN created_at > NOW() - INTERVAL '20 minutes' THEN INTERVAL '5 minutes'
WHEN created_at > NOW() - INTERVAL '1 hour' THEN INTERVAL '15 minutes'
ELSE INTERVAL '30 minutes' END
WHERE id IN (SELECT f.id FROM warmup_pending_filings f
JOIN email_accounts a ON a.id = f.email_account_id
WHERE f.next_attempt_at <= NOW() AND a.worker_id IS NOT NULL
@@ -27,10 +27,15 @@ const unstartedWarmupDispatch = `t.task_type='warmup' AND t.status='completed'
// Only nonce-gated commands that never began execution can be safely retired.
func (r *taskRepository) RecoverUnstartedWarmupDispatches(ctx context.Context, before time.Time, limit int) (int, error) {
recovered, _, err := r.RecoverUnstartedWarmupDispatchBatch(ctx, before, limit)
return recovered, err
}
func (r *taskRepository) RecoverUnstartedWarmupDispatchBatch(ctx context.Context, before time.Time, limit int) (recovered, selected int, err error) {
rows, err := r.db.Query(ctx, `SELECT t.id,t.email_account_id FROM tasks t JOIN warmup_tasks w ON w.task_id=t.id
WHERE `+unstartedWarmupDispatch+` ORDER BY COALESCE(t.send_reserved_at,t.completed_at),t.id LIMIT $2`, before, limit)
if err != nil {
return 0, err
return 0, 0, err
}
type candidate struct{ task, mailbox uuid.UUID }
var candidates []candidate
@@ -38,26 +43,26 @@ func (r *taskRepository) RecoverUnstartedWarmupDispatches(ctx context.Context, b
var c candidate
if err = rows.Scan(&c.task, &c.mailbox); err != nil {
rows.Close()
return 0, err
return 0, 0, err
}
candidates = append(candidates, c)
}
err = rows.Err()
rows.Close()
if err != nil {
return 0, err
return 0, 0, err
}
recovered := 0
selected = len(candidates)
for _, c := range candidates {
ok, err := r.retireUnstartedWarmupDispatch(ctx, c.task, c.mailbox, before)
if err != nil {
return recovered, err
return recovered, selected, err
}
if ok {
recovered++
}
}
return recovered, nil
return recovered, selected, nil
}
func (r *taskRepository) retireUnstartedWarmupDispatch(ctx context.Context, task, mailbox uuid.UUID, before time.Time) (bool, error) {
@@ -173,6 +173,30 @@ func TestLiveWarmupFilingRequiresAcknowledgement(t *testing.T) {
if claimed, err := repo.ClaimFilings(ctx, 100); err != nil || len(claimed) != 0 {
t.Fatalf("lease did not prevent duplicate dispatch: %+v %v", claimed, err)
}
for _, backoff := range []struct {
age time.Duration
delay time.Duration
}{{10 * time.Minute, 5 * time.Minute}, {25 * time.Minute, 15 * time.Minute}, {2 * time.Hour, 30 * time.Minute}} {
if _, err := pool.Exec(ctx, `UPDATE warmup_pending_filings SET created_at=$2,next_attempt_at=NOW()-INTERVAL '1 second' WHERE id=$1`, id, time.Now().Add(-backoff.age)); err != nil {
t.Fatal(err)
}
if _, err := repo.EnqueueFiling(ctx, action); err != nil {
t.Fatal(err)
}
if claimed, err := repo.ClaimFilings(ctx, 100); err != nil || len(claimed) != 1 || claimed[0].FilingID != id.String() {
t.Fatalf("durable retry was lost: %+v %v", claimed, err)
}
var next time.Time
if err := pool.QueryRow(ctx, `SELECT next_attempt_at FROM warmup_pending_filings WHERE id=$1`, id).Scan(&next); err != nil {
t.Fatal(err)
}
if delay := time.Until(next); delay < backoff.delay-time.Minute || delay > backoff.delay+time.Minute {
t.Fatalf("age=%v delay=%v want=%v", backoff.age, delay, backoff.delay)
}
if claimed, err := repo.ClaimFilings(ctx, 100); err != nil || len(claimed) != 0 {
t.Fatalf("old filing flooded the retry loop: %+v %v", claimed, err)
}
}
var waiting bool
if err := pool.QueryRow(ctx, `SELECT next_attempt_at < NOW() - INTERVAL '30 minutes' FROM warmup_pending_filings WHERE id = $1`, waitingID).Scan(&waiting); err != nil || !waiting {
t.Fatalf("unassigned backlog was leased instead of remaining pending: %v %v", waiting, err)
+29 -4
View File
@@ -23,13 +23,35 @@ const warmupReconcileBatch = 500
// campaign does not itself enqueue a task, so without this pass a freshly
// enabled mailbox would never start warming.
func (s *tasksService) ReconcileWarmupSchedules(ctx context.Context, limit int) (int, error) {
var recoverBatch func(context.Context, time.Time, int) (int, int, error)
if recovery, ok := s.taskRepo.(interface {
RecoverUnstartedWarmupDispatchBatch(context.Context, time.Time, int) (int, int, error)
}); ok {
recoverBatch = recovery.RecoverUnstartedWarmupDispatchBatch
} else if recovery, ok := s.taskRepo.(interface {
RecoverUnstartedWarmupDispatches(context.Context, time.Time, int) (int, error)
}); ok {
recovered, err := recovery.RecoverUnstartedWarmupDispatches(ctx, time.Now().Add(-10*time.Minute), limit)
if err != nil {
log.Warn().Err(err).Msg("warmup reconcile failed to retire unstarted dispatches")
} else if recovered > 0 {
recoverBatch = func(ctx context.Context, before time.Time, limit int) (int, int, error) {
n, err := recovery.RecoverUnstartedWarmupDispatches(ctx, before, limit)
return n, n, err
}
}
if recoverBatch != nil {
recovered := 0
for batch := 0; batch < 10 && ctx.Err() == nil; batch++ {
n, selected, err := recoverBatch(ctx, time.Now().Add(-10*time.Minute), limit)
recovered += n
if err != nil {
if recovered > 0 {
log.Info().Int("recovered", recovered).Msg("warmup reconcile retired unstarted dispatches before failing")
}
return 0, err
}
if selected < limit || limit <= 0 {
break
}
}
if recovered > 0 {
log.Info().Int("recovered", recovered).Msg("warmup reconcile retired unstarted dispatches")
}
}
@@ -95,6 +117,9 @@ func (s *tasksService) ReconcileWarmupSchedules(ctx context.Context, limit int)
// context is cancelled. Mirrors the other background sweeps (warmup health,
// dead-worker) and is started from the backend, which owns Cloud Tasks.
func (s *tasksService) StartWarmupReconciler(ctx context.Context, interval time.Duration) {
if interval <= 0 {
interval = time.Minute
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
@@ -0,0 +1,90 @@
package tasks
import (
"context"
"errors"
"testing"
"time"
"github.com/google/uuid"
"github.com/warmbly/warmbly/internal/repository"
)
type reconcileRecoveryRepo struct {
repository.TaskRepository
batches []int
calls int
err error
}
func (r *reconcileRecoveryRepo) RecoverUnstartedWarmupDispatches(_ context.Context, before time.Time, _ int) (int, error) {
if d := time.Since(before); d < 9*time.Minute || d > 11*time.Minute {
return 0, errors.New("changed the unstarted-send safety window")
}
r.calls++
if r.err != nil {
return 0, r.err
}
if len(r.batches) == 0 {
return 0, nil
}
n := r.batches[0]
r.batches = r.batches[1:]
return n, nil
}
type reconcileEmptyEmailRepo struct{ repository.EmailRepository }
type concurrentlyRetiredRecoveryRepo struct {
reconcileRecoveryRepo
}
func (r *concurrentlyRetiredRecoveryRepo) RecoverUnstartedWarmupDispatchBatch(_ context.Context, before time.Time, _ int) (int, int, error) {
if d := time.Since(before); d < 9*time.Minute || d > 11*time.Minute {
return 0, 0, errors.New("changed the safety window")
}
r.calls++
if r.calls == 1 {
return 0, 500, nil
}
return 50, 50, nil
}
func TestWarmupRecoveryDoesNotStopWhenOtherReconcilerRetiredSelectedBatch(t *testing.T) {
r := &concurrentlyRetiredRecoveryRepo{}
s := &tasksService{taskRepo: r, emailRepo: &reconcileEmptyEmailRepo{}}
if _, err := s.ReconcileWarmupSchedules(t.Context(), 500); err != nil || r.calls != 2 {
t.Fatalf("concurrent retirements cut recovery short: calls=%d err=%v", r.calls, err)
}
}
func (*reconcileEmptyEmailRepo) ListWarmupScheduleCandidates(context.Context, int) ([]uuid.UUID, error) {
return nil, nil
}
func TestWarmupRecoveryDrainsMultipleOldBatchesOnStartupWithinBound(t *testing.T) {
for _, test := range []struct {
name string
batches []int
want int
}{
{"existing fleet", []int{500, 500, 50}, 3},
{"bounded backlog", []int{500, 500, 500, 500, 500, 500, 500, 500, 500, 500, 500}, 10},
} {
t.Run(test.name, func(t *testing.T) {
r := &reconcileRecoveryRepo{batches: test.batches}
s := &tasksService{taskRepo: r, emailRepo: &reconcileEmptyEmailRepo{}}
if _, err := s.ReconcileWarmupSchedules(t.Context(), 500); err != nil || r.calls != test.want {
t.Fatalf("calls=%d want=%d err=%v", r.calls, test.want, err)
}
})
}
}
func TestWarmupRecoveryUnavailableDoesNotProceedAsSuccessfulRecovery(t *testing.T) {
want := errors.New("database unavailable")
s := &tasksService{taskRepo: &reconcileRecoveryRepo{err: want}}
if _, err := s.ReconcileWarmupSchedules(t.Context(), 500); !errors.Is(err, want) {
t.Fatalf("lost recovery failure: %v", err)
}
}