pin: Future-proof against refcount marshaled size changes

License: MIT
Signed-off-by: Jeromy <jeromyj@gmail.com>

pin/set.go

@@ -44,14 +44,29 @@ type itemIterator func() (k key.Key, data []byte, ok bool)
 
 type keyObserver func(key.Key)
 
+// refcount is the marshaled format of refcounts. It may change
+// between versions; this is valid for version 1. Changing it may
+// become desirable if there are many links with refcount > 255.
+//
+// There are two guarantees that need to be preserved, if this is
+// changed:
+//
+//     - the marshaled format is of fixed size, matching
+//       unsafe.Sizeof(refcount(0))
+//     - methods of refcount handle endianness, and may
+//       in later versions need encoding/binary.
 type refcount uint8
 
 func (r refcount) Bytes() []byte {
-	// refcount size can change in later versions; this may need
-	// encoding/binary
 	return []byte{byte(r)}
 }
 
+// readRefcount returns the idx'th refcount in []byte, which is
+// assumed to be a sequence of refcount.Bytes results.
+func (r *refcount) ReadFromIdx(buf []byte, idx int) {
+	*r = refcount(buf[idx])
+}
+
 type sortByHash struct {
 	links []*merkledag.Link
 	data  []byte
@@ -70,9 +85,9 @@ func (s sortByHash) Swap(a, b int) {
 	if len(s.data) != 0 {
 		const n = int(unsafe.Sizeof(refcount(0)))
 		tmp := make([]byte, n)
-		copy(tmp, s.data[a:a+n])
-		copy(s.data[a:a+n], s.data[b:b+n])
-		copy(s.data[b:b+n], tmp)
+		copy(tmp, s.data[a*n:a*n+n])
+		copy(s.data[a*n:a*n+n], s.data[b*n:b*n+n])
+		copy(s.data[b*n:b*n+n], tmp)
 	}
 }
 
@@ -267,7 +282,9 @@ func loadMultiset(ctx context.Context, dag merkledag.DAGService, root *merkledag
 
 	refcounts := make(map[key.Key]uint64)
 	walk := func(buf []byte, idx int, link *merkledag.Link) error {
-		refcounts[key.Key(link.Hash)] += uint64(buf[idx])
+		var r refcount
+		r.ReadFromIdx(buf, idx)
+		refcounts[key.Key(link.Hash)] += uint64(r)
 		return nil
 	}
 	if err := walkItems(ctx, dag, n, walk, internalKeys); err != nil {

pin/set_test.go

@@ -0,0 +1,85 @@
+package pin
+
+import (
+	"testing"
+	"testing/quick"
+
+	"github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore"
+	dssync "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore/sync"
+	"github.com/ipfs/go-ipfs/blocks/blockstore"
+	"github.com/ipfs/go-ipfs/blocks/key"
+	"github.com/ipfs/go-ipfs/blockservice"
+	"github.com/ipfs/go-ipfs/exchange/offline"
+	"github.com/ipfs/go-ipfs/merkledag"
+	"golang.org/x/net/context"
+)
+
+func ignoreKeys(key.Key) {}
+
+func copyMap(m map[key.Key]uint16) map[key.Key]uint64 {
+	c := make(map[key.Key]uint64, len(m))
+	for k, v := range m {
+		c[k] = uint64(v)
+	}
+	return c
+}
+
+func TestMultisetRoundtrip(t *testing.T) {
+	dstore := dssync.MutexWrap(datastore.NewMapDatastore())
+	bstore := blockstore.NewBlockstore(dstore)
+	bserv, err := blockservice.New(bstore, offline.Exchange(bstore))
+	if err != nil {
+		t.Fatal(err)
+	}
+	dag := merkledag.NewDAGService(bserv)
+
+	fn := func(m map[key.Key]uint16) bool {
+		// Generate a smaller range for refcounts than full uint64, as
+		// otherwise this just becomes overly cpu heavy, splitting it
+		// out into too many items. That means we need to convert to
+		// the right kind of map. As storeMultiset mutates the map as
+		// part of its bookkeeping, this is actually good.
+		refcounts := copyMap(m)
+
+		ctx := context.Background()
+		n, err := storeMultiset(ctx, dag, refcounts, ignoreKeys)
+		if err != nil {
+			t.Fatalf("storing multiset: %v", err)
+		}
+		root := &merkledag.Node{}
+		const linkName = "dummylink"
+		if err := root.AddNodeLink(linkName, n); err != nil {
+			t.Fatalf("adding link to root node: %v", err)
+		}
+
+		roundtrip, err := loadMultiset(ctx, dag, root, linkName, ignoreKeys)
+		if err != nil {
+			t.Fatalf("loading multiset: %v", err)
+		}
+
+		orig := copyMap(m)
+		success := true
+		for k, want := range orig {
+			if got, ok := roundtrip[k]; ok {
+				if got != want {
+					success = false
+					t.Logf("refcount changed: %v -> %v for %q", want, got, k)
+				}
+				delete(orig, k)
+				delete(roundtrip, k)
+			}
+		}
+		for k, v := range orig {
+			success = false
+			t.Logf("refcount missing: %v for %q", v, k)
+		}
+		for k, v := range roundtrip {
+			success = false
+			t.Logf("refcount extra: %v for %q", v, k)
+		}
+		return success
+	}
+	if err := quick.Check(fn, nil); err != nil {
+		t.Fatal(err)
+	}
+}