From 13863722418944f54e91023e1b9012ed798aec78 Mon Sep 17 00:00:00 2001 From: Jesse Robbins Date: Wed, 29 Jul 2026 08:44:14 -0400 Subject: [PATCH 1/2] perf: index messages.rfc822_message_id for dedup lookups GetDuplicateGroupMessages ran one unindexed query per RFC822 duplicate group; large archives could exhaust the CLI plan-request timeout before content-hash comparison began. - perf: batch GetDuplicateGroupMessages into chunked queries - perf: switch Engine.Scan to GetDuplicateGroupMessagesBatch - refactor: address final review findings on dedup rfc822 fix - merge current main before review Authored-By: Jesse Robbins (@jesserobbins) Generated with Codex Co-authored-by: Codex --- internal/dedup/dedup.go | 26 +++++--- internal/store/dedup.go | 100 +++++++++++++++++++++++++++-- internal/store/dedup_index_test.go | 46 +++++++++++++ internal/store/dedup_test.go | 54 ++++++++++++++++ internal/store/store.go | 24 +++++++ 5 files changed, 235 insertions(+), 15 deletions(-) create mode 100644 internal/store/dedup_index_test.go diff --git a/internal/dedup/dedup.go b/internal/dedup/dedup.go index 4d72e4286..dc9d018da 100644 --- a/internal/dedup/dedup.go +++ b/internal/dedup/dedup.go @@ -324,6 +324,17 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) { return nil, fmt.Errorf("find duplicates: %w", err) } + rfc822IDs := make([]string, len(storeGroups)) + for i, sg := range storeGroups { + rfc822IDs[i] = sg.RFC822MessageID + } + msgsByGroup, err := e.store.GetDuplicateGroupMessagesBatch( + rfc822IDs, e.config.AccountSourceIDs..., + ) + if err != nil { + return nil, fmt.Errorf("get duplicate group messages: %w", err) + } + report := &Report{ TotalMessages: totalMessages, BySourcePair: make(map[string]int), @@ -334,15 +345,7 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) { if ctx.Err() != nil { return nil, ctx.Err() } - msgs, err := e.store.GetDuplicateGroupMessages( - sg.RFC822MessageID, e.config.AccountSourceIDs..., - ) - if err != nil { - return nil, fmt.Errorf( - "get group messages for %s: %w", - sg.RFC822MessageID, err, - ) - } + msgs := msgsByGroup[sg.RFC822MessageID] if len(msgs) < 2 { continue } @@ -380,6 +383,11 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) { report.Groups = append(report.Groups, group) report.BySourcePair[sourcePairKey(group.Messages)]++ } + // Release the batch result map now that every group has been copied + // into report.Groups: on a large archive it can hold tens of + // thousands of rows, and the content-hash pass below can run long. + //nolint:ineffassign,wastedassign // deliberately drops the reference so the GC can reclaim it before the content-hash pass, not a leftover assignment + msgsByGroup = nil if e.config.ContentHashFallback { // Exclude only losers (messages already selected for pruning) from diff --git a/internal/store/dedup.go b/internal/store/dedup.go index ea96526b5..430971c0d 100644 --- a/internal/store/dedup.go +++ b/internal/store/dedup.go @@ -107,11 +107,14 @@ func (s *Store) FindDuplicatesByRFC822ID(sourceIDs ...int64) ([]DuplicateGroupKe return groups, rows.Err() } -func (s *Store) GetDuplicateGroupMessages( - rfc822ID string, sourceIDs ...int64, -) ([]DuplicateMessageRow, error) { - query := ` - SELECT m.id, m.source_id, s.source_type, s.identifier, +// duplicateGroupMessageColumns is the SELECT column list shared by +// GetDuplicateGroupMessages and GetDuplicateGroupMessagesBatch: the message +// metadata, the two correlated subqueries (label count, from address), and +// the EXISTS clause used to detect the Gmail SENT label. Both methods build +// their SELECT clause from this constant so the two queries can't drift +// apart; GetDuplicateGroupMessagesBatch prepends m.rfc822_message_id (needed +// to key its result map) since it has no per-call rfc822ID to bind. +const duplicateGroupMessageColumns = `m.id, m.source_id, s.source_type, s.identifier, m.source_message_id, COALESCE(m.subject, ''), m.sent_at, m.archived_at, (CASE WHEN mr.message_id IS NOT NULL THEN 1 ELSE 0 END) AS has_raw, @@ -132,7 +135,19 @@ func (s *Store) GetDuplicateGroupMessages( WHERE mr_from.message_id = m.id AND mr_from.recipient_type = 'from' LIMIT 1 - ), '') AS from_email + ), '') AS from_email` + +// GetDuplicateGroupMessages fetches every message row for a single RFC822 +// duplicate group in one query. It is retained as the reference +// implementation that the equivalence test in dedup_test.go checks +// GetDuplicateGroupMessagesBatch against, and it is still exercised by its +// own pre-existing direct tests, even though Engine.Scan now calls +// GetDuplicateGroupMessagesBatch instead. +func (s *Store) GetDuplicateGroupMessages( + rfc822ID string, sourceIDs ...int64, +) ([]DuplicateMessageRow, error) { + query := ` + SELECT ` + duplicateGroupMessageColumns + ` FROM messages m JOIN sources s ON s.id = m.source_id LEFT JOIN message_raw mr ON mr.message_id = m.id @@ -181,6 +196,79 @@ func (s *Store) GetDuplicateGroupMessages( return msgs, rows.Err() } +// GetDuplicateGroupMessagesBatch is the batched form of +// GetDuplicateGroupMessages: it fetches every message row for many RFC822 +// duplicate groups in a handful of chunked queries instead of one query per +// group (see kenn-io/msgvault#510 — 22,025 groups meant 22,025 unindexed +// queries). Returns a map keyed by RFC822 message ID; each value preserves +// the same per-group id-ascending order as GetDuplicateGroupMessages. +func (s *Store) GetDuplicateGroupMessagesBatch( + rfc822IDs []string, sourceIDs ...int64, +) (map[string][]DuplicateMessageRow, error) { + result := make(map[string][]DuplicateMessageRow) + if len(rfc822IDs) == 0 { + return result, nil + } + + const selectCols = "\n\t\tm.rfc822_message_id, " + duplicateGroupMessageColumns + + // queryInChunks binds prefixArgs BEFORE the chunked %s placeholder, so + // the source_id filter (prefixArgs) must appear textually before the + // rfc822_message_id IN (%s) clause below — the reverse of + // GetDuplicateGroupMessages's clause order, which binds rfc822ID first + // via a plain "=" and has no ordering constraint to satisfy. + var prefixArgs []any + sourceFilter := "" + if len(sourceIDs) > 0 { + placeholders := make([]string, len(sourceIDs)) + for i, id := range sourceIDs { + placeholders[i] = "?" + prefixArgs = append(prefixArgs, id) + } + sourceFilter = "m.source_id IN (" + strings.Join(placeholders, ",") + ") AND " + } + + queryTemplate := ` + SELECT` + selectCols + ` + FROM messages m + JOIN sources s ON s.id = m.source_id + LEFT JOIN message_raw mr ON mr.message_id = m.id + WHERE ` + sourceFilter + `m.rfc822_message_id IN (%s) + AND ` + LiveMessagesWhere("m", true) + ` + ORDER BY m.rfc822_message_id, m.id` + + err := queryInChunks(s.db, rfc822IDs, prefixArgs, queryTemplate, + func(rows *loggedRows) error { + var dm DuplicateMessageRow + var rfc822ID string + var sentAt, archivedAt sql.NullTime + var hasRaw, isFromMe, hasSent int + if err := rows.Scan( + &rfc822ID, &dm.ID, &dm.SourceID, &dm.SourceType, &dm.SourceIdentifier, + &dm.SourceMessageID, &dm.Subject, &sentAt, &archivedAt, + &hasRaw, &dm.LabelCount, &isFromMe, &hasSent, + &dm.FromEmail, + ); err != nil { + return err + } + if sentAt.Valid { + dm.SentAt = sentAt.Time + } + if archivedAt.Valid { + dm.ArchivedAt = archivedAt.Time + } + dm.HasRawMIME = hasRaw == 1 + dm.IsFromMe = isFromMe == 1 + dm.HasSentLabel = hasSent == 1 + result[rfc822ID] = append(result[rfc822ID], dm) + return nil + }) + if err != nil { + return nil, fmt.Errorf("get duplicate group messages batch: %w", err) + } + return result, nil +} + func (s *Store) MergeDuplicates( survivorID int64, duplicateIDs []int64, batchID string, ) (*MergeResult, error) { diff --git a/internal/store/dedup_index_test.go b/internal/store/dedup_index_test.go new file mode 100644 index 000000000..ba1a498c9 --- /dev/null +++ b/internal/store/dedup_index_test.go @@ -0,0 +1,46 @@ +package store + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestGetDuplicateGroupMessages_UsesRFC822Index verifies InitSchema creates +// idx_messages_rfc822_message_id and that the dedup per-group lookup query +// plans as an index search, not a full table scan. Locks in the fix for +// kenn-io/msgvault#510: 22,025 unindexed lookups (~190ms each) burned the +// CLI's 30-minute plan-request timeout before content-hash comparison +// started. +func TestGetDuplicateGroupMessages_UsesRFC822Index(t *testing.T) { + require := require.New(t) + assert := assert.New(t) + + dir := t.TempDir() + s, err := OpenForTest(filepath.Join(dir, "dedup_index.db")) + require.NoError(err) + defer func() { _ = s.Close() }() + require.NoError(s.InitSchema()) + + var idxCount int + require.NoError(s.db.QueryRow( + `SELECT COUNT(*) FROM sqlite_master WHERE type='index' AND name='idx_messages_rfc822_message_id'`, + ).Scan(&idxCount)) + assert.Equal(1, idxCount, "idx_messages_rfc822_message_id should be created by InitSchema") + + plan := explainPlan(t, s, + `SELECT id FROM messages WHERE rfc822_message_id = ?`, "some-id") + assert.Contains(plan, "idx_messages_rfc822_message_id", + "lookup by rfc822_message_id should use the index, not a full scan:\n%s", plan) + assert.NotContains(plan, "SCAN messages", + "lookup should not do a full table scan:\n%s", plan) + + batchPlan := explainPlan(t, s, + `SELECT id FROM messages WHERE rfc822_message_id IN (?,?)`, "a", "b") + assert.Contains(batchPlan, "idx_messages_rfc822_message_id", + "batched IN lookup should also use the index:\n%s", batchPlan) + assert.NotContains(batchPlan, "SCAN messages", + "batched lookup should not do a full table scan:\n%s", batchPlan) +} diff --git a/internal/store/dedup_test.go b/internal/store/dedup_test.go index 4febbe91a..331095bdf 100644 --- a/internal/store/dedup_test.go +++ b/internal/store/dedup_test.go @@ -332,3 +332,57 @@ func TestStore_GetAllRawMIMECandidates_PreservesFromCase(t *testing.T) { require.NotNil(got, "test message %d not in candidates: %+v", id, cands) assert.Equal(t, mxid, got.FromEmail, "FromEmail (case must be preserved)") } + +func TestStore_GetDuplicateGroupMessagesBatch_MatchesPerGroupQuery(t *testing.T) { + require := require.New(t) + assert := assert.New(t) + f := storetest.New(t) + + // 600 groups exceeds queryInChunks's chunk size of 500, so this + // exercises at least two chunk rounds and would catch a bug where the + // result map is reset per chunk instead of accumulated across chunks. + const numGroups = 600 + rfc822IDs := make([]string, numGroups) + for i := 0; i < numGroups; i++ { + rfc822ID := fmt.Sprintf("rfc822-batch-%d", i) + rfc822IDs[i] = rfc822ID + newRFC822Message(t, f, fmt.Sprintf("src-%d-a", i), rfc822ID) + newRFC822Message(t, f, fmt.Sprintf("src-%d-b", i), rfc822ID) + } + + batched, err := f.Store.GetDuplicateGroupMessagesBatch(rfc822IDs, f.Source.ID) + require.NoError(err, "GetDuplicateGroupMessagesBatch") + require.Len(batched, numGroups, "batched group count") + + for _, rfc822ID := range rfc822IDs { + want, err := f.Store.GetDuplicateGroupMessages(rfc822ID, f.Source.ID) + require.NoError(err, "GetDuplicateGroupMessages reference for %s", rfc822ID) + assert.Equal(want, batched[rfc822ID], "mismatch for group %s", rfc822ID) + } +} + +func TestStore_GetDuplicateGroupMessagesBatch_EmptyInput(t *testing.T) { + require := require.New(t) + assert := assert.New(t) + f := storetest.New(t) + + batched, err := f.Store.GetDuplicateGroupMessagesBatch(nil) + require.NoError(err, "GetDuplicateGroupMessagesBatch with nil input") + assert.Empty(batched, "no groups requested, no groups returned") +} + +func TestStore_GetDuplicateGroupMessagesBatch_FiltersBySourceID(t *testing.T) { + require := require.New(t) + assert := assert.New(t) + f := storetest.New(t) + idA := newRFC822Message(t, f, "src-a", "rfc822-scoped") + idB := newRFC822Message(t, f, "src-b", "rfc822-scoped") + + scoped, err := f.Store.GetDuplicateGroupMessagesBatch( + []string{"rfc822-scoped"}, f.Source.ID, + ) + require.NoError(err, "GetDuplicateGroupMessagesBatch scoped to source") + require.Len(scoped["rfc822-scoped"], 2) + assert.Equal(idA, scoped["rfc822-scoped"][0].ID) + assert.Equal(idB, scoped["rfc822-scoped"][1].ID) +} diff --git a/internal/store/store.go b/internal/store/store.go index 3320de646..fd4cca8d9 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -940,6 +940,30 @@ func (s *Store) InitSchema() error { } } + // Index over rfc822_message_id serves dedup's per-group message lookup + // (GetDuplicateGroupMessages / GetDuplicateGroupMessagesBatch). Without + // it, each lookup was a full scan of the messages table — measured at + // ~190ms/lookup, with one lookup per duplicate group, so a scan with + // 22k groups burned the entire 30-minute CLI plan-request timeout + // before content-hash comparison even started (kenn-io/msgvault#510). + // Plain (non-partial) index: a partial WHERE rfc822_message_id IS NOT + // NULL AND != '' form is not usable by the planner for this table's + // bound `= ?` / `IN (...)` lookups — SQLite can't prove col = ? implies + // col != '' since ? could bind to '' — so it would silently fall back + // to SCAN (verified via EXPLAIN QUERY PLAN before writing this). + // Identical DDL on both backends; runMaintenance already handles the + // PostgreSQL statement_timeout exemption internally (finding S1). IF + // NOT EXISTS is idempotent per start. + if err := s.runMaintenance(context.Background(), func(ctx context.Context, tx *loggedTx) error { + _, err := tx.ExecContext(ctx, ` + CREATE INDEX IF NOT EXISTS idx_messages_rfc822_message_id + ON messages(rfc822_message_id) + `) + return err + }); err != nil { + return fmt.Errorf("create rfc822 message id index: %w", err) + } + // Backfill last_modified for rows that predate the column. SQLite cannot // ADD COLUMN with a non-constant default, so the legacy ADD COLUMN above // leaves existing rows NULL; this one-shot UPDATE sets them to From 2c25920e010c7dc0c187f4926b46db11ce1f1f18 Mon Sep 17 00:00:00 2001 From: Wes McKinney Date: Wed, 29 Jul 2026 08:51:30 -0400 Subject: [PATCH 2/2] fix: honor cancellation during dedup batch lookup A timed-out dedup scan could continue issuing every RFC822 lookup chunk because the batching helper replaced the request context with a background context. Large archives therefore kept consuming database work after the caller disconnected. Thread the scan context through the batch API and chunk queries so in-flight database work is interruptible and no later chunk starts after cancellation. The background-context wrapper remains for callers that do not yet carry request scope. Generated with Codex Co-authored-by: Codex --- internal/dedup/dedup.go | 4 ++-- internal/store/dedup.go | 12 +++++++++++- internal/store/dedup_test.go | 18 +++++++++++++++++- internal/store/store.go | 11 ++++++++++- 4 files changed, 40 insertions(+), 5 deletions(-) diff --git a/internal/dedup/dedup.go b/internal/dedup/dedup.go index dc9d018da..ca485d9bb 100644 --- a/internal/dedup/dedup.go +++ b/internal/dedup/dedup.go @@ -328,8 +328,8 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) { for i, sg := range storeGroups { rfc822IDs[i] = sg.RFC822MessageID } - msgsByGroup, err := e.store.GetDuplicateGroupMessagesBatch( - rfc822IDs, e.config.AccountSourceIDs..., + msgsByGroup, err := e.store.GetDuplicateGroupMessagesBatchContext( + ctx, rfc822IDs, e.config.AccountSourceIDs..., ) if err != nil { return nil, fmt.Errorf("get duplicate group messages: %w", err) diff --git a/internal/store/dedup.go b/internal/store/dedup.go index 430971c0d..3c0769a72 100644 --- a/internal/store/dedup.go +++ b/internal/store/dedup.go @@ -204,6 +204,16 @@ func (s *Store) GetDuplicateGroupMessages( // the same per-group id-ascending order as GetDuplicateGroupMessages. func (s *Store) GetDuplicateGroupMessagesBatch( rfc822IDs []string, sourceIDs ...int64, +) (map[string][]DuplicateMessageRow, error) { + return s.GetDuplicateGroupMessagesBatchContext( + context.Background(), rfc822IDs, sourceIDs..., + ) +} + +// GetDuplicateGroupMessagesBatchContext is the request-aware form of +// GetDuplicateGroupMessagesBatch. +func (s *Store) GetDuplicateGroupMessagesBatchContext( + ctx context.Context, rfc822IDs []string, sourceIDs ...int64, ) (map[string][]DuplicateMessageRow, error) { result := make(map[string][]DuplicateMessageRow) if len(rfc822IDs) == 0 { @@ -237,7 +247,7 @@ func (s *Store) GetDuplicateGroupMessagesBatch( AND ` + LiveMessagesWhere("m", true) + ` ORDER BY m.rfc822_message_id, m.id` - err := queryInChunks(s.db, rfc822IDs, prefixArgs, queryTemplate, + err := queryInChunksContext(ctx, s.db, rfc822IDs, prefixArgs, queryTemplate, func(rows *loggedRows) error { var dm DuplicateMessageRow var rfc822ID string diff --git a/internal/store/dedup_test.go b/internal/store/dedup_test.go index 331095bdf..73ff97653 100644 --- a/internal/store/dedup_test.go +++ b/internal/store/dedup_test.go @@ -1,6 +1,7 @@ package store_test import ( + "context" "database/sql" "fmt" "testing" @@ -343,7 +344,7 @@ func TestStore_GetDuplicateGroupMessagesBatch_MatchesPerGroupQuery(t *testing.T) // result map is reset per chunk instead of accumulated across chunks. const numGroups = 600 rfc822IDs := make([]string, numGroups) - for i := 0; i < numGroups; i++ { + for i := range numGroups { rfc822ID := fmt.Sprintf("rfc822-batch-%d", i) rfc822IDs[i] = rfc822ID newRFC822Message(t, f, fmt.Sprintf("src-%d-a", i), rfc822ID) @@ -371,6 +372,21 @@ func TestStore_GetDuplicateGroupMessagesBatch_EmptyInput(t *testing.T) { assert.Empty(batched, "no groups requested, no groups returned") } +func TestStore_GetDuplicateGroupMessagesBatchContext_Canceled(t *testing.T) { + require := require.New(t) + f := storetest.New(t) + newRFC822Message(t, f, "src-a", "rfc822-canceled") + newRFC822Message(t, f, "src-b", "rfc822-canceled") + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := f.Store.GetDuplicateGroupMessagesBatchContext( + ctx, []string{"rfc822-canceled"}, f.Source.ID, + ) + require.ErrorIs(err, context.Canceled) +} + func TestStore_GetDuplicateGroupMessagesBatch_FiltersBySourceID(t *testing.T) { require := require.New(t) assert := assert.New(t) diff --git a/internal/store/store.go b/internal/store/store.go index fd4cca8d9..67e3f28ca 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -603,12 +603,21 @@ func (s *Store) runMaintenance(ctx context.Context, fn func(ctx context.Context, // so streaming-query timing reflects scan-close, not just prepare. type chunkQuerier interface { Query(query string, args ...any) (*loggedRows, error) + QueryContext(ctx context.Context, query string, args ...any) (*loggedRows, error) Exec(query string, args ...any) (sql.Result, error) } func queryInChunks[T any](db chunkQuerier, ids []T, prefixArgs []any, queryTemplate string, fn func(*loggedRows) error) error { + return queryInChunksContext(context.Background(), db, ids, prefixArgs, queryTemplate, fn) +} + +func queryInChunksContext[T any](ctx context.Context, db chunkQuerier, ids []T, prefixArgs []any, queryTemplate string, fn func(*loggedRows) error) error { const chunkSize = 500 for i := 0; i < len(ids); i += chunkSize { + if err := ctx.Err(); err != nil { + return err + } + end := min(i+chunkSize, len(ids)) chunk := ids[i:end] @@ -620,7 +629,7 @@ func queryInChunks[T any](db chunkQuerier, ids []T, prefixArgs []any, queryTempl } query := fmt.Sprintf(queryTemplate, strings.Join(placeholders, ",")) - rows, err := db.Query(query, args...) + rows, err := db.QueryContext(ctx, query, args...) if err != nil { return err }