graph/db: allow async cache population

We also add a functional option to opt out of this behaviour so that any
tests that rely on the graph cache to continue to work reliably.

Author: ellemouton

graph/db/graph.go

@@ -31,6 +31,8 @@ type ChannelGraph struct {
 	started atomic.Bool
 	stopped atomic.Bool
 
+	opts *chanGraphOptions
+
 	// cacheLoaded is true if the initial graphCache population has
 	// finished. We use this to ensure that when performing any reads,
 	// we only read from the graphCache if it has been fully populated.
@@ -54,6 +56,7 @@ func NewChannelGraph(v1Store V1Store,
 	}
 
 	g := &ChannelGraph{
+		opts:            opts,
 		V1Store:         v1Store,
 		topologyManager: newTopologyManager(),
 		quit:            make(chan struct{}),
@@ -78,8 +81,22 @@ func (c *ChannelGraph) Start() error {
 	log.Debugf("ChannelGraph starting")
 	defer log.Debug("ChannelGraph started")
 
-	if err := c.populateCache(context.TODO()); err != nil {
-		return fmt.Errorf("could not populate the graph cache: %w", err)
+	ctx := context.TODO()
+	if c.opts.asyncGraphCachePopulation {
+		c.wg.Add(1)
+		go func() {
+			defer c.wg.Done()
+
+			if err := c.populateCache(ctx); err != nil {
+				log.Errorf("Could not populate the graph "+
+					"cache: %v", err)
+			}
+		}()
+	} else {
+		if err := c.populateCache(ctx); err != nil {
+			return fmt.Errorf("could not populate the graph "+
+				"cache: %w", err)
+		}
 	}
 
 	c.wg.Add(1)

graph/db/graph_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/lightningnetwork/lnd/fn/v2"
 	"github.com/lightningnetwork/lnd/graph/db/models"
 	"github.com/lightningnetwork/lnd/kvdb"
+	"github.com/lightningnetwork/lnd/lntest/wait"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/routing/route"
 	"github.com/stretchr/testify/require"
@@ -102,7 +103,7 @@ func TestNodeInsertionAndDeletion(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
-	graph := MakeTestGraph(t)
+	graph := MakeTestGraph(t, WithSyncGraphCachePopulation())
 
 	// We'd like to test basic insertion/deletion for vertexes from the
 	// graph, so we'll create a test vertex to start with.
@@ -260,7 +261,7 @@ func TestPartialNode(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
-	graph := MakeTestGraph(t)
+	graph := MakeTestGraph(t, WithSyncGraphCachePopulation())
 
 	// To insert a partial node, we need to add a channel edge that has
 	// node keys for nodes we are not yet aware
@@ -386,7 +387,7 @@ func TestEdgeInsertionDeletion(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
-	graph := MakeTestGraph(t)
+	graph := MakeTestGraph(t, WithSyncGraphCachePopulation())
 
 	// We'd like to test the insertion/deletion of edges, so we create two
 	// vertexes to connect.
@@ -511,7 +512,7 @@ func TestDisconnectBlockAtHeight(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
-	graph := MakeTestGraph(t)
+	graph := MakeTestGraph(t, WithSyncGraphCachePopulation())
 
 	sourceNode := createTestVertex(t)
 	if err := graph.SetSourceNode(ctx, sourceNode); err != nil {
@@ -802,7 +803,7 @@ func TestEdgeInfoUpdates(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
-	graph := MakeTestGraph(t)
+	graph := MakeTestGraph(t, WithSyncGraphCachePopulation())
 
 	// We'd like to test the update of edges inserted into the database, so
 	// we create two vertexes to connect.
@@ -4350,7 +4351,9 @@ func TestGraphLoading(t *testing.T) {
 	// Next, create the graph for the first time.
 	graphStore := NewTestDB(t)
 
-	graph, err := NewChannelGraph(graphStore)
+	graph, err := NewChannelGraph(
+		graphStore, WithSyncGraphCachePopulation(),
+	)
 	require.NoError(t, err)
 	require.NoError(t, graph.Start())
 	t.Cleanup(func() {
@@ -4364,7 +4367,9 @@ func TestGraphLoading(t *testing.T) {
 
 	// Recreate the graph. This should cause the graph cache to be
 	// populated.
-	graphReloaded, err := NewChannelGraph(graphStore)
+	graphReloaded, err := NewChannelGraph(
+		graphStore, WithSyncGraphCachePopulation(),
+	)
 	require.NoError(t, err)
 	require.NoError(t, graphReloaded.Start())
 	t.Cleanup(func() {
@@ -4383,6 +4388,100 @@ func TestGraphLoading(t *testing.T) {
 	)
 }
 
+// TestAsyncGraphCache tests the behaviour of the ChannelGraph when the graph
+// cache is populated asynchronously.
+func TestAsyncGraphCache(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	const (
+		numNodes    = 100
+		numChannels = 3
+	)
+
+	// Next, create the graph for the first time.
+	graphStore := NewTestDB(t)
+
+	// The first time we spin up the graph, we Start is as normal and fill
+	// it with test data. This will ensure that the graph cache has
+	// something to load on the next Start.
+	graph, err := NewChannelGraph(graphStore)
+	require.NoError(t, err)
+	require.NoError(t, graph.Start())
+	channels, nodes := fillTestGraph(t, graph, numNodes, numChannels)
+
+	assertGraphState := func() {
+		var (
+			numNodes  int
+			chanIndex = make(map[uint64]struct{}, numChannels)
+		)
+		// We query the graph for all nodes and channels, and
+		// assert that we get the expected number of nodes and
+		// channels.
+		err := graph.ForEachNodeCached(
+			ctx, func(node route.Vertex,
+				chans map[uint64]*DirectedChannel) error {
+
+				numNodes++
+				for chanID := range chans {
+					chanIndex[chanID] = struct{}{}
+				}
+
+				return nil
+			},
+		)
+		require.NoError(t, err)
+
+		require.Equal(t, len(nodes), numNodes)
+		require.Equal(t, len(channels), len(chanIndex))
+	}
+
+	assertGraphState()
+
+	// Now we stop the graph.
+	require.NoError(t, graph.Stop())
+
+	// Recreate it but don't start it yet.
+	graph, err = NewChannelGraph(graphStore)
+	require.NoError(t, err)
+
+	// Spin off a goroutine that starts to make queries to the ChannelGraph.
+	// We start this before we start the graph, so that we can ensure that
+	// the queries are made while the graph cache is being populated.
+	var (
+		wg      sync.WaitGroup
+		numRuns = 10
+	)
+	for i := 0; i < numRuns; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			assertGraphState()
+		}()
+	}
+
+	require.NoError(t, graph.Start())
+	t.Cleanup(func() {
+		require.NoError(t, graph.Stop())
+	})
+
+	wg.Wait()
+
+	// Wait for the cache to be fully populated.
+	err = wait.Predicate(func() bool {
+		return graph.cacheLoaded.Load()
+	}, wait.DefaultTimeout)
+	require.NoError(t, err)
+
+	// And then assert that all the expected nodes and channels are
+	// present in the graph cache.
+	for _, node := range nodes {
+		_, ok := graph.graphCache.nodeChannels[node.PubKeyBytes]
+		require.True(t, ok)
+	}
+}
+
 // TestClosedScid tests that we can correctly insert a SCID into the index of
 // closed short channel ids.
 func TestClosedScid(t *testing.T) {

graph/db/options.go

@@ -31,14 +31,20 @@ type chanGraphOptions struct {
 	// preAllocCacheNumNodes is the number of nodes we expect to be in the
 	// graph cache, so we can pre-allocate the map accordingly.
 	preAllocCacheNumNodes int
+
+	// asyncGraphCachePopulation indicates whether the graph cache
+	// should be populated asynchronously or if the Start method should
+	// block until the cache is fully populated.
+	asyncGraphCachePopulation bool
 }
 
 // defaultChanGraphOptions returns a new chanGraphOptions instance populated
 // with default values.
 func defaultChanGraphOptions() *chanGraphOptions {
 	return &chanGraphOptions{
-		useGraphCache:         true,
-		preAllocCacheNumNodes: DefaultPreAllocCacheNumNodes,
+		useGraphCache:             true,
+		asyncGraphCachePopulation: true,
+		preAllocCacheNumNodes:     DefaultPreAllocCacheNumNodes,
 	}
 }
 
@@ -61,6 +67,25 @@ func WithPreAllocCacheNumNodes(n int) ChanGraphOption {
 	}
 }
 
+// WithAsyncGraphCachePopulation sets whether the graph cache should be
+// populated asynchronously or if the Start method should block until the
+// cache is fully populated.
+func WithAsyncGraphCachePopulation(async bool) ChanGraphOption {
+	return func(o *chanGraphOptions) {
+		o.asyncGraphCachePopulation = async
+	}
+}
+
+// WithSyncGraphCachePopulation will cause the ChannelGraph to block
+// until the graph cache is fully populated before returning from the Start
+// method. This is useful for tests that need to ensure the graph cache is
+// fully populated before proceeding with further operations.
+func WithSyncGraphCachePopulation() ChanGraphOption {
+	return func(o *chanGraphOptions) {
+		o.asyncGraphCachePopulation = false
+	}
+}
+
 // StoreOptions holds parameters for tuning and customizing a graph DB.
 type StoreOptions struct {
 	// RejectCacheSize is the maximum number of rejectCacheEntries to hold