diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ebd361748..7b91966c2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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) diff --git a/cmd/backend/main.go b/cmd/backend/main.go index 6822a6005..4dbd47c44 100644 --- a/cmd/backend/main.go +++ b/cmd/backend/main.go @@ -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, diff --git a/cmd/worker/main.go b/cmd/worker/main.go index 2fb4531c5..7819a77f0 100644 --- a/cmd/worker/main.go +++ b/cmd/worker/main.go @@ -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) } diff --git a/docs/content/docs/development/configuration.mdx b/docs/content/docs/development/configuration.mdx index ef81a3edb..85efa5bfb 100644 --- a/docs/content/docs/development/configuration.mdx +++ b/docs/content/docs/development/configuration.mdx @@ -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. + 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. @@ -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 | diff --git a/docs/content/docs/guides/warmup.mdx b/docs/content/docs/guides/warmup.mdx index 25fad0a52..06c71cb00 100644 --- a/docs/content/docs/guides/warmup.mdx +++ b/docs/content/docs/guides/warmup.mdx @@ -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. diff --git a/internal/app/worker/command_key.go b/internal/app/worker/command_key.go new file mode 100644 index 000000000..037df6f21 --- /dev/null +++ b/internal/app/worker/command_key.go @@ -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 +} diff --git a/internal/app/worker/command_key_avro_test.go b/internal/app/worker/command_key_avro_test.go new file mode 100644 index 000000000..ddcbf9005 --- /dev/null +++ b/internal/app/worker/command_key_avro_test.go @@ -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}) +} diff --git a/internal/app/worker/command_key_test.go b/internal/app/worker/command_key_test.go new file mode 100644 index 000000000..5a7972db1 --- /dev/null +++ b/internal/app/worker/command_key_test.go @@ -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) + } + } +} diff --git a/internal/app/worker/event_send_email.go b/internal/app/worker/event_send_email.go index 239d0f379..23fa6e8dc 100644 --- a/internal/app/worker/event_send_email.go +++ b/internal/app/worker/event_send_email.go @@ -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 } diff --git a/internal/app/worker/not_loaded_text_test.go b/internal/app/worker/not_loaded_text_test.go index 119b32313..cc9fc1e07 100644 --- a/internal/app/worker/not_loaded_text_test.go +++ b/internal/app/worker/not_loaded_text_test.go @@ -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) + } +} diff --git a/internal/events/publisher.go b/internal/events/publisher.go index ef651248e..b460690a9 100644 --- a/internal/events/publisher.go +++ b/internal/events/publisher.go @@ -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 diff --git a/internal/infrastructure/eventbus/bus.go b/internal/infrastructure/eventbus/bus.go index e8c63d68a..fa5ec1c24 100644 --- a/internal/infrastructure/eventbus/bus.go +++ b/internal/infrastructure/eventbus/bus.go @@ -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 diff --git a/internal/infrastructure/eventbus/kafka.go b/internal/infrastructure/eventbus/kafka.go index 427d14ed4..040a0dea2 100644 --- a/internal/infrastructure/eventbus/kafka.go +++ b/internal/infrastructure/eventbus/kafka.go @@ -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 { diff --git a/internal/infrastructure/eventbus/kafka_test.go b/internal/infrastructure/eventbus/kafka_test.go index 23752240b..f8c65087c 100644 --- a/internal/infrastructure/eventbus/kafka_test.go +++ b/internal/infrastructure/eventbus/kafka_test.go @@ -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() diff --git a/internal/infrastructure/kafka/consumer.go b/internal/infrastructure/kafka/consumer.go index c1e85883a..f7d03c1cf 100644 --- a/internal/infrastructure/kafka/consumer.go +++ b/internal/infrastructure/kafka/consumer.go @@ -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) + } +} diff --git a/internal/infrastructure/kafka/consumer_concurrent.go b/internal/infrastructure/kafka/consumer_concurrent.go new file mode 100644 index 000000000..b7e208e28 --- /dev/null +++ b/internal/infrastructure/kafka/consumer_concurrent.go @@ -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) + } +} diff --git a/internal/infrastructure/kafka/consumer_concurrent_test.go b/internal/infrastructure/kafka/consumer_concurrent_test.go new file mode 100644 index 000000000..b5ab00d83 --- /dev/null +++ b/internal/infrastructure/kafka/consumer_concurrent_test.go @@ -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) + } +} diff --git a/internal/infrastructure/kafka/consumer_retry_test.go b/internal/infrastructure/kafka/consumer_retry_test.go index 1ec234de3..16213dc83 100644 --- a/internal/infrastructure/kafka/consumer_retry_test.go +++ b/internal/infrastructure/kafka/consumer_retry_test.go @@ -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()) diff --git a/internal/infrastructure/kafka/producer.go b/internal/infrastructure/kafka/producer.go index 314237f7c..876c6b91c 100644 --- a/internal/infrastructure/kafka/producer.go +++ b/internal/infrastructure/kafka/producer.go @@ -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 + } +} diff --git a/internal/infrastructure/kafka/producer_delivery_test.go b/internal/infrastructure/kafka/producer_delivery_test.go new file mode 100644 index 000000000..c4d135b55 --- /dev/null +++ b/internal/infrastructure/kafka/producer_delivery_test.go @@ -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") + } +} diff --git a/internal/repository/pg_warmup_recovery.go b/internal/repository/pg_warmup_recovery.go index 5f15b3f2f..4c047aebc 100644 --- a/internal/repository/pg_warmup_recovery.go +++ b/internal/repository/pg_warmup_recovery.go @@ -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 diff --git a/internal/repository/warmup_dispatch_recovery.go b/internal/repository/warmup_dispatch_recovery.go index 2a8ea34d9..7a634eae7 100644 --- a/internal/repository/warmup_dispatch_recovery.go +++ b/internal/repository/warmup_dispatch_recovery.go @@ -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) { diff --git a/internal/repository/warmup_recovery_live_test.go b/internal/repository/warmup_recovery_live_test.go index 840d8c355..03beec3db 100644 --- a/internal/repository/warmup_recovery_live_test.go +++ b/internal/repository/warmup_recovery_live_test.go @@ -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) diff --git a/internal/tasks/warmup_reconciler.go b/internal/tasks/warmup_reconciler.go index 57f9ca18d..55bb3d68b 100644 --- a/internal/tasks/warmup_reconciler.go +++ b/internal/tasks/warmup_reconciler.go @@ -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() diff --git a/internal/tasks/warmup_reconciler_recovery_test.go b/internal/tasks/warmup_reconciler_recovery_test.go new file mode 100644 index 000000000..35f712a71 --- /dev/null +++ b/internal/tasks/warmup_reconciler_recovery_test.go @@ -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) + } +}