From dd4b405f8e6d0baa5d0cc1fe3aca41e28678bc88 Mon Sep 17 00:00:00 2001 From: Matthew Meszaros Date: Thu, 27 Aug 2026 20:34:55 -0700 Subject: [PATCH] feat: make the mailbox removal this PR adds actually reach the worker: deactivateAccount read the assignment off the row EmailRepository.Update returns, whose RETURNING list carries the mailbox as the dashboard sees it and no worker_id, so account.WorkerID was always nil, the guard fired every time and PublishRemoveEmail still had no caller in the codebase; it now asks GetWorkerID directly, the three deactivation handlers hand it the ids they already parsed, LoadAccountOntoWorker refuses to ship a mailbox that is not active so the reconciler cannot put back the mailbox the consumer just removed, and tests drive all three handlers end to end against a repository stub that withholds worker_id exactly like the real one --- .../app/consumer/deactivate_account_test.go | 223 ++++++++++++++++++ internal/app/consumer/event_email_error.go | 49 ++-- internal/app/email/loader.go | 10 +- internal/app/email/loader_status_test.go | 101 ++++++++ 4 files changed, 365 insertions(+), 18 deletions(-) create mode 100644 internal/app/consumer/deactivate_account_test.go create mode 100644 internal/app/email/loader_status_test.go diff --git a/internal/app/consumer/deactivate_account_test.go b/internal/app/consumer/deactivate_account_test.go new file mode 100644 index 00000000..8ab571f6 --- /dev/null +++ b/internal/app/consumer/deactivate_account_test.go @@ -0,0 +1,223 @@ +package jobs + +import ( + "context" + "errors" + "testing" + + "github.com/google/uuid" + "github.com/warmbly/warmbly/internal/errx" + "github.com/warmbly/warmbly/internal/events" + "github.com/warmbly/warmbly/internal/models" + "github.com/warmbly/warmbly/internal/repository" +) + +// stubEmailRepo answers the two calls deactivation makes. Embedding the +// interface keeps the stub to those two; anything else panics loudly rather +// than passing silently. +type stubEmailRepo struct { + repository.EmailRepository + + workerID *uuid.UUID + workerErr *errx.Error + updateErr *errx.Error + statusSet []string + updateCalls int + workerCalls int +} + +func (s *stubEmailRepo) Update(ctx context.Context, userID, emailAccountID string, udata *models.UpdateEmail) (*models.Email, *errx.Error) { + s.updateCalls++ + if udata.Status != nil { + s.statusSet = append(s.statusSet, *udata.Status) + } + if s.updateErr != nil { + return nil, s.updateErr + } + // Deliberately mirrors the real repository: its RETURNING list carries the + // mailbox as the dashboard sees it and no worker assignment. Reading + // WorkerID off this row is what made the removal unreachable the first + // time, so the stub must keep lying about it in exactly the same way. + id, _ := uuid.Parse(emailAccountID) + return &models.Email{ID: id, Status: "inactive"}, nil +} + +func (s *stubEmailRepo) GetWorkerID(ctx context.Context, emailAccountID uuid.UUID) (*uuid.UUID, *errx.Error) { + s.workerCalls++ + return s.workerID, s.workerErr +} + +// stubPublisher records the removals published to workers. +type stubPublisher struct { + events.Publisher + + err error + removed []removal +} + +type removal struct { + workerID uuid.UUID + userID string + emailID string +} + +func (p *stubPublisher) PublishRemoveEmail(ctx context.Context, workerID uuid.UUID, remove *models.RemoveWorkerEmail) error { + p.removed = append(p.removed, removal{workerID: workerID, userID: remove.UserID, emailID: remove.EmailID}) + return p.err +} + +func newDeactivationFixture(worker *uuid.UUID) (*JobsService, *stubEmailRepo, *stubPublisher) { + repo := &stubEmailRepo{workerID: worker} + pub := &stubPublisher{} + return &JobsService{EmailRepository: repo, Publisher: pub}, repo, pub +} + +func TestDeactivateAccountTellsTheWorkerToDropTheMailbox(t *testing.T) { + workerID := uuid.New() + userID := uuid.New() + emailID := uuid.New() + s, repo, pub := newDeactivationFixture(&workerID) + + s.deactivateAccount(context.Background(), userID, emailID) + + if len(repo.statusSet) != 1 || repo.statusSet[0] != "inactive" { + t.Errorf("status writes = %v, want one \"inactive\"", repo.statusSet) + } + if len(pub.removed) != 1 { + t.Fatalf("published %d removals, want 1: the worker keeps syncing a dead mailbox until it restarts", len(pub.removed)) + } + got := pub.removed[0] + if got.workerID != workerID { + t.Errorf("removal sent to worker %s, want %s", got.workerID, workerID) + } + if got.userID != userID.String() || got.emailID != emailID.String() { + t.Errorf("removal carried user=%s email=%s, want user=%s email=%s", got.userID, got.emailID, userID, emailID) + } +} + +func TestDeactivateAccountSkipsAnUnassignedMailbox(t *testing.T) { + s, repo, pub := newDeactivationFixture(nil) + + s.deactivateAccount(context.Background(), uuid.New(), uuid.New()) + + if repo.updateCalls != 1 { + t.Errorf("update called %d times, want 1", repo.updateCalls) + } + if len(pub.removed) != 0 { + t.Errorf("published %d removals for a mailbox on no worker, want 0", len(pub.removed)) + } +} + +// A status write that failed means the mailbox is still active, so publishing +// the removal would strand it: loaded on no worker, active in the database, +// and only the reconciler's next pass to put it back. +func TestDeactivateAccountDoesNotRemoveWhenTheStatusWriteFails(t *testing.T) { + workerID := uuid.New() + s, repo, pub := newDeactivationFixture(&workerID) + repo.updateErr = errx.InternalError() + + s.deactivateAccount(context.Background(), uuid.New(), uuid.New()) + + if len(pub.removed) != 0 { + t.Errorf("published %d removals after a failed status write, want 0", len(pub.removed)) + } + if repo.workerCalls != 0 { + t.Errorf("looked up the assignment %d times after a failed status write, want 0", repo.workerCalls) + } +} + +func TestDeactivateAccountSurvivesAnAssignmentLookupFailure(t *testing.T) { + s, repo, pub := newDeactivationFixture(nil) + repo.workerErr = errx.ErrNotFound + + s.deactivateAccount(context.Background(), uuid.New(), uuid.New()) + + if len(pub.removed) != 0 { + t.Errorf("published %d removals on a failed lookup, want 0", len(pub.removed)) + } +} + +// A publish failure is logged, not fatal: the account is already inactive and +// the worker drops the mailbox on its next restart. +func TestDeactivateAccountToleratesAPublishFailure(t *testing.T) { + workerID := uuid.New() + s, _, pub := newDeactivationFixture(&workerID) + pub.err = errors.New("bus down") + + s.deactivateAccount(context.Background(), uuid.New(), uuid.New()) + + if len(pub.removed) != 1 { + t.Errorf("published %d removals, want 1 attempt", len(pub.removed)) + } +} + +// The consumer runs with pieces unwired in tests and in reduced deployments. +func TestDeactivateAccountWithNothingWired(t *testing.T) { + (&JobsService{}).deactivateAccount(context.Background(), uuid.New(), uuid.New()) + + repo := &stubEmailRepo{} + (&JobsService{EmailRepository: repo}).deactivateAccount(context.Background(), uuid.New(), uuid.New()) + if repo.updateCalls != 1 { + t.Errorf("update called %d times with no publisher wired, want 1", repo.updateCalls) + } + if repo.workerCalls != 0 { + t.Errorf("looked up the assignment %d times with no publisher wired, want 0", repo.workerCalls) + } +} + +// Every worker-raised deactivation has to reach the worker, not just the one +// that happened to be debugged. These run the real handlers end to end. +func TestEveryDeactivationPathRemovesTheMailboxFromItsWorker(t *testing.T) { + handlers := map[string]func(*JobsService, context.Context, models.EmailErrorEvent) error{ + "auth error": (*JobsService).HandleEmailAuthError, + "disabled": (*JobsService).HandleEmailDisabled, + "rate limited": (*JobsService).HandleEmailRateLimited, + } + + for name, handle := range handlers { + t.Run(name, func(t *testing.T) { + workerID := uuid.New() + userID := uuid.New() + emailID := uuid.New() + s, repo, pub := newDeactivationFixture(&workerID) + + event := models.EmailErrorEvent{ + EmailAccountID: emailID.String(), + UserID: userID.String(), + ErrorCode: "AUTHENTICATION_FAILED", + ErrorType: "critical", + Message: "the grant is gone", + } + if err := handle(s, context.Background(), event); err != nil { + t.Fatalf("handler returned %v", err) + } + + if len(repo.statusSet) != 1 || repo.statusSet[0] != "inactive" { + t.Errorf("status writes = %v, want one \"inactive\"", repo.statusSet) + } + if len(pub.removed) != 1 { + t.Fatalf("published %d removals, want 1", len(pub.removed)) + } + if pub.removed[0].workerID != workerID || pub.removed[0].emailID != emailID.String() { + t.Errorf("removal = %+v, want worker %s and mailbox %s", pub.removed[0], workerID, emailID) + } + }) + } +} + +// A malformed event must not be answered with a status write or a removal. +func TestHandlersRejectAMalformedEventBeforeDeactivating(t *testing.T) { + workerID := uuid.New() + s, repo, pub := newDeactivationFixture(&workerID) + + err := s.HandleEmailAuthError(context.Background(), models.EmailErrorEvent{ + EmailAccountID: "not-a-uuid", + UserID: uuid.New().String(), + }) + if err == nil { + t.Error("a malformed email account id was accepted") + } + if repo.updateCalls != 0 || len(pub.removed) != 0 { + t.Errorf("acted on a malformed event: %d updates, %d removals", repo.updateCalls, len(pub.removed)) + } +} diff --git a/internal/app/consumer/event_email_error.go b/internal/app/consumer/event_email_error.go index 666c8cb3..086deaa2 100644 --- a/internal/app/consumer/event_email_error.go +++ b/internal/app/consumer/event_email_error.go @@ -57,7 +57,7 @@ func (s *JobsService) HandleEmailAuthError(ctx context.Context, event models.Ema } // Mark email account as needing re-auth (set status to inactive) - s.deactivateAccount(ctx, event.UserID, event.EmailAccountID) + s.deactivateAccount(ctx, userID, emailAccountID) // Send Pub/Sub notification to user if s.StreamingPublisher != nil && event.UserVisible { @@ -122,7 +122,7 @@ func (s *JobsService) HandleEmailDisabled(ctx context.Context, event models.Emai } // Mark email account as inactive - s.deactivateAccount(ctx, event.UserID, event.EmailAccountID) + s.deactivateAccount(ctx, userID, emailAccountID) // Send Pub/Sub notification to user if s.StreamingPublisher != nil && event.UserVisible { @@ -187,7 +187,7 @@ func (s *JobsService) HandleEmailRateLimited(ctx context.Context, event models.E } // Mark email account as inactive (terminated due to abuse) - s.deactivateAccount(ctx, event.UserID, event.EmailAccountID) + s.deactivateAccount(ctx, userID, emailAccountID) if s.WarmupService != nil { health, _ := s.WarmupService.ApplyRateLimitExceeded(ctx, emailAccountID, "worker sync/email rate limit exceeded") @@ -280,32 +280,47 @@ func ptrString(s string) *string { return &s } -// deactivateAccount marks a mailbox inactive AND tells the worker holding it to -// drop it. The per-worker load filters on status, but that only decides what a -// worker picks up when it starts: without this the running worker keeps syncing -// a mailbox that has just been disabled, failing every pass and writing an -// email_account_errors row each time, until someone restarts it. -func (s *JobsService) deactivateAccount(ctx context.Context, userID, emailAccountID string) { +// deactivateAccount marks a mailbox inactive AND tells the worker holding it +// to drop it. The per-worker load filters on status, but that only decides what +// a worker picks up when it starts, and the send path relays an account-level +// failure without stopping the mailbox's own sync loop: without this the worker +// keeps syncing a mailbox that has just been deactivated, failing every pass +// and writing an email_account_errors row each time, until someone restarts it. +func (s *JobsService) deactivateAccount(ctx context.Context, userID, emailAccountID uuid.UUID) { if s.EmailRepository == nil { return } inactive := "inactive" - account, xerr := s.EmailRepository.Update(ctx, userID, emailAccountID, &models.UpdateEmail{ + if _, xerr := s.EmailRepository.Update(ctx, userID.String(), emailAccountID.String(), &models.UpdateEmail{ Status: &inactive, - }) - if xerr != nil { + }); xerr != nil { log.Error().Str("error", xerr.Message).Msg("Failed to update email account status") return } - if s.Publisher == nil || account == nil || account.WorkerID == nil { + if s.Publisher == nil { return } - if err := s.Publisher.PublishRemoveEmail(ctx, *account.WorkerID, &models.RemoveWorkerEmail{ - UserID: userID, - EmailID: emailAccountID, + + // Asked for on its own rather than read off the row Update returned: that + // row is the mailbox as the dashboard sees it and carries no assignment. + workerID, xerr := s.EmailRepository.GetWorkerID(ctx, emailAccountID) + if xerr != nil { + log.Warn(). + Str("error", xerr.Message). + Str("email_account_id", emailAccountID.String()). + Msg("Cannot tell the worker to drop the deactivated mailbox: assignment lookup failed") + return + } + if workerID == nil { + return + } + if err := s.Publisher.PublishRemoveEmail(ctx, *workerID, &models.RemoveWorkerEmail{ + UserID: userID.String(), + EmailID: emailAccountID.String(), }); err != nil { log.Warn().Err(err). - Str("email_account_id", emailAccountID). + Str("email_account_id", emailAccountID.String()). + Str("worker_id", workerID.String()). Msg("Failed to tell the worker to drop the deactivated mailbox") } } diff --git a/internal/app/email/loader.go b/internal/app/email/loader.go index 609e03ea..5f491e4b 100644 --- a/internal/app/email/loader.go +++ b/internal/app/email/loader.go @@ -119,7 +119,8 @@ func (s *emailService) loadAccountBestEffort(ctx context.Context, accountID uuid // LoadAccountOntoWorker assigns a worker if the account has none, rebuilds the // account's decrypted credentials into an AddWorkerEmail payload, and publishes -// it so the worker loads the account into memory. Safe to call repeatedly. +// it so the worker loads the account into memory. Safe to call repeatedly, and +// a no-op for a mailbox that is not active. func (s *emailService) LoadAccountOntoWorker(ctx context.Context, accountID uuid.UUID) error { acc, xerr := s.emailRepository.GetByID(ctx, accountID) if xerr != nil { @@ -128,6 +129,13 @@ func (s *emailService) LoadAccountOntoWorker(ctx context.Context, accountID uuid if acc == nil { return nil } + // A mailbox that is not active must never be shipped to a worker. The + // reconciler reads the active list a tick before it publishes, so without + // this it can put back a mailbox that was deactivated in between and undo + // the removal the consumer just sent. + if acc.Status != "active" { + return nil + } workerID, rerr := s.releaseDeadWorker(ctx, acc.ID, acc.WorkerID) if rerr != nil { diff --git a/internal/app/email/loader_status_test.go b/internal/app/email/loader_status_test.go new file mode 100644 index 00000000..c1bef0c4 --- /dev/null +++ b/internal/app/email/loader_status_test.go @@ -0,0 +1,101 @@ +package email + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/warmbly/warmbly/internal/app/worker" + "github.com/warmbly/warmbly/internal/errx" + "github.com/warmbly/warmbly/internal/events" + "github.com/warmbly/warmbly/internal/models" + "github.com/warmbly/warmbly/internal/repository" +) + +// stubLoaderRepo serves one mailbox to the loader. +type stubLoaderRepo struct { + repository.EmailRepository + + account *models.Email +} + +func (s *stubLoaderRepo) GetByID(ctx context.Context, emailAccountID uuid.UUID) (*models.Email, *errx.Error) { + return s.account, nil +} + +// countingAssignment records placement attempts. Reaching it at all means the +// loader decided the mailbox belongs on a worker. +type countingAssignment struct { + worker.WorkerAssignmentService + + assigned int +} + +func (c *countingAssignment) IsWorkerLive(ctx context.Context, workerID uuid.UUID) (bool, error) { + return true, nil +} + +func (c *countingAssignment) AssignWorkerToEmail(ctx context.Context, emailAccountID, orgID uuid.UUID) (*uuid.UUID, error) { + c.assigned++ + id := uuid.New() + return &id, nil +} + +type countingPublisher struct { + events.Publisher + + added int +} + +func (p *countingPublisher) PublishAddEmail(ctx context.Context, workerID uuid.UUID, email *models.AddWorkerEmail) error { + p.added++ + return nil +} + +func loaderFor(status string) (*emailService, *countingAssignment, *countingPublisher) { + org := uuid.New() + assign := &countingAssignment{} + pub := &countingPublisher{} + repo := &stubLoaderRepo{account: &models.Email{ + ID: uuid.New(), + UserID: uuid.New().String(), + OrganizationID: &org, + Email: "sender@example.com", + Status: status, + }} + return &emailService{emailRepository: repo, workerAssignment: assign, publisher: pub}, assign, pub +} + +// The reconciler lists active mailboxes a tick before it publishes them, so a +// mailbox deactivated in between would otherwise be shipped straight back onto +// the worker the consumer just told to drop it. +func TestLoadAccountOntoWorkerSkipsAMailboxThatIsNotActive(t *testing.T) { + for _, status := range []string{"inactive", "revoked"} { + t.Run(status, func(t *testing.T) { + s, assign, pub := loaderFor(status) + + if err := s.LoadAccountOntoWorker(context.Background(), uuid.New()); err != nil { + t.Fatalf("LoadAccountOntoWorker: %v", err) + } + if assign.assigned != 0 { + t.Errorf("placed a %s mailbox on a worker (%d assignments)", status, assign.assigned) + } + if pub.added != 0 { + t.Errorf("shipped a %s mailbox to a worker (%d publishes)", status, pub.added) + } + }) + } +} + +// The guard must not stop the mailboxes that do belong on a worker: an active +// one still reaches placement. +func TestLoadAccountOntoWorkerStillPlacesAnActiveMailbox(t *testing.T) { + s, assign, _ := loaderFor("active") + + if err := s.LoadAccountOntoWorker(context.Background(), uuid.New()); err != nil { + t.Fatalf("LoadAccountOntoWorker: %v", err) + } + if assign.assigned != 1 { + t.Errorf("assignments = %d, want 1", assign.assigned) + } +}