rapide: fix termination logic

rapide/rapide.go

@@ -3,6 +3,7 @@ package rapide
 import (
 	"context"
 	"fmt"
+	"io"
 	"sync"
 
 	"github.com/ipfs/go-cid"
@@ -28,6 +29,7 @@ func (c *Client) Get(ctx context.Context, root cid.Cid, traversal ipsl.Traversal
 	out := make(chan blocks.BlockOrError)
 	d := &download{
 		out:    out,
+		ctx:    ctx,
 		cancel: cancel,
 		root: node{
 			state:     todo,
@@ -45,20 +47,39 @@ func (c *Client) Get(ctx context.Context, root cid.Cid, traversal ipsl.Traversal
 
 type download struct {
 	out       chan<- blocks.BlockOrError
+	ctx       context.Context
 	cancel    context.CancelFunc
 	root      node
 	closeOnce sync.Once
 }
 
 // err cuts out the download and make it return an error, this is intended for unrecoverable errors.
 func (d *download) err(err error) {
+	d.closeOnce.Do(func() {
+		select {
+		case d.out <- blocks.IsNot(err):
+		case <-d.ctx.Done():
+		}
+		d.cancel()
+		close(d.out)
+	})
+}
+
+func (d *download) finish() {
 	d.closeOnce.Do(func() {
 		d.cancel()
-		d.out <- blocks.IsNot(err)
 		close(d.out)
 	})
 }
 
+func (d *download) workerFinished() {
+	d.root.mu.Lock()
+	defer d.root.mu.Unlock()
+	if d.root.state == done && len(d.root.childrens) == 0 {
+		d.finish() // file was downloaded !
+	}
+}
+
 type node struct {
 	// parent is not protected by the mutex and is readonly after creation
 	parent *node
@@ -99,6 +120,12 @@ func (n *node) expand(d *download, b blocks.Block) error {
 	n.childrens = childrens
 
 	for node, parent := n, n.parent; len(node.childrens) == 0; node, parent = parent, parent.parent {
+		if parent == nil {
+			// finished!
+			d.finish()
+			return io.EOF
+		}
+
 		// nothing to do, backtrack
 		parent.mu.Lock()
 		for i, v := range parent.childrens {

rapide/serverdriven.go

@@ -32,7 +32,8 @@ func (d *download) startServerDrivenWorker(ctx context.Context, impl ServerDrive
 }
 
 func (w *serverDrivenWorker) work(ctx context.Context) {
-	defer w.resetCurrentNodesWorkState()
+	defer w.download.workerFinished()
+	defer w.resetCurrentChildsNodeWorkState()
 
 workLoop:
 	for {
@@ -59,7 +60,7 @@ workLoop:
 		for {
 			if len(tasks) == 0 {
 				cancelCurrentRequest()
-				w.resetCurrentNodesWorkState()
+				w.resetCurrentChildsNodeWorkState()
 				continue workLoop
 			}
 			b, err := stream.Next()
@@ -70,7 +71,7 @@ workLoop:
 					return // request canceled
 				case err == io.EOF:
 					cancelCurrentRequest()
-					w.resetCurrentNodesWorkState()
+					w.resetCurrentChildsNodeWorkState()
 					continue workLoop
 				default:
 					// FIXME: support ignoring erroring parts of the tree when searching (dontGoThere)
@@ -99,13 +100,13 @@ workLoop:
 				task.mu.Unlock()
 				// we finished all parts of our tree, cancel current work and restart a new request.
 				cancelCurrentRequest()
-				w.resetCurrentNodesWorkState()
+				w.resetCurrentChildsNodeWorkState()
 				continue workLoop
 			}
 			if err := task.expand(w.download, b); err != nil {
 				task.mu.Unlock()
 				cancelCurrentRequest()
-				w.resetCurrentNodesWorkState()
+				w.resetCurrentChildsNodeWorkState()
 				return
 			}
 
@@ -210,16 +211,20 @@ func (w *serverDrivenWorker) findWork() (cid.Cid, ipsl.Traversal, bool) {
 	}
 }
 
-// resetCurrentNodesWorkState updates the state of the current node to longer count towards it.
-func (w *serverDrivenWorker) resetCurrentNodesWorkState() {
+// resetCurrentChildsNodeWorkState updates the state of the current node to longer count towards it.
+func (w *serverDrivenWorker) resetCurrentChildsNodeWorkState() {
 	c := w.current
 	if c == nil {
 		return // nothing to do
 	}
-	w.current = nil
 
 	// recursively walk the state and remove ourself from counters
-	w.recurseCancelNode(c)
+	// This is pretty contensius but that should be fine because server driven downloads should be cancel rarely, also most of thoses are gonna go on the fast path anyway.	c.mu.Lock()
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	for _, child := range c.childrens {
+		w.recurseCancelNode(child)
+	}
 }
 
 func (w *serverDrivenWorker) recurseCancelNode(c *node) {
@@ -229,11 +234,11 @@ func (w *serverDrivenWorker) recurseCancelNode(c *node) {
 
 	// This is pretty contensius but that should be fine because server driven downloads should be cancel rarely, also most of thoses are gonna go on the fast path anyway.	c.mu.Lock()
 	c.mu.Lock()
+	defer c.mu.Unlock()
 	c.workers -= 1
 	for _, child := range c.childrens {
 		w.recurseCancelNode(child)
 	}
-	defer c.mu.Unlock()
 }
 
 func (w *serverDrivenWorker) isOurTask(c *node) bool {