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

@@ -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 {

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) {

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

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