From 430a645858f1de0f1bd67fe1ff0bc9ea0bd2ba43 Mon Sep 17 00:00:00 2001 From: Matthew Meszaros Date: Mon, 24 Aug 2026 20:16:03 -0700 Subject: [PATCH] feat: stop the IMAP incremental pass fetching further batches once a sync lane is denied, so a mailbox unfrozen by the LIST-STATUS release fix does not walk its whole invisible backlog into the flood detector and deactivate itself: imapIncremental now returns not-complete on the first batch that could not be fully stored (holding the folder's mod-sequence for the next tick instead of setting stats.aborted, which would also skip the backfill and every other folder), ReleaseMailbox takes Client.mu like every other selected-state command so it cannot interleave with a warmup MOVE/STORE, every SELECT is funnelled through Client.selectMailbox so UNSELECT is skipped when nothing is selected and a strict server never answers BAD, and SmtpImapData.ImapClient plus WMail.gov become the narrow ImapConn and syncBudget interfaces so a full IMAP pass can be driven against a fake and its fetch round trips counted in TestImapSyncStopsFetchingOnceTheLiveLaneIsDenied, TestImapSyncKeepsWhatFitBeforeTheDenial and TestImapSyncWalksEveryBatchWithinBudget --- internal/app/worker/wmail/governor.go | 13 + internal/app/worker/wmail/imap_conn.go | 37 +++ internal/app/worker/wmail/sync_imap.go | 12 +- internal/app/worker/wmail/sync_imap_test.go | 238 ++++++++++++++++++ internal/app/worker/wmail/wmail.go | 9 +- internal/client/smtpimap/imap/client.go | 34 ++- .../client/smtpimap/imap/warmup_actions.go | 6 +- 7 files changed, 332 insertions(+), 17 deletions(-) create mode 100644 internal/app/worker/wmail/imap_conn.go create mode 100644 internal/app/worker/wmail/sync_imap_test.go diff --git a/internal/app/worker/wmail/governor.go b/internal/app/worker/wmail/governor.go index a4bd8e1d..11d754fe 100644 --- a/internal/app/worker/wmail/governor.go +++ b/internal/app/worker/wmail/governor.go @@ -38,6 +38,19 @@ type Admission struct { Until time.Time } +// syncBudget is the fair-use engine a sync pass charges. *governor is the +// only implementation; the interface exists so a pass can be driven against a +// fixed budget in tests, where the governor's Redis windows are unreachable. +type syncBudget interface { + Policy() models.SyncPolicy + SetPolicy(policy models.SyncPolicy) + Admit(ctx context.Context, lane SyncLane) Admission + ObserveLive(ctx context.Context, n int) bool + RecordThrottledDay(ctx context.Context) bool +} + +var _ syncBudget = (*governor)(nil) + // governor is the per-mailbox fair-use engine. Counters live in Redis (shared // across workers, so an organization budget holds even when its mailboxes sit // on different machines) as fixed windows: cheap INCRs with a TTL, no sorted diff --git a/internal/app/worker/wmail/imap_conn.go b/internal/app/worker/wmail/imap_conn.go new file mode 100644 index 00000000..a94d6461 --- /dev/null +++ b/internal/app/worker/wmail/imap_conn.go @@ -0,0 +1,37 @@ +package wmail + +import ( + "context" + "time" + + goimap "github.com/emersion/go-imap/v2" + "github.com/warmbly/warmbly/internal/client/smtpimap/imap" + "github.com/warmbly/warmbly/internal/errx" + "github.com/warmbly/warmbly/internal/models" +) + +// ImapConn is everything the worker drives on one IMAP connection. +// SmtpImapData holds this interface rather than *imap.Client so a sync pass +// can be run against a fake and its round trips counted; production always +// holds an *imap.Client. +type ImapConn interface { + // Sync pass. + Folders() ([]models.Mailbox, *errx.MailError) + ReleaseMailbox() + SelectForSync(mailbox string) (uint32, *errx.MailError) + SearchChangedSince(modSeq uint64) ([]goimap.UID, *errx.MailError) + SearchSince(since time.Time) ([]goimap.UID, *errx.MailError) + FetchEnvelopes(ctx context.Context, uids []goimap.UID) ([]*imap.Fetched, *errx.MailError) + FetchBody(f *imap.Fetched) + + // Send path. + AppendToSent(ctx context.Context, raw []byte, sentAt time.Time) error + + // Warmup actions. + MarkAsRead(ctx context.Context, mailboxName string, uid uint32) error + MarkImportant(ctx context.Context, mailboxName string, uid uint32) error + MoveToFolder(ctx context.Context, sourceMailbox, dstFolder string, uid uint32) error + RemoveFromSpam(ctx context.Context, sourceMailbox, inboxName string, uid uint32) error +} + +var _ ImapConn = (*imap.Client)(nil) diff --git a/internal/app/worker/wmail/sync_imap.go b/internal/app/worker/wmail/sync_imap.go index b7fecf3b..5b70589a 100644 --- a/internal/app/worker/wmail/sync_imap.go +++ b/internal/app/worker/wmail/sync_imap.go @@ -151,7 +151,6 @@ func (w *WMail) imapIncremental(ctx context.Context, box *models.Mailbox, modSeq // Newest first: when budget is short, the freshest mail lands first. sort.Slice(uids, func(i, j int) bool { return uids[i] > uids[j] }) - complete := true for lo := 0; lo < len(uids); lo += config.ImapFetchBatchSize { hi := min(lo+config.ImapFetchBatchSize, len(uids)) fetched, err := client.FetchEnvelopes(ctx, uids[lo:hi]) @@ -162,14 +161,15 @@ func (w *WMail) imapIncremental(ctx context.Context, box *models.Mailbox, modSeq if err != nil { return false, err } - if !done { - complete = false - } - if stats.aborted || ctx.Err() != nil { + // A denied lane means no later batch can be stored either, and every + // extra batch still counts new mail toward flood detection: a mailbox + // unfrozen on a long backlog would deactivate itself walking mail it + // cannot keep. Stop here; the held mod-sequence re-offers the rest. + if !done || stats.aborted || ctx.Err() != nil { return false, nil } } - return complete, nil + return true, nil } // imapApply routes one fetched batch: known messages get an UPDATE_EMAIL, diff --git a/internal/app/worker/wmail/sync_imap_test.go b/internal/app/worker/wmail/sync_imap_test.go new file mode 100644 index 00000000..11ce0862 --- /dev/null +++ b/internal/app/worker/wmail/sync_imap_test.go @@ -0,0 +1,238 @@ +package wmail + +import ( + "context" + "fmt" + "testing" + "time" + + goimap "github.com/emersion/go-imap/v2" + "github.com/google/uuid" + "github.com/warmbly/warmbly/internal/client/smtpimap/imap" + "github.com/warmbly/warmbly/internal/config" + "github.com/warmbly/warmbly/internal/errx" + "github.com/warmbly/warmbly/internal/models" +) + +// fakeImapConn is the sync pass's view of a server. Only the methods a pass +// calls are implemented; the embedded nil interface makes anything else panic +// rather than silently pass. +type fakeImapConn struct { + ImapConn + folders []models.Mailbox + changed []goimap.UID + fetches int + released int +} + +func (c *fakeImapConn) Folders() ([]models.Mailbox, *errx.MailError) { return c.folders, nil } + +func (c *fakeImapConn) ReleaseMailbox() { c.released++ } + +func (c *fakeImapConn) SelectForSync(string) (uint32, *errx.MailError) { + return uint32(len(c.changed)), nil +} + +func (c *fakeImapConn) SearchChangedSince(uint64) ([]goimap.UID, *errx.MailError) { + return append([]goimap.UID(nil), c.changed...), nil +} + +func (c *fakeImapConn) FetchEnvelopes(_ context.Context, uids []goimap.UID) ([]*imap.Fetched, *errx.MailError) { + c.fetches++ + out := make([]*imap.Fetched, 0, len(uids)) + for _, uid := range uids { + out = append(out, &imap.Fetched{Email: &models.EmailMessageData{ + UID: uint32(uid), + MessageID: fmt.Sprintf("<%d@fake.test>", uid), + Subject: "hello", + }}) + } + return out, nil +} + +func (c *fakeImapConn) FetchBody(*imap.Fetched) {} + +// fixedBudget is a syncBudget that admits a fixed number of messages and then +// denies on the daily window, which is how a real governor answers once the +// mailbox's day is spent. Redis is unreachable from a unit test, so the fake +// stands in for the counters, not for the decision the pass makes from them. +type fixedBudget struct { + allow int + admitted int + // observed totals what ObserveLive was told, so a test can check that a + // held backlog is not re-counted toward the flood threshold every pass. + observed int +} + +func (b *fixedBudget) Policy() models.SyncPolicy { return normalizePolicy(models.SyncPolicy{}) } +func (b *fixedBudget) SetPolicy(models.SyncPolicy) {} +func (b *fixedBudget) RecordThrottledDay(context.Context) bool { return false } + +func (b *fixedBudget) Admit(context.Context, SyncLane) Admission { + if b.admitted >= b.allow { + return Admission{Reason: models.SyncThrottleDaily, Until: time.Now().Add(time.Hour)} + } + b.admitted++ + return Admission{OK: true} +} + +func (b *fixedBudget) ObserveLive(_ context.Context, n int) bool { + b.observed += n + return false +} + +// newIMAPTestMail builds the smallest WMail that can run an IMAP pass. +func newIMAPTestMail(conn ImapConn, budget syncBudget, saved *models.Mailbox) (*WMail, *[]captured) { + var events []captured + w := &WMail{ + UserID: uuid.New(), + ID: uuid.New(), + Storage: fakeStore{}, + EmailMessageMapRepository: fakeMessageMap{}, + gov: budget, + SmtpImapData: &SmtpImapData{ + ImapClient: conn, + Mailboxes: []*models.Mailbox{saved}, + }, + } + w.onEvent = func(kind models.JobEventType, body any) error { + events = append(events, captured{eventType: kind, body: body}) + return nil + } + // The backfill is a separate lane with its own early return; keep it out + // of the way so these tests only exercise the live batch loop. + w.tracker = newSyncTracker( + &models.SyncState{BackfillStatus: models.SyncBackfillComplete}, + func(models.SyncState) error { return nil }, + ) + return w, &events +} + +func uidRange(n int) []goimap.UID { + out := make([]goimap.UID, 0, n) + for i := 1; i <= n; i++ { + out = append(out, goimap.UID(i)) + } + return out +} + +func relayedModSeq(t *testing.T, events []captured) uint64 { + t.Helper() + for i := len(events) - 1; i >= 0; i-- { + if events[i].eventType != models.JobEventTypeMailboxUpdate { + continue + } + return events[i].body.(*models.JobEventMailboxUpdate).Data.HighestModSeq + } + t.Fatal("no MAILBOX_UPDATE was relayed") + return 0 +} + +func hasEvent(events []captured, kind models.JobEventType) bool { + for _, e := range events { + if e.eventType == kind { + return true + } + } + return false +} + +// A mailbox unfrozen on a long backlog must not walk the whole thing: once +// the live lane is denied, the folder stops fetching and holds its +// mod-sequence, so the backlog is not re-offered to the flood detector batch +// after batch until the mailbox deactivates itself. +func TestImapSyncStopsFetchingOnceTheLiveLaneIsDenied(t *testing.T) { + conn := &fakeImapConn{ + folders: []models.Mailbox{{Name: "INBOX", UIDValidity: 7, HighestModSeq: 90_000}}, + changed: uidRange(3 * config.ImapFetchBatchSize), + } + budget := &fixedBudget{allow: 0} + w, events := newIMAPTestMail(conn, budget, &models.Mailbox{Name: "INBOX", UIDValidity: 7, HighestModSeq: 100}) + + if err := w.Sync(t.Context()); err != nil { + t.Fatalf("Sync: %v", err) + } + + if conn.fetches != 1 { + t.Errorf("fetched %d batches after the lane was denied, want 1", conn.fetches) + } + if got := w.SmtpImapData.Mailboxes[0].HighestModSeq; got != 100 { + t.Errorf("mod-sequence advanced to %d; a deferred backlog must hold it at 100", got) + } + if got := relayedModSeq(t, *events); got != 100 { + t.Errorf("relayed mod-sequence = %d, want the held 100", got) + } + if hasEvent(*events, models.JobEventTypeEmailRateLimited) { + t.Error("mailbox was deactivated by its own deferred backlog") + } + + // The next pass is re-offered the same mail. It must not read as a fresh + // flood: only messages the pass has never classified are observed. + seenAfterFirst := budget.observed + if seenAfterFirst != config.ImapFetchBatchSize { + t.Fatalf("observed %d new messages, want one batch", seenAfterFirst) + } + if err := w.Sync(t.Context()); err != nil { + t.Fatalf("second Sync: %v", err) + } + if budget.observed != seenAfterFirst { + t.Errorf("observed %d after the second pass, want %d: the held backlog was counted twice", + budget.observed, seenAfterFirst) + } + if conn.fetches != 2 { + t.Errorf("fetched %d batches over two passes, want 2", conn.fetches) + } +} + +// The denial stops the loop at the batch it happened in, not before it: +// everything admitted up to that point is stored, and only the rest waits. +func TestImapSyncKeepsWhatFitBeforeTheDenial(t *testing.T) { + conn := &fakeImapConn{ + folders: []models.Mailbox{{Name: "INBOX", UIDValidity: 7, HighestModSeq: 90_000}}, + changed: uidRange(3 * config.ImapFetchBatchSize), + } + budget := &fixedBudget{allow: config.ImapFetchBatchSize + 50} + w, events := newIMAPTestMail(conn, budget, &models.Mailbox{Name: "INBOX", UIDValidity: 7, HighestModSeq: 100}) + + if err := w.Sync(t.Context()); err != nil { + t.Fatalf("Sync: %v", err) + } + + if conn.fetches != 2 { + t.Errorf("fetched %d batches, want 2 (the full first and the one that ran out)", conn.fetches) + } + if budget.admitted != config.ImapFetchBatchSize+50 { + t.Errorf("admitted %d, want the whole budget spent", budget.admitted) + } + if got := w.SmtpImapData.Mailboxes[0].HighestModSeq; got != 100 { + t.Errorf("mod-sequence advanced to %d with mail still on the server", got) + } + if hasEvent(*events, models.JobEventTypeEmailRateLimited) { + t.Error("mailbox was deactivated by a plain budget denial") + } +} + +// The control case: with budget to spare the pass still walks every batch and +// the folder's mod-sequence moves to what the server reported. +func TestImapSyncWalksEveryBatchWithinBudget(t *testing.T) { + conn := &fakeImapConn{ + folders: []models.Mailbox{{Name: "INBOX", UIDValidity: 7, HighestModSeq: 90_000}}, + changed: uidRange(3 * config.ImapFetchBatchSize), + } + budget := &fixedBudget{allow: 10 * config.ImapFetchBatchSize} + w, _ := newIMAPTestMail(conn, budget, &models.Mailbox{Name: "INBOX", UIDValidity: 7, HighestModSeq: 100}) + + if err := w.Sync(t.Context()); err != nil { + t.Fatalf("Sync: %v", err) + } + + if conn.fetches != 3 { + t.Errorf("fetched %d batches, want 3", conn.fetches) + } + if got := w.SmtpImapData.Mailboxes[0].HighestModSeq; got != 90_000 { + t.Errorf("mod-sequence = %d, want 90000 once every change was stored", got) + } + if conn.released != 1 { + t.Errorf("released the mailbox %d times, want once before LIST-STATUS", conn.released) + } +} diff --git a/internal/app/worker/wmail/wmail.go b/internal/app/worker/wmail/wmail.go index 0e4355df..4cfba7eb 100644 --- a/internal/app/worker/wmail/wmail.go +++ b/internal/app/worker/wmail/wmail.go @@ -38,7 +38,7 @@ type GraphData struct { } type SmtpImapData struct { - ImapClient *imap.Client + ImapClient ImapConn SmtpClient *smtp.Client Mailboxes []*models.Mailbox mailbox uint32 @@ -76,7 +76,7 @@ type WMail struct { // Sync fair use: the budget engine and the relayed state (governor.go, // sync_state.go). Built in NewWMail from the ADD_EMAIL payload. - gov *governor + gov syncBudget tracker *syncTracker // laneCache remembers deferred messages' lanes across passes; googleTick // and graphTick carry the running pass's stats into provider callbacks. @@ -231,14 +231,15 @@ func NewWMail( mail.SmtpImapData = &SmtpImapData{} if data.ImapSync { - mail.SmtpImapData.ImapClient = &imap.Client{ + conn := &imap.Client{ Email: data.Email, AuthType: models.AuthPlain, Credentials: data.SmtpImap.Credentials.IMAP, } - if err := mail.SmtpImapData.ImapClient.Connect(); err != nil { + if err := conn.Connect(); err != nil { return nil, err } + mail.SmtpImapData.ImapClient = conn // Saved folder cursors: live sync resumes from each folder's stored // HIGHESTMODSEQ instead of re-baselining (and, before this, instead // of re-walking every folder on every worker restart). diff --git a/internal/client/smtpimap/imap/client.go b/internal/client/smtpimap/imap/client.go index 479e6897..e0698055 100644 --- a/internal/client/smtpimap/imap/client.go +++ b/internal/client/smtpimap/imap/client.go @@ -11,6 +11,7 @@ import ( "net/textproto" "strings" "sync" + "sync/atomic" "time" "github.com/emersion/go-imap/v2" @@ -47,6 +48,12 @@ type Client struct { // Guarded by mu. sentMailboxName string + // selected records whether a mailbox is currently SELECTed, so + // ReleaseMailbox does not send UNSELECT in authenticated state, where a + // strict server answers BAD. Atomic because the sync path selects without + // holding mu while warmup actions select under it. + selected atomic.Bool + // BindIP optionally pins outbound TCP to a specific local source address. // When nil, WORKER_BIND_IP is consulted; when still unset, the OS default // route is used. @@ -69,6 +76,7 @@ func (c *Client) Connect() *errx.MailError { } c.client = imapclient.New(conn, nil) + c.selected.Store(false) var xerr *errx.MailError @@ -173,20 +181,30 @@ func (c *Client) Folders() ([]models.Mailbox, *errx.MailError) { } func (c *Client) Mailbox(mailbox string, uidvali, opts *imap.SelectOptions) error { - if _, err := c.client.Select(mailbox, opts).Wait(); err != nil { + if _, err := c.selectMailbox(mailbox, opts); err != nil { return err } return nil } +// selectMailbox is the single SELECT funnel: every path that changes the +// selected mailbox goes through it so ReleaseMailbox knows whether there is +// one to release. A failed SELECT leaves the session with no mailbox +// selected (RFC 3501 6.3.1). +func (c *Client) selectMailbox(mailbox string, opts *imap.SelectOptions) (*imap.SelectData, error) { + data, err := c.client.Select(mailbox, opts).Wait() + c.selected.Store(err == nil) + return data, err +} + // SelectForSync opens a mailbox read-only with CONDSTORE enabled and returns // its message count. FETCH is only valid against a selected mailbox, so the // sync loop must call this before FetchChanges; CONDSTORE on the SELECT is // what arms ChangedSince. The count lets the caller skip the fetch entirely // for an empty mailbox, where a 1:* set is a server error. func (c *Client) SelectForSync(mailbox string) (uint32, *errx.MailError) { - data, err := c.client.Select(mailbox, &imap.SelectOptions{ReadOnly: true, CondStore: true}).Wait() + data, err := c.selectMailbox(mailbox, &imap.SelectOptions{ReadOnly: true, CondStore: true}) if err != nil { return 0, c.handleError(err) } @@ -197,11 +215,19 @@ func (c *Client) SelectForSync(mailbox string) (uint32, *errx.MailError) { // the selected mailbox with the values it held at SELECT, so a loop that keeps // INBOX selected never sees another change land. Servers without UNSELECT keep // the previous behaviour. +// +// It takes mu because it changes selected state, which is exactly what mu +// exists to serialize against an in-flight warmup MOVE/STORE. func (c *Client) ReleaseMailbox() { - if c.client == nil || !c.client.Caps().Has(imap.CapUnselect) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.client == nil || !c.selected.Load() || !c.client.Caps().Has(imap.CapUnselect) { return } - _ = c.client.Unselect().Wait() + if err := c.client.Unselect().Wait(); err == nil { + c.selected.Store(false) + } } // Fetched is one message's envelope as read by FetchEnvelopes, plus what diff --git a/internal/client/smtpimap/imap/warmup_actions.go b/internal/client/smtpimap/imap/warmup_actions.go index 6dc26ff1..b51b4611 100644 --- a/internal/client/smtpimap/imap/warmup_actions.go +++ b/internal/client/smtpimap/imap/warmup_actions.go @@ -15,7 +15,7 @@ func (c *Client) MarkAsRead(ctx context.Context, mailboxName string, uid uint32) c.mu.Lock() defer c.mu.Unlock() - if _, err := c.client.Select(mailboxName, nil).Wait(); err != nil { + if _, err := c.selectMailbox(mailboxName, nil); err != nil { return fmt.Errorf("select %q: %w", mailboxName, err) } @@ -36,7 +36,7 @@ func (c *Client) MarkImportant(ctx context.Context, mailboxName string, uid uint c.mu.Lock() defer c.mu.Unlock() - if _, err := c.client.Select(mailboxName, nil).Wait(); err != nil { + if _, err := c.selectMailbox(mailboxName, nil); err != nil { return fmt.Errorf("select %q: %w", mailboxName, err) } @@ -81,7 +81,7 @@ func (c *Client) moveUID(ctx context.Context, src, dst string, uid uint32) error } func (c *Client) moveUIDLocked(src, dst string, uid uint32) error { - if _, err := c.client.Select(src, nil).Wait(); err != nil { + if _, err := c.selectMailbox(src, nil); err != nil { return fmt.Errorf("select %q: %w", src, err) }