Merge pull request #10644 from djdongjin/improve-gc-tricolor

Improve gc tricolor by avoiding revisiting nodes

pkg/gc/gc.go

@@ -69,12 +69,14 @@ func Tricolor(roots []Node, refs func(ref Node) ([]Node, error)) (map[Node]struc
 	)
 
 	grays = append(grays, roots...)
+	for _, root := range roots {
+		seen[root] = struct{}{} // pre-mark this as not-white
+	}
 
 	for len(grays) > 0 {
 		// Pick any gray object
 		id := grays[len(grays)-1] // effectively "depth first" because first element
 		grays = grays[:len(grays)-1]
-		seen[id] = struct{}{} // post-mark this as not-white
 		rs, err := refs(id)
 		if err != nil {
 			return nil, err
@@ -84,6 +86,7 @@ func Tricolor(roots []Node, refs func(ref Node) ([]Node, error)) (map[Node]struc
 		for _, target := range rs {
 			if _, ok := seen[target]; !ok {
 				grays = append(grays, target)
+				seen[target] = struct{}{}
 			}
 		}
 

pkg/gc/gc_test.go

@@ -18,6 +18,7 @@ package gc
 
 import (
 	"context"
+	"fmt"
 	"reflect"
 	"testing"
 )
@@ -52,6 +53,28 @@ func TestTricolorBasic(t *testing.T) {
 	}
 }
 
+func BenchmarkTricolor(b *testing.B) {
+	roots := []string{"A", "C"}
+	refs := map[string][]string{
+		"A": {"B", "C"},
+		"B": {"A", "C"},
+		"C": {"D", "F", "B"},
+		"E": {"F", "G", "C"},
+		"F": {"H", "C"},
+	}
+	// assume 100 nodes can reach D.
+	for i := 0; i < 100; i++ {
+		ref := fmt.Sprintf("X%d", i)
+		refs["A"] = append(refs["A"], ref)
+		refs[ref] = []string{"D"}
+	}
+
+	// Run the Add function b.N times to benchmark it.
+	for i := 0; i < b.N; i++ {
+		Tricolor(toNodes(roots), lookup(refs))
+	}
+}
+
 func TestConcurrentBasic(t *testing.T) {
 	roots := []string{"A", "C"}
 	all := []string{"A", "B", "C", "D", "E", "F", "G", "H", "I"}