reprovider: reduce pinned strategy i/o overhead

License: MIT
Signed-off-by: Łukasz Magiera <magik6k@gmail.com>

exchange/reprovide/providers.go

@@ -2,7 +2,6 @@ package reprovide
 
 import (
 	"context"
-	"errors"
 
 	blocks "github.com/ipfs/go-ipfs/blocks/blockstore"
 	merkledag "github.com/ipfs/go-ipfs/merkledag"
@@ -29,36 +28,72 @@ func NewPinnedProvider(pinning pin.Pinner, dag merkledag.DAGService, onlyRoots b
 		outCh := make(chan *cid.Cid)
 		go func() {
 			defer close(outCh)
-			set.ForEach(func(c *cid.Cid) error {
+			for c := range set.new {
 				select {
 				case <-ctx.Done():
-					return errors.New("context cancelled")
+					return
 				case outCh <- c:
 				}
-				return nil
-			})
+			}
+
 		}()
 
 		return outCh, nil
 	}
 }
 
-func pinSet(ctx context.Context, pinning pin.Pinner, dag merkledag.DAGService, onlyRoots bool) (*cid.Set, error) {
-	set := cid.NewSet()
-	for _, key := range pinning.DirectKeys() {
-		set.Add(key)
-	}
+func pinSet(ctx context.Context, pinning pin.Pinner, dag merkledag.DAGService, onlyRoots bool) (*streamingSet, error) {
+	set := newStreamingSet()
+
+	go func() {
+		for _, key := range pinning.DirectKeys() {
+			set.add(key)
+		}
 
-	for _, key := range pinning.RecursiveKeys() {
-		set.Add(key)
+		for _, key := range pinning.RecursiveKeys() {
+			set.add(key)
 
-		if !onlyRoots {
-			err := merkledag.EnumerateChildren(ctx, dag.GetLinks, key, set.Visit)
-			if err != nil {
-				return nil, err
+			if !onlyRoots {
+				err := merkledag.EnumerateChildren(ctx, dag.GetLinks, key, set.add)
+				if err != nil {
+					return //TODO: propagate to chan / log?
+				}
 			}
 		}
-	}
+
+		close(set.new)
+	}()
 
 	return set, nil
 }
+
+type streamingSet struct {
+	set map[string]struct{}
+	new chan *cid.Cid
+}
+
+// NewSet initializes and returns a new Set.
+func newStreamingSet() *streamingSet {
+	return &streamingSet{
+		set: make(map[string]struct{}),
+		new: make(chan *cid.Cid),
+	}
+}
+
+// has returns if the Set contains a given Cid.
+func (s *streamingSet) has(c *cid.Cid) bool {
+	_, ok := s.set[string(c.Bytes())]
+	return ok
+}
+
+// add adds a Cid to the set only if it is
+// not in it already.
+func (s *streamingSet) add(c *cid.Cid) bool {
+	if !s.has(c) {
+		s.set[string(c.Bytes())] = struct{}{}
+		s.new <- c
+		return true
+	}
+
+	return false
+}