From c8a4fd5ba36f2f911c2a7209a777b44e5ca9908f Mon Sep 17 00:00:00 2001 From: wizardchen Date: Thu, 9 Jul 2026 15:25:05 +0800 Subject: [PATCH] fix(wiki): prevent silent work loss and cross-batch doc races Two correctness fixes on the Phase 3 concurrent claim-based ingest path: - Slug-lock timeout / reduce failure no longer silently drops work. When a per-slug lock can't be acquired within wikiSlugLockWait, or reduce returns an error, the page's read-modify-write never lands. Previously the owning document was still counted as successful and its pending row was deleted at trim time, permanently losing the contribution (finalize only rebuilds the index / cross-links, it does not re-run reduce). Such documents are now folded into failedOps: their rows stay queued, their finalizing slot is held, and requeueFailedOps' fail_count budget retries them (dead-lettering a permanently hot/broken slug). - ClaimBatch now claims by dedup_key affinity so a document is never split across concurrent batches. limit counts distinct dedup_keys (documents), and ALL eligible rows sharing a key are claimed together. On Postgres the per-key anchor row is locked with FOR UPDATE SKIP LOCKED, handing concurrent claimers disjoint key sets. This restores in-document ordering (e.g. ingest-then-retract collapses correctly) that the per-batch dedup could not guarantee once the exclusive per-KB lock was removed, preventing a stale ingest from resurrecting a retracted document. Adds repository tests for same-key affinity and in-flight sibling rows. --- internal/application/repository/task_queue.go | 87 +++++++++++++++---- .../application/repository/task_queue_test.go | 66 ++++++++++++++ .../application/service/wiki_ingest_batch.go | 67 +++++++++++++- internal/types/interfaces/task_queue.go | 17 ++-- 4 files changed, 208 insertions(+), 29 deletions(-) diff --git a/internal/application/repository/task_queue.go b/internal/application/repository/task_queue.go index c9b63c8d..1d8848ff 100644 --- a/internal/application/repository/task_queue.go +++ b/internal/application/repository/task_queue.go @@ -10,7 +10,6 @@ import ( "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "gorm.io/gorm" - "gorm.io/gorm/clause" ) // taskPendingOpsRepository implements interfaces.TaskPendingOpsRepository. @@ -73,15 +72,31 @@ func (r *taskPendingOpsRepository) PeekBatch( return ops, nil } -// ClaimBatch atomically claims up to `limit` eligible rows for the tuple. +// ClaimBatch atomically claims eligible rows for the tuple, grouped by +// dedup_key. `limit` counts DISTINCT dedup_keys (i.e. documents), NOT rows: +// ALL eligible rows sharing a chosen dedup_key are claimed together and +// returned in the same batch. This is the invariant the concurrent wiki +// consumers rely on — a document with multiple queued ops (e.g. an ingest +// followed by a retract) must never be split across two concurrent batches, +// otherwise each batch's per-batch last-write-wins dedup can't collapse the +// pair and the two ops race (a stale ingest could resurrect a retracted doc). +// // Eligibility = unclaimed (claimed_at IS NULL) OR stale claim // (claimed_at < staleBefore). The whole thing runs in one transaction: -// we first lock the candidate ids (FOR UPDATE SKIP LOCKED on Postgres so -// concurrent claimers slide past each other's rows), stamp claimed_at, -// then read the claimed rows back. SKIP LOCKED is Postgres-only; on other -// dialects (SQLite, used by unit tests / Lite) writes are serialized by -// the engine so the plain SELECT+UPDATE is already race-free within the -// single-writer model. +// +// - Postgres: we lock the ANCHOR row (earliest eligible id) of each +// candidate dedup_key with FOR UPDATE SKIP LOCKED. Because the anchor +// uniquely represents its key, SKIP LOCKED hands concurrent claimers +// DISJOINT key sets — a key whose anchor is already locked by another +// in-flight claim is skipped entirely rather than half-claimed. We then +// stamp every eligible row of the chosen keys and read them back. +// - Other dialects (SQLite, used by unit tests / Lite mode): writes are +// serialized by the single-writer engine, so a plain grouped SELECT + +// UPDATE is already race-free. +// +// Rows are claimed by explicit id (only the eligible ones), so a freshly +// enqueued or still-in-flight sibling row of a chosen key is never handed +// out twice. func (r *taskPendingOpsRepository) ClaimBatch( ctx context.Context, taskType, scope, scopeID string, @@ -97,19 +112,55 @@ func (r *taskPendingOpsRepository) ClaimBatch( now := time.Now() var claimed []*types.TaskPendingOp err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - sel := tx.Model(&types.TaskPendingOp{}). + // 1. Pick up to `limit` distinct dedup_keys to claim, oldest first. + var keys []string + if tx.Dialector.Name() == "postgres" { + // Lock the anchor (earliest eligible) row of each key with SKIP + // LOCKED so concurrent claimers get disjoint KEY sets, then map + // the locked anchors back to their dedup_keys. + const anchorSQL = ` +SELECT dedup_key FROM task_pending_ops +WHERE id IN ( + SELECT id FROM ( + SELECT id, ROW_NUMBER() OVER (PARTITION BY dedup_key ORDER BY id) AS rn + FROM task_pending_ops + WHERE task_type = ? AND scope = ? AND scope_id = ? + AND (claimed_at IS NULL OR claimed_at < ?) + ) anchors WHERE anchors.rn = 1 +) +ORDER BY id ASC +LIMIT ? +FOR UPDATE SKIP LOCKED` + if err := tx.Raw(anchorSQL, taskType, scope, scopeID, staleBefore, limit). + Scan(&keys).Error; err != nil { + return err + } + } else { + if err := tx.Model(&types.TaskPendingOp{}). + Where("task_type = ? AND scope = ? AND scope_id = ?", taskType, scope, scopeID). + Where("(claimed_at IS NULL OR claimed_at < ?)", staleBefore). + Group("dedup_key"). + Order("MIN(id) ASC"). + Limit(limit). + Pluck("dedup_key", &keys).Error; err != nil { + return err + } + } + if len(keys) == 0 { + return nil + } + + // 2. Resolve the exact eligible rows for the chosen keys and claim + // them by id. Claiming by id (not by "dedup_key IN keys") means a + // sibling row that is still in flight elsewhere (claimed & fresh) + // is left untouched and never returned to this batch. + var ids []int64 + if err := tx.Model(&types.TaskPendingOp{}). Where("task_type = ? AND scope = ? AND scope_id = ?", taskType, scope, scopeID). - // Parentheses are explicit so the OR binds tightly and doesn't - // leak past the tuple filter (GORM ANDs Where clauses but does - // not parenthesize raw fragments). + Where("dedup_key IN ?", keys). Where("(claimed_at IS NULL OR claimed_at < ?)", staleBefore). Order("id ASC"). - Limit(limit) - if tx.Dialector.Name() == "postgres" { - sel = sel.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}) - } - var ids []int64 - if err := sel.Pluck("id", &ids).Error; err != nil { + Pluck("id", &ids).Error; err != nil { return err } if len(ids) == 0 { diff --git a/internal/application/repository/task_queue_test.go b/internal/application/repository/task_queue_test.go index 3a45f380..9e633ad2 100644 --- a/internal/application/repository/task_queue_test.go +++ b/internal/application/repository/task_queue_test.go @@ -307,6 +307,72 @@ func TestTaskPendingOps_ClaimBatch_MarksAndReturnsDisjoint(t *testing.T) { assert.Len(t, third, 0) } +// TestTaskPendingOps_ClaimBatch_KeepsSameKeyTogether verifies the +// dedup_key affinity invariant: all rows sharing a knowledge_id are claimed +// in the SAME batch (never split), and `limit` counts distinct keys, not +// rows. This is what stops a concurrent batch from processing one op of a +// document while another batch processes a second op of the same document. +func TestTaskPendingOps_ClaimBatch_KeepsSameKeyTogether(t *testing.T) { + db := setupTaskQueueTestDB(t) + repo := NewTaskPendingOpsRepository(db) + ctx := context.Background() + + // Document k1 has TWO queued ops (ingest then retract); k2 and k3 have + // one each. Enqueue order interleaves them so a naive row-ordered claim + // would split k1 across batches. + require.NoError(t, repo.Enqueue(ctx, makePendingOp("wiki:ingest", "knowledge_base", "kb", "ingest", "k1", nil))) + require.NoError(t, repo.Enqueue(ctx, makePendingOp("wiki:ingest", "knowledge_base", "kb", "ingest", "k2", nil))) + require.NoError(t, repo.Enqueue(ctx, makePendingOp("wiki:ingest", "knowledge_base", "kb", "retract", "k1", nil))) + require.NoError(t, repo.Enqueue(ctx, makePendingOp("wiki:ingest", "knowledge_base", "kb", "ingest", "k3", nil))) + + stale := time.Now().Add(-time.Hour) + + // limit=2 keys → k1 (both rows) + k2. k1 must NOT be split. + first, err := repo.ClaimBatch(ctx, "wiki:ingest", "knowledge_base", "kb", 2, stale) + require.NoError(t, err) + require.Len(t, first, 3, "k1's two rows + k2's one row") + byKey := map[string]int{} + for _, r := range first { + byKey[r.DedupKey]++ + } + assert.Equal(t, 2, byKey["k1"], "both k1 ops claimed together") + assert.Equal(t, 1, byKey["k2"]) + assert.Zero(t, byKey["k3"], "k3 belongs to the next batch (limit was 2 keys)") + + // A concurrent-style second claim gets the remaining key only — it can + // never see k1's rows again (disjoint). + second, err := repo.ClaimBatch(ctx, "wiki:ingest", "knowledge_base", "kb", 10, stale) + require.NoError(t, err) + require.Len(t, second, 1) + assert.Equal(t, "k3", second[0].DedupKey) +} + +// TestTaskPendingOps_ClaimBatch_NewSiblingRowNotDoubleClaimed verifies a +// row enqueued for an already-claimed (in-flight) document is left for a +// later batch rather than being handed out alongside the fresh claim. +func TestTaskPendingOps_ClaimBatch_NewSiblingRowNotDoubleClaimed(t *testing.T) { + db := setupTaskQueueTestDB(t) + repo := NewTaskPendingOpsRepository(db) + ctx := context.Background() + + require.NoError(t, repo.Enqueue(ctx, makePendingOp("wiki:ingest", "knowledge_base", "kb", "ingest", "k1", nil))) + stale := time.Now().Add(-time.Hour) + + first, err := repo.ClaimBatch(ctx, "wiki:ingest", "knowledge_base", "kb", 5, stale) + require.NoError(t, err) + require.Len(t, first, 1) + + // A retract for the same, still-in-flight document arrives. + require.NoError(t, repo.Enqueue(ctx, makePendingOp("wiki:ingest", "knowledge_base", "kb", "retract", "k1", nil))) + + // The next claim gets ONLY the new (unclaimed) retract row — the + // still-claimed ingest row is not re-handed-out. + second, err := repo.ClaimBatch(ctx, "wiki:ingest", "knowledge_base", "kb", 5, stale) + require.NoError(t, err) + require.Len(t, second, 1) + assert.Equal(t, "retract", second[0].Op) +} + // TestTaskPendingOps_ClaimBatch_ReclaimsStale verifies a claim older than // staleBefore is re-claimable (crash recovery), while a fresh claim is not. func TestTaskPendingOps_ClaimBatch_ReclaimsStale(t *testing.T) { diff --git a/internal/application/service/wiki_ingest_batch.go b/internal/application/service/wiki_ingest_batch.go index b9f0edd6..349d5b96 100644 --- a/internal/application/service/wiki_ingest_batch.go +++ b/internal/application/service/wiki_ingest_batch.go @@ -512,6 +512,27 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task // slugs out of the wiki log feed so users don't see clickable entries // pointing at missing pages. failedAdditionSlugs := make(map[string]struct{}) + // unappliedSlugKIDs collects the knowledge_ids that contributed to a + // slug whose update never landed — either because we could NOT acquire + // the per-slug lock within wikiSlugLockWait, or because reduce returned + // an error. In both cases the page keeps its prior content, so the + // owning document(s) must be re-queued rather than trimmed — otherwise + // the row is deleted and the contribution is silently lost forever + // (finalize only rebuilds the index / cross-links, it does not re-run + // reduce). requeueFailedOps' fail_count budget bounds the retries and + // dead-letters a document whose slug fails/stays hot permanently. + unappliedSlugKIDs := make(map[string]struct{}) + // collectUnapplied records every knowledge_id backing a slug we failed + // to apply. Caller holds no lock; it takes reduceMu itself. + collectUnapplied := func(updates []SlugUpdate) { + reduceMu.Lock() + for _, u := range updates { + if u.KnowledgeID != "" { + unappliedSlugKIDs[u.KnowledgeID] = struct{}{} + } + } + reduceMu.Unlock() + } // Build the kid → wikiSpan lookup before kicking off reduce. Each // per-slug reduce attaches a postprocess.wiki.page[slug] subspan @@ -546,14 +567,22 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task return nil } if !acquired { - // Contended slug we couldn't get in time. Best-effort skip, - // matching reduce-failure semantics: the page keeps its prior - // content and the finalize/dead-link passes reconcile later. + // Contended slug we couldn't get in time. The page keeps its + // prior content, so the documents that fed this slug are NOT + // done: record their knowledge_ids so the trim phase re-queues + // them (via the failed-op retry budget) for a later, + // hopefully-uncontended batch instead of deleting their rows. logger.Warnf(reduceCtx, "wiki ingest: slug %s busy > %s, deferring update", slug, wikiSlugLockWait) + collectUnapplied(updates) return nil } if reduceErr != nil { + // The page's read-modify-write failed, so this slug's update + // never landed. Re-queue the contributing documents (same + // fail_count budget as a map failure) so a later batch retries + // instead of silently dropping the row at trim time. logger.Warnf(reduceCtx, "wiki ingest: reduce failed for slug %s: %v", slug, reduceErr) + collectUnapplied(updates) if isLikelyRateLimitError(reduceErr) { reduceMu.Lock() rateLimited = true @@ -712,7 +741,14 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task // the WikiSpan nil-check below so a doc that had no attempt to // attach a span to still drains its counter slot. The matching +1 // is seeded by KnowledgePostProcess.SetFinalizing. - s.finalizeWikiSubtask(ctx, r.KnowledgeID) + // + // EXCEPT docs with an unapplied slug (contended lock or reduce + // error): those are re-queued below, so keep their finalizing slot + // held — the retry (or the dead-letter drain in requeueFailedOps) + // releases it once the op reaches a real terminal state. + if _, unapplied := unappliedSlugKIDs[r.KnowledgeID]; !unapplied { + s.finalizeWikiSubtask(ctx, r.KnowledgeID) + } if r.WikiSpan == nil { continue } @@ -748,6 +784,29 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task // mapOneDocument (the failedOps path returns before reaching // docResults). Nothing extra to do here for them. + // Fold documents with an unapplied slug (contended lock or reduce error) + // into failedOps so they are neither trimmed nor promoted to completed: + // requeueFailedOps then runs them through the same fail_count budget + // (retry now, dead-letter once the slug stays permanently hot/broken) as + // a map-phase failure. A doc already counted as a map failure is skipped + // to avoid a double fail_count bump. + if len(unappliedSlugKIDs) > 0 { + failedKIDs := make(map[string]struct{}, len(failedOps)) + for _, op := range failedOps { + failedKIDs[op.KnowledgeID] = struct{}{} + } + for _, op := range pendingOps { + if _, unapplied := unappliedSlugKIDs[op.KnowledgeID]; !unapplied { + continue + } + if _, already := failedKIDs[op.KnowledgeID]; already { + continue + } + failedOps = append(failedOps, op) + failedKIDs[op.KnowledgeID] = struct{}{} + } + } + // Build the trim set: rows that should be removed from // task_pending_ops. We start from the full peekedIDs (every row we // pulled, even ones de-duplicated by knowledge_id) and subtract diff --git a/internal/types/interfaces/task_queue.go b/internal/types/interfaces/task_queue.go index 1018c95e..22b56074 100644 --- a/internal/types/interfaces/task_queue.go +++ b/internal/types/interfaces/task_queue.go @@ -35,13 +35,16 @@ type TaskPendingOpsRepository interface { // IncrFailCount and leave them for the next pass). PeekBatch(ctx context.Context, taskType, scope, scopeID string, limit int) ([]*types.TaskPendingOp, error) - // ClaimBatch atomically claims up to `limit` rows for the tuple, - // ordered by id ASC (FIFO), and returns them with claimed_at set to - // now. A row is eligible when it is unclaimed (claimed_at IS NULL) or - // its claim is stale (claimed_at < staleBefore) — the latter recovers - // rows abandoned by a crashed worker. On Postgres the selection uses - // FOR UPDATE SKIP LOCKED so concurrent claimers never block on or - // double-claim each other's rows. + // ClaimBatch atomically claims eligible rows for the tuple, grouped by + // dedup_key, and returns them with claimed_at set to now. `limit` + // counts DISTINCT dedup_keys (documents), NOT rows: ALL eligible rows + // sharing a chosen dedup_key are claimed together so a document with + // several queued ops is never split across two concurrent batches. + // A row is eligible when it is unclaimed (claimed_at IS NULL) or its + // claim is stale (claimed_at < staleBefore) — the latter recovers rows + // abandoned by a crashed worker. On Postgres the per-key anchor row is + // locked with FOR UPDATE SKIP LOCKED so concurrent claimers take + // disjoint key sets without blocking or double-claiming. // // Claimed rows are NOT removed: the consumer must DeleteByIDs on // success, or ReleaseByIDs to hand a still-retryable row back to the