rewrite enumerate children async to be less fragile

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>

merkledag.go

@@ -389,103 +389,92 @@ func EnumerateChildren(ctx context.Context, ds LinkService, root *cid.Cid, visit
 	return nil
 }
 
-func EnumerateChildrenAsync(ctx context.Context, ds DAGService, c *cid.Cid, visit func(*cid.Cid) bool) error {
-	toprocess := make(chan []*cid.Cid, 8)
-	nodes := make(chan *NodeOption, 8)
-
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	defer close(toprocess)
+// FetchGraphConcurrency is total number of concurrent fetches that
+// 'fetchNodes' will start at a time
+var FetchGraphConcurrency = 8
 
-	go fetchNodes(ctx, ds, toprocess, nodes)
+func EnumerateChildrenAsync(ctx context.Context, ds DAGService, c *cid.Cid, visit func(*cid.Cid) bool) error {
+	if !visit(c) {
+		return nil
+	}
 
 	root, err := ds.Get(ctx, c)
 	if err != nil {
 		return err
 	}
 
-	nodes <- &NodeOption{Node: root}
-	live := 1
+	feed := make(chan node.Node)
+	out := make(chan *NodeOption)
+	done := make(chan struct{})
 
-	for {
-		select {
-		case opt, ok := <-nodes:
-			if !ok {
-				return nil
-			}
-
-			if opt.Err != nil {
-				return opt.Err
-			}
-
-			nd := opt.Node
+	var setlk sync.Mutex
 
-			// a node has been fetched
-			live--
-
-			var cids []*cid.Cid
-			for _, lnk := range nd.Links() {
-				c := lnk.Cid
-				if visit(c) {
-					live++
-					cids = append(cids, c)
-				}
+	for i := 0; i < FetchGraphConcurrency; i++ {
+		go func() {
+			for n := range feed {
+				links := n.Links()
+				cids := make([]*cid.Cid, 0, len(links))
+				for _, l := range links {
+					setlk.Lock()
+					unseen := visit(l.Cid)
+					setlk.Unlock()
+					if unseen {
+						cids = append(cids, l.Cid)
 					}
-
-			if live == 0 {
-				return nil
 				}
 
-			if len(cids) > 0 {
+				for nopt := range ds.GetMany(ctx, cids) {
 					select {
-				case toprocess <- cids:
+					case out <- nopt:
 					case <-ctx.Done():
-					return ctx.Err()
+						return
 					}
 				}
+				select {
+				case done <- struct{}{}:
 				case <-ctx.Done():
-			return ctx.Err()
 				}
 			}
-}
-
-// FetchGraphConcurrency is total number of concurrenct fetches that
-// 'fetchNodes' will start at a time
-var FetchGraphConcurrency = 8
-
-func fetchNodes(ctx context.Context, ds DAGService, in <-chan []*cid.Cid, out chan<- *NodeOption) {
-	var wg sync.WaitGroup
-	defer func() {
-		// wait for all 'get' calls to complete so we don't accidentally send
-		// on a closed channel
-		wg.Wait()
-		close(out)
 		}()
+	}
+	defer close(feed)
 
-	rateLimit := make(chan struct{}, FetchGraphConcurrency)
+	send := feed
+	var todobuffer []node.Node
+	var inProgress int
 
-	get := func(ks []*cid.Cid) {
-		defer wg.Done()
-		defer func() {
-			<-rateLimit
-		}()
-		nodes := ds.GetMany(ctx, ks)
-		for opt := range nodes {
+	next := root
+	for {
 		select {
-			case out <- opt:
-			case <-ctx.Done():
-				return
+		case send <- next:
+			inProgress++
+			if len(todobuffer) > 0 {
+				next = todobuffer[0]
+				todobuffer = todobuffer[1:]
+			} else {
+				next = nil
+				send = nil
 			}
+		case <-done:
+			inProgress--
+			if inProgress == 0 && next == nil {
+				return nil
 			}
+		case nc := <-out:
+			if nc.Err != nil {
+				return nc.Err
+			}
+
+			if next == nil {
+				next = nc.Node
+				send = feed
+			} else {
+				todobuffer = append(todobuffer, nc.Node)
 			}
 
-	for ks := range in {
-		select {
-		case rateLimit <- struct{}{}:
 		case <-ctx.Done():
-			return
+			return ctx.Err()
 		}
-		wg.Add(1)
-		go get(ks)
 	}
+
 }

merkledag_test.go

@@ -504,3 +504,46 @@ func TestCidRawDoesnNeedData(t *testing.T) {
 		t.Fatal("raw node shouldn't have any links")
 	}
 }
+
+func TestEnumerateAsyncFailsNotFound(t *testing.T) {
+	a := NodeWithData([]byte("foo1"))
+	b := NodeWithData([]byte("foo2"))
+	c := NodeWithData([]byte("foo3"))
+	d := NodeWithData([]byte("foo4"))
+
+	ds := dstest.Mock()
+	for _, n := range []node.Node{a, b, c} {
+		_, err := ds.Add(n)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	parent := new(ProtoNode)
+	if err := parent.AddNodeLinkClean("a", a); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := parent.AddNodeLinkClean("b", b); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := parent.AddNodeLinkClean("c", c); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := parent.AddNodeLinkClean("d", d); err != nil {
+		t.Fatal(err)
+	}
+
+	pcid, err := ds.Add(parent)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cset := cid.NewSet()
+	err = EnumerateChildrenAsync(context.Background(), ds, pcid, cset.Visit)
+	if err == nil {
+		t.Fatal("this should have failed")
+	}
+}