plumbing: add context to allow cancel on diff/patch computing

Signed-off-by: Marc Barussaud <marc.barussaud@orange.com>

1 files changed, 21 insertions, 0 deletions

diff --git a/utils/merkletrie/difftree.go b/utils/merkletrie/difftree.go
index 2294096..d57ed13 100644
--- a/utils/merkletrie/difftree.go
+++ b/utils/merkletrie/difftree.go
@@ -248,15 +248,30 @@ package merkletrie
 // h: else of i
 
 import (
+	"context"
+	"errors"
 	"fmt"
 
 	"gopkg.in/src-d/go-git.v4/utils/merkletrie/noder"
 )
 
+var (
+	ErrCanceled = errors.New("operation canceled")
+)
+
 // DiffTree calculates the list of changes between two merkletries.  It
 // uses the provided hashEqual callback to compare noders.
 func DiffTree(fromTree, toTree noder.Noder,
 	hashEqual noder.Equal) (Changes, error) {
+	return DiffTreeContext(context.Background(), fromTree, toTree, hashEqual)
+}
+
+// DiffTree calculates the list of changes between two merkletries.  It
+// uses the provided hashEqual callback to compare noders.
+// Error will be returned if context expires
+// Provided context must be non nil
+func DiffTreeContext(ctx context.Context, fromTree, toTree noder.Noder,
+	hashEqual noder.Equal) (Changes, error) {
 	ret := NewChanges()
 
 	ii, err := newDoubleIter(fromTree, toTree, hashEqual)
@@ -265,6 +280,12 @@ func DiffTree(fromTree, toTree noder.Noder,
 	}
 
 	for {
+		select {
+		case <-ctx.Done():
+			return nil, ErrCanceled
+		default:
+		}
+
 		from := ii.from.current
 		to := ii.to.current