Merge pull request #3809 from ipfs/faster-pin-ls

Fix #3783: Improve IsPinned() lookups for indirect pins

Author: whyrusleeping

pin/pin.go

@@ -276,8 +276,9 @@ func (p *pinner) isPinnedWithType(c *cid.Cid, mode PinMode) (string, bool, error
 	}
 
 	// Default is Indirect
+	visitedSet := cid.NewSet()
 	for _, rc := range p.recursePin.Keys() {
-		has, err := hasChild(p.dserv, rc, c)
+		has, err := hasChild(p.dserv, rc, c, visitedSet.Visit)
 		if err != nil {
 			return "", false, err
 		}
@@ -519,7 +520,9 @@ func (p *pinner) PinWithMode(c *cid.Cid, mode PinMode) {
 	}
 }
 
-func hasChild(ds mdag.LinkService, root *cid.Cid, child *cid.Cid) (bool, error) {
+// hasChild recursively looks for a Cid among the children of a root Cid.
+// The visit function can be used to shortcut already-visited branches.
+func hasChild(ds mdag.LinkService, root *cid.Cid, child *cid.Cid, visit func(*cid.Cid) bool) (bool, error) {
 	links, err := ds.GetLinks(context.Background(), root)
 	if err != nil {
 		return false, err
@@ -529,14 +532,15 @@ func hasChild(ds mdag.LinkService, root *cid.Cid, child *cid.Cid) (bool, error)
 		if lnk.Cid.Equals(child) {
 			return true, nil
 		}
+		if visit(c) {
+			has, err := hasChild(ds, c, child, visit)
+			if err != nil {
+				return false, err
+			}
 
-		has, err := hasChild(ds, c, child)
-		if err != nil {
-			return false, err
-		}
-
-		if has {
-			return has, nil
+			if has {
+				return has, nil
+			}
 		}
 	}
 	return false, nil

pin/pin_test.go

@@ -35,6 +35,17 @@ func assertPinned(t *testing.T, p Pinner, c *cid.Cid, failmsg string) {
 	}
 }
 
+func assertUnpinned(t *testing.T, p Pinner, c *cid.Cid, failmsg string) {
+	_, pinned, err := p.IsPinned(c)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if pinned {
+		t.Fatal(failmsg)
+	}
+}
+
 func TestPinnerBasic(t *testing.T) {
 	ctx := context.Background()
 
@@ -145,6 +156,122 @@ func TestPinnerBasic(t *testing.T) {
 	assertPinned(t, np, bk, "could not find recursively pinned node")
 }
 
+func TestIsPinnedLookup(t *testing.T) {
+	// We are going to test that lookups work in pins which share
+	// the same branches. For that we will construct this tree:
+	//
+	// A5->A4->A3->A2->A1->A0
+	//         /           /
+	// B-------           /
+	//  \                /
+	//   C---------------
+	//
+	// We will ensure that IsPinned works for all objects both when they
+	// are pinned and once they have been unpinned.
+	aBranchLen := 6
+	if aBranchLen < 3 {
+		t.Fatal("set aBranchLen to at least 3")
+	}
+
+	ctx := context.Background()
+	dstore := dssync.MutexWrap(ds.NewMapDatastore())
+	bstore := blockstore.NewBlockstore(dstore)
+	bserv := bs.New(bstore, offline.Exchange(bstore))
+
+	dserv := mdag.NewDAGService(bserv)
+
+	// TODO does pinner need to share datastore with blockservice?
+	p := NewPinner(dstore, dserv, dserv)
+
+	aNodes := make([]*mdag.ProtoNode, aBranchLen, aBranchLen)
+	aKeys := make([]*cid.Cid, aBranchLen, aBranchLen)
+	for i := 0; i < aBranchLen; i++ {
+		a, _ := randNode()
+		if i >= 1 {
+			err := a.AddNodeLink("child", aNodes[i-1])
+			if err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		ak, err := dserv.Add(a)
+		if err != nil {
+			t.Fatal(err)
+		}
+		//t.Logf("a[%d] is %s", i, ak)
+		aNodes[i] = a
+		aKeys[i] = ak
+	}
+
+	// Pin A5 recursively
+	if err := p.Pin(ctx, aNodes[aBranchLen-1], true); err != nil {
+		t.Fatal(err)
+	}
+
+	// Create node B and add A3 as child
+	b, _ := randNode()
+	if err := b.AddNodeLink("mychild", aNodes[3]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Create C node
+	c, _ := randNode()
+	// Add A0 as child of C
+	if err := c.AddNodeLink("child", aNodes[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Add C
+	ck, err := dserv.Add(c)
+	if err != nil {
+		t.Fatal(err)
+	}
+	//t.Logf("C is %s", ck)
+
+	// Add C to B and Add B
+	if err := b.AddNodeLink("myotherchild", c); err != nil {
+		t.Fatal(err)
+	}
+	bk, err := dserv.Add(b)
+	if err != nil {
+		t.Fatal(err)
+	}
+	//t.Logf("B is %s", bk)
+
+	// Pin C recursively
+
+	if err := p.Pin(ctx, c, true); err != nil {
+		t.Fatal(err)
+	}
+
+	// Pin B recursively
+
+	if err := p.Pin(ctx, b, true); err != nil {
+		t.Fatal(err)
+	}
+
+	assertPinned(t, p, aKeys[0], "A0 should be pinned")
+	assertPinned(t, p, aKeys[1], "A1 should be pinned")
+	assertPinned(t, p, ck, "C should be pinned")
+	assertPinned(t, p, bk, "B should be pinned")
+
+	// Unpin A5 recursively
+	if err := p.Unpin(ctx, aKeys[5], true); err != nil {
+		t.Fatal(err)
+	}
+
+	assertPinned(t, p, aKeys[0], "A0 should still be pinned through B")
+	assertUnpinned(t, p, aKeys[4], "A4 should be unpinned")
+
+	// Unpin B recursively
+	if err := p.Unpin(ctx, bk, true); err != nil {
+		t.Fatal(err)
+	}
+	assertUnpinned(t, p, bk, "B should be unpinned")
+	assertUnpinned(t, p, aKeys[1], "A1 should be unpinned")
+	assertPinned(t, p, aKeys[0], "A0 should still be pinned through C")
+}
+
 func TestDuplicateSemantics(t *testing.T) {
 	ctx := context.Background()
 	dstore := dssync.MutexWrap(ds.NewMapDatastore())