From b4b26faefdd0163c19ee86d7d61ec0103a1b24b5 Mon Sep 17 00:00:00 2001
From: Jeromy <jeromyj@gmail.com>
Date: Fri, 31 Oct 2014 18:48:48 +0000
Subject: [PATCH] make FetchGraph waitable

---
 merkledag/merkledag.go | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/merkledag/merkledag.go b/merkledag/merkledag.go
index d9221f0da..26431dea2 100644
--- a/merkledag/merkledag.go
+++ b/merkledag/merkledag.go
@@ -2,6 +2,7 @@ package merkledag
 
 import (
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
@@ -242,9 +243,16 @@ func (n *dagService) Remove(nd *Node) error {
 	return n.Blocks.DeleteBlock(k)
 }
 
-func FetchGraph(ctx context.Context, root *Node, serv DAGService) {
+func FetchGraph(ctx context.Context, root *Node, serv DAGService) chan struct{} {
+	var wg sync.WaitGroup
+	done := make(chan struct{})
+
 	for _, l := range root.Links {
+		wg.Add(1)
 		go func(lnk *Link) {
+
+			// Signal child is done on way out
+			defer wg.Done()
 			select {
 			case <-ctx.Done():
 				return
@@ -255,7 +263,16 @@ func FetchGraph(ctx context.Context, root *Node, serv DAGService) {
 				log.Error(err)
 				return
 			}
-			FetchGraph(ctx, nd, serv)
+
+			// Wait for children to finish
+			<-FetchGraph(ctx, nd, serv)
 		}(l)
 	}
+
+	go func() {
+		wg.Wait()
+		done <- struct{}{}
+	}()
+
+	return done
 }