make FetchGraph waitable
Jeromy committed
Oct 31, 2014 at 18:48 UTC
b4b26faefdd0163c19ee86d7d61ec0103a1b24b5
1 file changed
+19
-2
merkledag/merkledag.go
+19
-2
@@ -2,6 +2,7 @@ package merkledag
2
3
import (
4
"fmt"
5
+ "sync"
6
"time"
7
8
"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 {
243
return n.Blocks.DeleteBlock(k)
244
}
245
245
-func FetchGraph(ctx context.Context, root *Node, serv DAGService) {
246
+func FetchGraph(ctx context.Context, root *Node, serv DAGService) chan struct{} {
247
+ var wg sync.WaitGroup
248
+ done := make(chan struct{})
249
+
250
for _, l := range root.Links {
251
+ wg.Add(1)
252
go func(lnk *Link) {
253
+
254
+ // Signal child is done on way out
255
+ defer wg.Done()
256
select {
257
case <-ctx.Done():
258
return
@@ -255,7 +263,16 @@ func FetchGraph(ctx context.Context, root *Node, serv DAGService) {
263
log.Error(err)
264
return
265
}
258
- FetchGraph(ctx, nd, serv)
266
+
267
+ // Wait for children to finish
268
+ <-FetchGraph(ctx, nd, serv)
269
}(l)
270
}
271
+
272
+ go func() {
273
+ wg.Wait()
274
+ done <- struct{}{}
275
+ }()
276
+
277
+ return done
278
}