[v0.21.x-branch] Backport #11413: htlcswitch: avoid per-row queries when pruning network result store

docs/release-notes/unreleased/11413.md

@@ -0,0 +1,9 @@
+# Bug Fixes
+
+* [Fixed slow startup cleanup on SQL-backed
+  nodes](https://github.com/lightningnetwork/lnd/pull/11413): pruning the
+  switch's network result store at startup issued one database round-trip per
+  stored entry. The store accumulates one entry per locally initiated outbound
+  HTLC attempt, so cleanup on a busy node using remote Postgres could take
+  hours. The store is now pruned with a single streamed read and one cascading
+  delete instead of a round-trip per entry.

htlcswitch/payment_result.go

@@ -267,32 +267,72 @@ func (store *networkResultStore) cleanStore(keep map[uint64]struct{}) error {
 			return err
 		}
 
-		// Iterate through the bucket, deleting all items not in the
-		// keep map.
-		var toClean [][]byte
-		if err := networkResults.ForEach(func(k, _ []byte) error {
+		// Snapshot the entries we want to keep (the results of still
+		// in-flight attempts) and count the rest. The store is only
+		// pruned here at startup, so it can hold one entry per resolved
+		// HTLC attempt and grow very large. Deleting those one key at a
+		// time is O(N) round-trips on a SQL backend; since the keep set
+		// is small, we instead read the keepers with a single streaming
+		// query, drop the whole bucket in one cascading delete, and
+		// re-insert the keepers.
+		type keptResult struct {
+			key   []byte
+			value []byte
+		}
+		var (
+			kept    []keptResult
+			removed int
+		)
+		if err := kvdb.ForAll(networkResults, func(k, v []byte) error {
 			pid := binary.BigEndian.Uint64(k)
-			if _, ok := keep[pid]; ok {
+			if _, ok := keep[pid]; !ok {
+				removed++
 				return nil
 			}
 
-			toClean = append(toClean, k)
+			// The bucket is deleted below, so clone the retained
+			// entries before stashing them.
+			kept = append(kept, keptResult{
+				key:   bytes.Clone(k),
+				value: bytes.Clone(v),
+			})
+
 			return nil
 		}); err != nil {
 			return err
 		}
 
-		for _, k := range toClean {
-			err := networkResults.Delete(k)
+		// If there is nothing to remove, avoid dropping and recreating
+		// the bucket.
+		if removed == 0 {
+			return nil
+		}
+
+		// Drop the whole bucket in a single query (on the SQL backend
+		// this cascades to every row), then recreate it and restore the
+		// entries we are keeping.
+		if err := tx.DeleteTopLevelBucket(
+			networkResultStoreBucketKey,
+		); err != nil {
+			return err
+		}
+
+		networkResults, err = tx.CreateTopLevelBucket(
+			networkResultStoreBucketKey,
+		)
+		if err != nil {
+			return err
+		}
+
+		for _, r := range kept {
+			err := networkResults.Put(r.key, r.value)
 			if err != nil {
 				return err
 			}
 		}
 
-		if len(toClean) > 0 {
-			log.Infof("Removed %d stale entries from network "+
-				"result store", len(toClean))
-		}
+		log.Infof("Removed %d stale entries from network result store",
+			removed)
 
 		return nil
 	}, func() {})

htlcswitch/payment_result_test.go

@@ -2,13 +2,15 @@ package htlcswitch
 
 import (
 	"bytes"
+	"encoding/binary"
 	"math/rand"
 	"reflect"
 	"testing"
 	"time"
 
 	"github.com/davecgh/go-spew/spew"
 	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/kvdb"
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/stretchr/testify/require"
@@ -180,24 +182,201 @@ func TestNetworkResultStore(t *testing.T) {
 			t.Fatalf("no result received")
 		}
 	}
+}
+
+// TestNetworkResultStoreCleanStore tests that cleanStore prunes the store down
+// to exactly the requested attempt IDs. It opens the store via
+// kvdb.GetTestBackend so that, under the postgres/sqlite build tags, the SQL
+// drop-and-recreate path is exercised (TestNetworkResultStore above goes
+// through channeldb.OpenForTesting, which always opens bbolt).
+func TestNetworkResultStoreCleanStore(t *testing.T) {
+	if kvdb.PostgresBackend {
+		stop, err := kvdb.StartEmbeddedPostgres()
+		require.NoError(t, err)
 
-	// Clean the store keeping the first two results.
-	toKeep := map[uint64]struct{}{
-		0: {},
-		1: {},
+		t.Cleanup(func() {
+			require.NoError(t, stop())
+		})
+	} else {
+		t.Parallel()
 	}
-	// Finally, delete the result.
-	err = store.cleanStore(toKeep)
-	require.NoError(t, err)
 
-	// Payment IDs 0 and 1 should be found, 2 and 3 should be deleted.
-	for i := uint64(0); i < numResults; i++ {
-		_, err = store.getResult(i)
-		if i <= 1 {
-			require.NoError(t, err, "unable to get result")
-		}
-		if i >= 2 && err != ErrPaymentIDNotFound {
-			t.Fatalf("expected ErrPaymentIDNotFound, got %v", err)
-		}
+	tests := []struct {
+		name       string
+		numResults uint64
+		keep       map[uint64]struct{}
+	}{
+		{
+			name:       "partial keep",
+			numResults: 5,
+			keep:       map[uint64]struct{}{0: {}, 2: {}, 4: {}},
+		},
+		{
+			name:       "empty keep",
+			numResults: 5,
+			keep:       map[uint64]struct{}{},
+		},
+		{
+			name:       "keep all",
+			numResults: 5,
+			keep: map[uint64]struct{}{
+				0: {}, 1: {}, 2: {}, 3: {}, 4: {},
+			},
+		},
+		{
+			name:       "empty store",
+			numResults: 0,
+			keep:       map[uint64]struct{}{0: {}},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			// GetTestBackend keys its fixture off the unique temp
+			// dir, so parallel subtests get independent databases.
+			backend, cleanup, err := kvdb.GetTestBackend(
+				t.TempDir(), "db",
+			)
+			require.NoError(t, err)
+			t.Cleanup(func() {
+				require.NoError(t, backend.Close())
+				cleanup()
+			})
+
+			store := newNetworkResultStore(backend)
+
+			// Store a result for every attempt ID.
+			for i := range test.numResults {
+				msg := &lnwire.UpdateAddHTLC{ID: i}
+				n := &networkResult{
+					msg:          msg,
+					unencrypted:  true,
+					isResolution: true,
+				}
+				require.NoError(t, store.storeResult(i, n))
+			}
+
+			// Prune, keeping only the IDs in the keep set.
+			require.NoError(t, store.cleanStore(test.keep))
+
+			// Every kept ID must still be present; every other ID
+			// must have been removed.
+			for i := range test.numResults {
+				result, err := store.getResult(i)
+				if _, ok := test.keep[i]; ok {
+					require.NoErrorf(t, err,
+						"kept id %d missing", i)
+					msg := result.msg
+					add, ok := msg.(*lnwire.UpdateAddHTLC)
+					require.True(t, ok)
+					require.Equal(t, i, add.ID)
+
+					continue
+				}
+
+				require.ErrorIsf(t, err, ErrPaymentIDNotFound,
+					"id %d not removed", i)
+			}
+		})
+	}
+}
+
+// BenchmarkPostgresNetworkResultStoreCleanStore benchmarks pruning the network
+// result store at different sizes on PostgreSQL, where individual bucket
+// operations require database round trips.
+func BenchmarkPostgresNetworkResultStoreCleanStore(b *testing.B) {
+	if !kvdb.PostgresBackend {
+		b.Skip("requires the kvdb_postgres build tag")
+	}
+
+	stop, err := kvdb.StartEmbeddedPostgres()
+	require.NoError(b, err)
+	b.Cleanup(func() {
+		require.NoError(b, stop())
+	})
+
+	tests := []struct {
+		name       string
+		numResults uint64
+		numKeep    uint64
+	}{
+		{
+			name:       "1k_all_stale",
+			numResults: 1_000,
+		},
+		{
+			name:       "10k_all_stale",
+			numResults: 10_000,
+		},
+		{
+			name:       "10k_keep_10",
+			numResults: 10_000,
+			numKeep:    10,
+		},
+	}
+
+	for _, test := range tests {
+		b.Run(test.name, func(b *testing.B) {
+			backend, cleanup, err := kvdb.GetTestBackend(
+				b.TempDir(), "db",
+			)
+			require.NoError(b, err)
+			b.Cleanup(func() {
+				require.NoError(b, backend.Close())
+				cleanup()
+			})
+
+			store := newNetworkResultStore(backend)
+			keep := make(map[uint64]struct{}, test.numKeep)
+			for i := range test.numKeep {
+				keep[i] = struct{}{}
+			}
+
+			var result bytes.Buffer
+			err = serializeNetworkResult(&result, &networkResult{
+				msg:          &lnwire.UpdateAddHTLC{},
+				unencrypted:  true,
+				isResolution: true,
+			})
+			require.NoError(b, err)
+
+			// Populate results in one transaction.
+			// This setup is outside the measured benchmark.
+			update := func(tx kvdb.RwTx) error {
+				bucket, err := tx.CreateTopLevelBucket(
+					networkResultStoreBucketKey,
+				)
+				if err != nil {
+					return err
+				}
+
+				var key [8]byte
+				for i := range test.numResults {
+					binary.BigEndian.PutUint64(key[:], i)
+					if err := bucket.Put(
+						key[:], result.Bytes(),
+					); err != nil {
+						return err
+					}
+				}
+
+				return nil
+			}
+			populate := func() error {
+				return kvdb.Update(backend, update, func() {})
+			}
+
+			b.ReportAllocs()
+			b.ResetTimer()
+			for range b.N {
+				b.StopTimer()
+				require.NoError(b, populate())
+				b.StartTimer()
+
+				require.NoError(b, store.cleanStore(keep))
+			}
+		})
 	}
 }