resource: build sub-graphs from a node's dependents instead of pruning a clone

resource/resource_graph.go

@@ -1091,19 +1091,66 @@ func (g *Graph) SubGraphFrom(node Name) (*Graph, error) {
 	return g.subGraphFromWithMutex(node)
 }
 
+// dependentsWithMutex returns node and every node that depends on it, directly or
+// transitively, read from the transitive closure. This method is NOT threadsafe: A client
+// must read-lock [Graph.mu] while calling this method.
+func (g *Graph) dependentsWithMutex(node Name) (map[Name]*GraphNode, error) {
+	if _, ok := g.nodes.Get(node); !ok {
+		return nil, errors.Errorf("cannot create sub-graph from non existing node %v ", node.Name)
+	}
+	dependents := map[Name]*GraphNode{}
+	for name, gNode := range g.nodes.All() {
+		if g.transitiveClosureMatrix[name][node] != 0 {
+			dependents[name] = gNode
+		}
+	}
+	return dependents, nil
+}
+
 // subGraphFrom returns a Sub-Graph containing all linked dependencies starting with node
 // [Name]. This method is NOT threadsafe: A client must read-lock [Graph.mu] while calling
 // this method.
+//
+// The sub-graph holds node and its dependents, the same *GraphNode values as g, and the
+// edges and transitive closure among them. Every path between two of those nodes stays
+// among them, so the closure is the same as g's. It is built from just those nodes rather
+// than by removing every other node from a clone of g, which costs a transitive closure
+// update per removed edge.
 func (g *Graph) subGraphFromWithMutex(node Name) (*Graph, error) {
-	if _, ok := g.nodes.Get(node); !ok {
-		return nil, errors.Errorf("cannot create sub-graph from non existing node %v ", node.Name)
+	dependents, err := g.dependentsWithMutex(node)
+	if err != nil {
+		return nil, err
+	}
+	subGraph := &Graph{
+		children:                resourceDependencies{},
+		nodes:                   g.nodes.Copy(),
+		parents:                 resourceDependencies{},
+		transitiveClosureMatrix: make(transitiveClosureMatrix, len(dependents)),
+		logicalClock:            g.logicalClock,
+	}
+	for name := range g.nodes.Keys() {
+		if _, ok := dependents[name]; !ok {
+			subGraph.nodes.Delete(name)
+		}
 	}
-	subGraph := g.clone()
-	sorted := subGraph.ReverseTopologicalSort()
-	for _, n := range sorted {
-		if !subGraph.isNodeDependingOn(node, n) {
-			subGraph.remove(n)
+	for name := range dependents {
+		// A dependent's children depend on it, so they are all dependents too.
+		if children, ok := g.children[name]; ok {
+			subGraph.children[name] = copyNodes(children)
+		}
+		for parent, edge := range g.parents[name] {
+			if _, ok := dependents[parent]; ok {
+				if subGraph.parents[name] == nil {
+					subGraph.parents[name] = graphNodes{}
+				}
+				subGraph.parents[name][parent] = edge
+			}
+		}
+		row := make(map[Name]int, len(dependents))
+		for other := range dependents {
+			row[other] = g.transitiveClosureMatrix[name][other]
 		}
+		subGraph.transitiveClosureMatrix[name] = row
 	}
 	return subGraph, nil
 }
@@ -1114,13 +1161,13 @@ func (g *Graph) MarkReachability(node Name, reachable bool) error {
 	g.mu.RLock()
 	defer g.mu.RUnlock()
 
-	subGraph, err := g.subGraphFromWithMutex(node)
+	dependents, err := g.dependentsWithMutex(node)
 	if err != nil {
 		return err
 	}
-	for name, n := range subGraph.nodes.All() {
+	for name, gNode := range dependents {
 		if name == node || name.Remote != "" {
-			n.markReachability(reachable)
+			gNode.markReachability(reachable)
 		}
 	}
 	return nil

resource/resource_graph_bench_test.go

@@ -0,0 +1,82 @@
+package resource
+
+import (
+	"fmt"
+	"testing"
+
+	"go.viam.com/test"
+
+	"go.viam.com/rdk/logging"
+)
+
+// robotSizedGraph builds a graph shaped like a robot with one remote: the
+// remote and 50 resources on it, plus 60 local resources in three layers
+// (components, services on components and remote resources, and services on
+// services). It returns the graph, the remote's name and a component with a
+// dozen dependents.
+func robotSizedGraph(tb testing.TB) (*Graph, Name, Name) {
+	tb.Helper()
+	g := NewGraph(logging.NewTestLogger(tb))
+	add := func(name Name, deps ...Name) {
+		test.That(tb, g.AddNode(name, &GraphNode{}), test.ShouldBeNil)
+		for _, dep := range deps {
+			test.That(tb, g.AddChild(name, dep), test.ShouldBeNil)
+		}
+	}
+
+	remote := NewName(APINamespace("rdk").WithType("component").WithSubtype("remote"), "remote1")
+	add(remote)
+	remoteRes := make([]Name, 50)
+	for i := range remoteRes {
+		remoteRes[i] = NewName(apiA, fmt.Sprintf("r%d", i)).PrependRemote(remote.Name)
+		add(remoteRes[i], remote)
+	}
+
+	components := make([]Name, 20)
+	for i := range components {
+		components[i] = NewName(apiA, fmt.Sprintf("c%d", i))
+		if i < 5 {
+			add(components[i])
+		} else {
+			add(components[i], components[i%5])
+		}
+	}
+	services := make([]Name, 25)
+	for i := range services {
+		services[i] = NewName(apiA, fmt.Sprintf("s%d", i))
+		deps := []Name{components[i%20], components[(i*7+1)%20], components[(i*13+2)%20]}
+		if i%2 == 0 {
+			deps = append(deps, remoteRes[i%50], remoteRes[(i*3+1)%50])
+		}
+		add(services[i], deps...)
+	}
+	for i := 0; i < 15; i++ {
+		add(NewName(apiA, fmt.Sprintf("h%d", i)), services[i%25], services[(i*7+3)%25], components[(i*11)%20])
+	}
+	return g, remote, components[3]
+}
+
+// The remote refresh marks the remote and the remote resources depending on it
+// reachable on every pass.
+func BenchmarkGraphMarkReachability(b *testing.B) {
+	g, remote, _ := robotSizedGraph(b)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if err := g.MarkReachability(remote, true); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// Marking a resource's dependents for update starts from its sub-graph.
+func BenchmarkGraphSubGraphFrom(b *testing.B) {
+	g, _, component := robotSizedGraph(b)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := g.SubGraphFrom(component); err != nil {
+			b.Fatal(err)
+		}
+	}
+}

resource/resource_graph_test.go

@@ -265,6 +265,21 @@ func TestResourceGraphSubGraph(t *testing.T) {
 		NewName(apiA, "F"),
 		NewName(apiA, "C"),
 	}...))
+
+	// D and F also depend on B and A, which are not in the sub-graph, so those edges are
+	// left out of it.
+	c, d, f := NewName(apiA, "C"), NewName(apiA, "D"), NewName(apiA, "F")
+	test.That(t, newResourceNameSet(sg.GetAllChildrenOf(c)...), test.ShouldResemble, newResourceNameSet(d, f))
+	test.That(t, sg.GetAllChildrenOf(d), test.ShouldBeEmpty)
+	test.That(t, sg.GetAllChildrenOf(f), test.ShouldBeEmpty)
+	test.That(t, sg.GetAllParentsOf(c), test.ShouldBeEmpty)
+	test.That(t, sg.GetAllParentsOf(d), test.ShouldResemble, []Name{c})
+	test.That(t, sg.GetAllParentsOf(f), test.ShouldResemble, []Name{c})
+	test.That(t, sg.transitiveClosureMatrix, test.ShouldResemble, transitiveClosureMatrix{
+		c: {c: 1, d: 0, f: 0},
+		d: {c: 1, d: 1, f: 0},
+		f: {c: 1, d: 0, f: 1},
+	})
 }
 
 func TestResourceGraphDepTree(t *testing.T) {