perf: purge in chunks of 8, not 100

Author: mennatnaga

internal/db/collection_purge.go

@@ -27,8 +27,9 @@ import (
 
 // purgeChunkSize is the number of documents purged per transaction when no caller
 // transaction is supplied. It keeps each commit well under the store's per-transaction
-// size limit.
-const purgeChunkSize = 100
+// size limit. Badger re-sorts the transaction's pending writes on every iterator open,
+// and each document opens several, so cost grows faster than chunk size.
+const purgeChunkSize = 8
 
 // PurgeByDocIDs permanently removes all state for the given documents from this node:
 // datastore values, headstore entries, and, when pruneHistory is true, every blockstore

internal/db/collection_purge_bench_test.go

@@ -0,0 +1,136 @@
+// Copyright 2026 Democratized Data Foundation
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package db
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+
+	badgerds "github.com/dgraph-io/badger/v4"
+	"github.com/stretchr/testify/require"
+
+	"github.com/sourcenetwork/corekv/badger"
+
+	"github.com/sourcenetwork/defradb/client"
+	acpDB "github.com/sourcenetwork/defradb/internal/db/acp"
+)
+
+// newOnDiskDB opens a DB on a badger directory with stock options. It returns a close func
+// rather than registering cleanup so a caller can release it per iteration. An in-memory
+// store has no value log and different write amplification, so it cannot stand in for a
+// deployed node when the cost of a write path is what is being measured.
+func newOnDiskDB(b *testing.B, ctx context.Context) (*DB, func()) {
+	b.Helper()
+
+	dir, err := os.MkdirTemp("", "purgebench")
+	require.NoError(b, err)
+
+	rootstore, err := badger.NewDatastore(dir, badgerds.DefaultOptions(dir))
+	require.NoError(b, err)
+
+	adminInfo, err := acpDB.NewNACInfo(ctx, "", false)
+	require.NoError(b, err)
+
+	db, err := newDB(ctx, rootstore, adminInfo)
+	require.NoError(b, err)
+
+	return db, func() {
+		db.Close()
+		_ = os.RemoveAll(dir)
+	}
+}
+
+// setupIndexedCollection builds a collection with three secondary indexes, one unique,
+// matching the shape of a heavily indexed production collection. The index count is what
+// decides how many writes each purged document adds to its transaction, and that is the
+// term the pending-write sort is quadratic in.
+func setupIndexedCollection(b *testing.B, ctx context.Context, db *DB) client.Collection {
+	b.Helper()
+
+	_, err := db.AddCollection(ctx, `type Record {
+		hash: String
+		blockNumber: Int
+		groupID: String
+		payload: String
+	}`)
+	require.NoError(b, err)
+
+	col, err := db.GetCollectionByName(ctx, "Record")
+	require.NoError(b, err)
+
+	for _, req := range []client.NewIndexRequest{
+		{Fields: []client.IndexedFieldDescription{{Name: "hash"}}, Unique: true},
+		{Fields: []client.IndexedFieldDescription{{Name: "blockNumber"}}},
+		{Fields: []client.IndexedFieldDescription{{Name: "groupID"}}},
+	} {
+		_, err := col.NewIndex(ctx, req)
+		require.NoError(b, err)
+	}
+
+	return col
+}
+
+// addRecords writes n documents and returns their IDs. The payload field is unindexed and
+// exists only to give each document enough size to reach the value log.
+func addRecords(b *testing.B, ctx context.Context, col client.Collection, n int) []client.DocID {
+	b.Helper()
+
+	docIDs := make([]client.DocID, 0, n)
+	for i := range n {
+		doc, err := client.NewDocFromJSON(ctx, fmt.Appendf(nil,
+			`{"hash":"0x%064x","blockNumber":%d,"groupID":"g-%d","payload":%q}`,
+			i, i/200, i/200, fmt.Sprintf("%0512d", i)), col.Version())
+		require.NoError(b, err)
+		require.NoError(b, col.AddDocument(ctx, doc))
+		docIDs = append(docIDs, doc.ID())
+	}
+
+	return docIDs
+}
+
+// BenchmarkPurgeByDocIDsChunkSize measures how long it takes to purge a fixed set of
+// documents as the number of them sharing a transaction changes. purgeChunkSize is not a
+// free choice: the per-chunk cost grows with the square of the chunk, so total purge time
+// is close to linear in it, and this reports that curve.
+//
+// pruneHistory is on because that is the deployed setting and it adds the per-document DAG
+// walk. Documents are written locally, so their DAGs are one commit deep; a document built
+// up over many merges walks further and costs more than this measures.
+func BenchmarkPurgeByDocIDsChunkSize(b *testing.B) {
+	const docs = 2000
+
+	for _, chunkSize := range []int{8, 25, 50, 100, 200} {
+		b.Run(fmt.Sprintf("chunk=%d", chunkSize), func(b *testing.B) {
+			ctx := context.Background()
+
+			for b.Loop() {
+				b.StopTimer()
+				db, closeDB := newOnDiskDB(b, ctx)
+				col := setupIndexedCollection(b, ctx, db)
+				docIDs := addRecords(b, ctx, col, docs)
+				concrete, ok := col.(*collection)
+				require.True(b, ok)
+				b.StartTimer()
+
+				for i := 0; i < len(docIDs); i += chunkSize {
+					end := min(i+chunkSize, len(docIDs))
+					require.NoError(b, concrete.purgeChunk(ctx, docIDs[i:end], true))
+				}
+
+				b.StopTimer()
+				closeDB()
+				b.StartTimer()
+			}
+		})
+	}
+}