fix: count only the blocks whose purge succeeded

Author: mennatnaga

pkg/pruner/pruner.go

@@ -259,6 +259,7 @@ func (p *Pruner) runEventQueuePrune(ctx context.Context, q *EventQueue) error {
 func (p *Pruner) purgeFromDrainResult(ctx context.Context, q *EventQueue, result *DrainResult) error {
 	startTime := time.Now()
 	totalSubmitted := int64(0)
+	blocksPruned := int64(0)
 
 	// Dependents before blocks, so a block is never removed ahead of the documents that
 	// reference it.
@@ -277,6 +278,9 @@ func (p *Pruner) purgeFromDrainResult(ctx context.Context, q *EventQueue, result
 			// A failed collection is re-queued and purged again later, so counting its
 			// partial progress here would count those documents twice.
 			totalSubmitted += submitted
+			if colName == p.collections.BlockCollection {
+				blocksPruned = int64(result.BlockCount)
+			}
 			continue
 		}
 
@@ -299,7 +303,9 @@ func (p *Pruner) purgeFromDrainResult(ctx context.Context, q *EventQueue, result
 		totalSubmitted, result.BlockCount, time.Since(startTime))
 
 	p.mu.Lock()
-	p.totalBlocksPruned += int64(result.BlockCount)
+	// Only blocks whose own purge succeeded. A re-queued block collection is drained and
+	// counted again on a later cycle, so counting it here counts those blocks twice.
+	p.totalBlocksPruned += blocksPruned
 	p.totalDocsSubmitted += totalSubmitted
 	p.lastPruneTime = time.Now()
 	p.mu.Unlock()

pkg/pruner/pruner_test.go

@@ -2,6 +2,7 @@ package pruner
 
 import (
 	"context"
+	"errors"
 	"os"
 	"path/filepath"
 	"testing"
@@ -160,6 +161,45 @@ func TestPurgeStopsOnCancelledContext(t *testing.T) {
 	require.Zero(t, submitted)
 }
 
+// A block collection whose purge fails is re-queued, so the same blocks are drained and
+// counted again later. Counting them on the failed cycle too makes the total exceed the
+// blocks that were ever pruned.
+func TestBlockCounterIgnoresAFailedBlockPurge(t *testing.T) {
+	cols := DefaultCollectionConfig()
+
+	q := NewEventQueue(cols)
+	q.Push(cols.BlockCollection, testDocID(1))
+	result := q.DrainDocs(1)
+	require.NotNil(t, result)
+	require.Equal(t, 1, result.BlockCount)
+
+	p := &Pruner{cfg: &Config{Enabled: true}, collections: cols, stopChan: make(chan struct{})}
+	p.purgeDocs = func(context.Context, []client.DocID) error {
+		return errors.New("purge failed")
+	}
+
+	require.NoError(t, p.purgeFromDrainResult(context.Background(), q, result))
+
+	require.Zero(t, p.totalBlocksPruned, "blocks that were re-queued must not count as pruned")
+	require.Equal(t, 1, q.BlockCount(), "guard: the blocks are back on the queue")
+}
+
+// The counter still moves on the path that did purge.
+func TestBlockCounterCountsASuccessfulBlockPurge(t *testing.T) {
+	cols := DefaultCollectionConfig()
+
+	q := NewEventQueue(cols)
+	q.Push(cols.BlockCollection, testDocID(1))
+	result := q.DrainDocs(1)
+	require.NotNil(t, result)
+
+	p := &Pruner{cfg: &Config{Enabled: true}, collections: cols, stopChan: make(chan struct{})}
+	p.purgeDocs = func(context.Context, []client.DocID) error { return nil }
+
+	require.NoError(t, p.purgeFromDrainResult(context.Background(), q, result))
+	require.Equal(t, int64(1), p.totalBlocksPruned)
+}
+
 // DrainDocs empties every collection up front, so a cycle that stops part-way has to re-queue the
 // collections it never reached as well as the one it stopped on. Nothing else re-adds a document
 // once it has replicated, so anything left behind is never pruned.