lib/model, lib/weakhash: Abort pulling quicker on folder stop (ref #5028)

lib/weakhash/benchmark_test.go

@@ -7,6 +7,7 @@
 package weakhash
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -28,7 +29,7 @@ func BenchmarkFind1MFile(b *testing.B) {
 		if err != nil {
 			b.Fatal(err)
 		}
-		_, err = Find(fd, []uint32{0, 1, 2}, size)
+		_, err = Find(context.Background(), fd, []uint32{0, 1, 2}, size)
 		if err != nil {
 			b.Fatal(err)
 		}

lib/weakhash/weakhash.go

@@ -8,6 +8,7 @@ package weakhash
 
 import (
 	"bufio"
+	"context"
 	"io"
 
 	"github.com/chmduquesne/rollinghash/adler32"
@@ -23,7 +24,7 @@ const (
 // Find finds all the blocks of the given size within io.Reader that matches
 // the hashes provided, and returns a hash -> slice of offsets within reader
 // map, that produces the same weak hash.
-func Find(ir io.Reader, hashesToFind []uint32, size int) (map[uint32][]int64, error) {
+func Find(ctx context.Context, ir io.Reader, hashesToFind []uint32, size int) (map[uint32][]int64, error) {
 	if ir == nil || len(hashesToFind) == 0 {
 		return nil, nil
 	}
@@ -50,6 +51,12 @@ func Find(ir io.Reader, hashesToFind []uint32, size int) (map[uint32][]int64, er
 	var i int64
 	var hash uint32
 	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+		}
+
 		hash = hf.Sum32()
 		if existing, ok := offsets[hash]; ok && len(existing) < maxWeakhashFinderHits {
 			offsets[hash] = append(existing, i)
@@ -67,8 +74,8 @@ func Find(ir io.Reader, hashesToFind []uint32, size int) (map[uint32][]int64, er
 	return offsets, nil
 }
 
-func NewFinder(ir io.ReadSeeker, size int, hashesToFind []uint32) (*Finder, error) {
-	offsets, err := Find(ir, hashesToFind, size)
+func NewFinder(ctx context.Context, ir io.ReadSeeker, size int, hashesToFind []uint32) (*Finder, error) {
+	offsets, err := Find(ctx, ir, hashesToFind, size)
 	if err != nil {
 		return nil, err
 	}

lib/weakhash/weakhash_test.go

@@ -11,6 +11,7 @@ package weakhash
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"io/ioutil"
 	"os"
@@ -36,7 +37,7 @@ func TestFinder(t *testing.T) {
 	}
 
 	hashes := []uint32{65143183, 65798547}
-	finder, err := NewFinder(f, 4, hashes)
+	finder, err := NewFinder(context.Background(), f, 4, hashes)
 	if err != nil {
 		t.Error(err)
 	}