Merge pull request #10121 from ellemouton/graphPerf6

[5] sqldb+graph/db: add and use new paginate & batch helper

Author: ellemouton

graph/db/graph_test.go

@@ -2510,6 +2510,11 @@ func TestFilterKnownChanIDs(t *testing.T) {
 // methods that acquire the cache mutex along with the DB mutex.
 func TestStressTestChannelGraphAPI(t *testing.T) {
 	t.Parallel()
+
+	if testing.Short() {
+		t.Skipf("Skipping test in short mode")
+	}
+
 	ctx := context.Background()
 
 	graph := MakeTestGraph(t)

graph/db/sql_migration_test.go

@@ -1177,6 +1177,10 @@ func assertResultState(t *testing.T, sql *SQLStore, expState dbState) {
 func TestMigrateGraphToSQLRapid(t *testing.T) {
 	t.Parallel()
 
+	if testing.Short() {
+		t.Skipf("skipping test in short mode")
+	}
+
 	dbFixture := NewTestDBFixture(t)
 
 	rapid.Check(t, func(rt *rapid.T) {

graph/db/sql_store.go

@@ -806,49 +806,18 @@ func (s *SQLStore) ForEachSourceNodeChannel(ctx context.Context,
 func (s *SQLStore) ForEachNode(ctx context.Context,
 	cb func(tx NodeRTx) error, reset func()) error {
 
-	var lastID int64
-
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
-		nodeCB := func(dbID int64, node *models.LightningNode) error {
-			err := cb(newSQLGraphNodeTx(
-				db, s.cfg.ChainHash, dbID, node,
-			))
-			if err != nil {
-				return fmt.Errorf("callback failed for "+
-					"node(id=%d): %w", dbID, err)
-			}
-			lastID = dbID
-
-			return nil
-		}
-
-		for {
-			nodes, err := db.ListNodesPaginated(
-				ctx, sqlc.ListNodesPaginatedParams{
-					Version: int16(ProtocolV1),
-					ID:      lastID,
-					Limit:   s.cfg.QueryCfg.MaxPageSize,
-				},
-			)
-			if err != nil {
-				return fmt.Errorf("unable to fetch nodes: %w",
-					err)
-			}
-
-			if len(nodes) == 0 {
-				break
-			}
-
-			err = forEachNodeInBatch(
-				ctx, s.cfg.QueryCfg, db, nodes, nodeCB,
-			)
-			if err != nil {
-				return fmt.Errorf("unable to iterate over "+
-					"nodes: %w", err)
-			}
-		}
+		return forEachNodePaginated(
+			ctx, s.cfg.QueryCfg, db,
+			ProtocolV1,
+			func(ctx context.Context, dbNodeID int64,
+				node *models.LightningNode) error {
 
-		return nil
+				return cb(newSQLGraphNodeTx(
+					db, s.cfg.ChainHash, dbNodeID, node,
+				))
+			},
+		)
 	}, reset)
 }
 
@@ -1328,115 +1297,8 @@ func (s *SQLStore) ForEachChannel(ctx context.Context,
 	cb func(*models.ChannelEdgeInfo, *models.ChannelEdgePolicy,
 		*models.ChannelEdgePolicy) error, reset func()) error {
 
-	handleChannel := func(db SQLQueries, batchData *batchChannelData,
-		row sqlc.ListChannelsWithPoliciesPaginatedRow) error {
-
-		node1, node2, err := buildNodeVertices(
-			row.Node1Pubkey, row.Node2Pubkey,
-		)
-		if err != nil {
-			return fmt.Errorf("unable to build node vertices: %w",
-				err)
-		}
-
-		edge, err := buildEdgeInfoWithBatchData(
-			s.cfg.ChainHash, row.GraphChannel, node1, node2,
-			batchData,
-		)
-		if err != nil {
-			return fmt.Errorf("unable to build channel info: %w",
-				err)
-		}
-
-		dbPol1, dbPol2, err := extractChannelPolicies(row)
-		if err != nil {
-			return fmt.Errorf("unable to extract channel "+
-				"policies: %w", err)
-		}
-
-		p1, p2, err := buildChanPoliciesWithBatchData(
-			dbPol1, dbPol2, edge.ChannelID, node1, node2, batchData,
-		)
-		if err != nil {
-			return fmt.Errorf("unable to build channel "+
-				"policies: %w", err)
-		}
-
-		err = cb(edge, p1, p2)
-		if err != nil {
-			return fmt.Errorf("callback failed for channel "+
-				"id=%d: %w", edge.ChannelID, err)
-		}
-
-		return nil
-	}
-
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
-		lastID := int64(-1)
-		for {
-			//nolint:ll
-			rows, err := db.ListChannelsWithPoliciesPaginated(
-				ctx, sqlc.ListChannelsWithPoliciesPaginatedParams{
-					Version: int16(ProtocolV1),
-					ID:      lastID,
-					Limit:   s.cfg.QueryCfg.MaxPageSize,
-				},
-			)
-			if err != nil {
-				return err
-			}
-
-			if len(rows) == 0 {
-				break
-			}
-
-			// Collect the channel & policy IDs that we want to
-			// do a batch collection for.
-			var (
-				channelIDs = make([]int64, len(rows))
-				policyIDs  = make([]int64, 0, len(rows)*2)
-			)
-			for i, row := range rows {
-				channelIDs[i] = row.GraphChannel.ID
-
-				// Extract policy IDs from the row
-				dbPol1, dbPol2, err := extractChannelPolicies(
-					row,
-				)
-				if err != nil {
-					return fmt.Errorf("unable to extract "+
-						"channel policies: %w", err)
-				}
-
-				if dbPol1 != nil {
-					policyIDs = append(policyIDs, dbPol1.ID)
-				}
-
-				if dbPol2 != nil {
-					policyIDs = append(policyIDs, dbPol2.ID)
-				}
-			}
-
-			batchData, err := batchLoadChannelData(
-				ctx, s.cfg.QueryCfg, db, channelIDs,
-				policyIDs,
-			)
-			if err != nil {
-				return fmt.Errorf("unable to batch load "+
-					"channel data: %w", err)
-			}
-
-			for _, row := range rows {
-				err := handleChannel(db, batchData, row)
-				if err != nil {
-					return err
-				}
-
-				lastID = row.GraphChannel.ID
-			}
-		}
-
-		return nil
+		return forEachChannelWithPolicies(ctx, db, s.cfg, cb)
 	}, reset)
 }
 
@@ -5082,3 +4944,169 @@ func batchLoadChannelPolicyExtrasHelper(ctx context.Context,
 		},
 	)
 }
+
+// forEachNodePaginated executes a paginated query to process each node in the
+// graph. It uses the provided SQLQueries interface to fetch nodes in batches
+// and applies the provided processNode function to each node.
+func forEachNodePaginated(ctx context.Context, cfg *sqldb.QueryConfig,
+	db SQLQueries, protocol ProtocolVersion,
+	processNode func(context.Context, int64,
+		*models.LightningNode) error) error {
+
+	pageQueryFunc := func(ctx context.Context, lastID int64,
+		limit int32) ([]sqlc.GraphNode, error) {
+
+		return db.ListNodesPaginated(
+			ctx, sqlc.ListNodesPaginatedParams{
+				Version: int16(protocol),
+				ID:      lastID,
+				Limit:   limit,
+			},
+		)
+	}
+
+	extractPageCursor := func(node sqlc.GraphNode) int64 {
+		return node.ID
+	}
+
+	collectFunc := func(node sqlc.GraphNode) (int64, error) {
+		return node.ID, nil
+	}
+
+	batchQueryFunc := func(ctx context.Context,
+		nodeIDs []int64) (*batchNodeData, error) {
+
+		return batchLoadNodeData(ctx, cfg, db, nodeIDs)
+	}
+
+	processItem := func(ctx context.Context, dbNode sqlc.GraphNode,
+		batchData *batchNodeData) error {
+
+		node, err := buildNodeWithBatchData(&dbNode, batchData)
+		if err != nil {
+			return fmt.Errorf("unable to build "+
+				"node(id=%d): %w", dbNode.ID, err)
+		}
+
+		return processNode(ctx, dbNode.ID, node)
+	}
+
+	return sqldb.ExecuteCollectAndBatchWithSharedDataQuery(
+		ctx, cfg, int64(-1), pageQueryFunc, extractPageCursor,
+		collectFunc, batchQueryFunc, processItem,
+	)
+}
+
+// forEachChannelWithPolicies executes a paginated query to process each channel
+// with policies in the graph.
+func forEachChannelWithPolicies(ctx context.Context, db SQLQueries,
+	cfg *SQLStoreConfig, processChannel func(*models.ChannelEdgeInfo,
+		*models.ChannelEdgePolicy,
+		*models.ChannelEdgePolicy) error) error {
+
+	type channelBatchIDs struct {
+		channelID int64
+		policyIDs []int64
+	}
+
+	pageQueryFunc := func(ctx context.Context, lastID int64,
+		limit int32) ([]sqlc.ListChannelsWithPoliciesPaginatedRow,
+		error) {
+
+		return db.ListChannelsWithPoliciesPaginated(
+			ctx, sqlc.ListChannelsWithPoliciesPaginatedParams{
+				Version: int16(ProtocolV1),
+				ID:      lastID,
+				Limit:   limit,
+			},
+		)
+	}
+
+	extractPageCursor := func(
+		row sqlc.ListChannelsWithPoliciesPaginatedRow) int64 {
+
+		return row.GraphChannel.ID
+	}
+
+	collectFunc := func(row sqlc.ListChannelsWithPoliciesPaginatedRow) (
+		channelBatchIDs, error) {
+
+		ids := channelBatchIDs{
+			channelID: row.GraphChannel.ID,
+		}
+
+		// Extract policy IDs from the row.
+		dbPol1, dbPol2, err := extractChannelPolicies(row)
+		if err != nil {
+			return ids, err
+		}
+
+		if dbPol1 != nil {
+			ids.policyIDs = append(ids.policyIDs, dbPol1.ID)
+		}
+		if dbPol2 != nil {
+			ids.policyIDs = append(ids.policyIDs, dbPol2.ID)
+		}
+
+		return ids, nil
+	}
+
+	batchDataFunc := func(ctx context.Context,
+		allIDs []channelBatchIDs) (*batchChannelData, error) {
+
+		// Separate channel IDs from policy IDs.
+		var (
+			channelIDs = make([]int64, len(allIDs))
+			policyIDs  = make([]int64, 0, len(allIDs)*2)
+		)
+
+		for i, ids := range allIDs {
+			channelIDs[i] = ids.channelID
+			policyIDs = append(policyIDs, ids.policyIDs...)
+		}
+
+		return batchLoadChannelData(
+			ctx, cfg.QueryCfg, db, channelIDs, policyIDs,
+		)
+	}
+
+	processItem := func(ctx context.Context,
+		row sqlc.ListChannelsWithPoliciesPaginatedRow,
+		batchData *batchChannelData) error {
+
+		node1, node2, err := buildNodeVertices(
+			row.Node1Pubkey, row.Node2Pubkey,
+		)
+		if err != nil {
+			return err
+		}
+
+		edge, err := buildEdgeInfoWithBatchData(
+			cfg.ChainHash, row.GraphChannel, node1, node2,
+			batchData,
+		)
+		if err != nil {
+			return fmt.Errorf("unable to build channel info: %w",
+				err)
+		}
+
+		dbPol1, dbPol2, err := extractChannelPolicies(row)
+		if err != nil {
+			return err
+		}
+
+		p1, p2, err := buildChanPoliciesWithBatchData(
+			dbPol1, dbPol2, edge.ChannelID, node1, node2, batchData,
+		)
+		if err != nil {
+			return err
+		}
+
+		return processChannel(edge, p1, p2)
+	}
+
+	return sqldb.ExecuteCollectAndBatchWithSharedDataQuery(
+		ctx, cfg.QueryCfg, int64(-1), pageQueryFunc, extractPageCursor,
+		collectFunc, batchDataFunc, processItem,
+	)
+}

make/testing_flags.mk

@@ -121,6 +121,12 @@ ifneq ($(nocache),)
 TEST_FLAGS += -test.count=1
 endif
 
+# If the short flag is added, then any unit tests marked with "testing.Short()"
+# will be skipped.
+ifneq ($(short),)
+TEST_FLAGS += -short
+endif
+
 GOLIST := $(GOCC) list -tags="$(DEV_TAGS)" -deps $(PKG)/... | grep '$(PKG)'| grep -v '/vendor/'
 
 # UNIT_TARGTED is undefined iff a specific package and/or unit test case is