Cleanup, fix broken restart, and more tests.

License: MIT
Signed-off-by: Erik Ingenito <erik@carbonfive.com>

provider/provider.go

@@ -62,28 +62,13 @@ func (p *provider) handleAnnouncements() {
 				case <-p.ctx.Done():
 					return
 				case entry := <-p.queue.Dequeue():
-					if err := doProvide(p.ctx, p.contentRouting, entry.cid); err != nil {
+					log.Info("announce - start - ", entry.cid)
+					if err := p.contentRouting.Provide(p.ctx, entry.cid, true); err != nil {
 						log.Warningf("Unable to provide entry: %s, %s", entry.cid, err)
 					}
-
-					if err := entry.Complete(); err != nil {
-						log.Warningf("Unable to complete queue entry when providing: %s, %s", entry.cid, err)
-					}
+					log.Info("announce - end - ", entry.cid)
 				}
 			}
 		}()
 	}
 }
-
-// TODO: better document this provide logic
-func doProvide(ctx context.Context, contentRouting routing.ContentRouting, key cid.Cid) error {
-	// announce
-	log.Info("announce - start - ", key)
-	if err := contentRouting.Provide(ctx, key, true); err != nil {
-		log.Warningf("Failed to provide cid: %s", err)
-		// TODO: Maybe put these failures onto a failures queue?
-		return err
-	}
-	log.Info("announce - end - ", key)
-	return nil
-}

provider/provider_test.go

@@ -2,13 +2,15 @@ package provider
 
 import (
 	"context"
-	"github.com/ipfs/go-cid"
-	"github.com/ipfs/go-datastore"
-	"github.com/ipfs/go-ipfs-blocksutil"
-	pstore "github.com/libp2p/go-libp2p-peerstore"
 	"math/rand"
 	"testing"
 	"time"
+
+	blocksutil "github.com/ipfs/go-ipfs-blocksutil"
+	cid "github.com/ipfs/go-cid"
+	datastore "github.com/ipfs/go-datastore"
+	pstore "github.com/libp2p/go-libp2p-peerstore"
+	sync "github.com/ipfs/go-datastore/sync"
 )
 
 var blockGenerator = blocksutil.NewBlockGenerator()
@@ -25,11 +27,10 @@ func mockContentRouting() *mockRouting {
 
 func TestAnnouncement(t *testing.T) {
 	ctx := context.Background()
-	defer func() {
-		ctx.Done()
-	}()
+	defer ctx.Done()
 
-	queue, err := NewQueue(ctx, "test", datastore.NewMapDatastore())
+	ds := sync.MutexWrap(datastore.NewMapDatastore())
+	queue, err := NewQueue(ctx, "test", ds)
 	if err != nil {
 		t.Fatal(err)
 	}

provider/queue.go

@@ -3,14 +3,15 @@ package provider
 import (
 	"context"
 	"errors"
-	"github.com/ipfs/go-cid"
-	ds "github.com/ipfs/go-datastore"
-	"github.com/ipfs/go-datastore/namespace"
-	"github.com/ipfs/go-datastore/query"
 	"math"
 	"strconv"
 	"strings"
 	"sync"
+
+	cid "github.com/ipfs/go-cid"
+	datastore "github.com/ipfs/go-datastore"
+	namespace "github.com/ipfs/go-datastore/namespace"
+	query "github.com/ipfs/go-datastore/query"
 )
 
 // Entry allows for the durability in the queue. When a cid is dequeued it is
@@ -18,17 +19,10 @@ import (
 // receive.
 type Entry struct {
 	cid   cid.Cid
-	key   ds.Key
+	key   datastore.Key
 	queue *Queue
 }
 
-// Complete the entry by removing it from the queue
-func (e *Entry) Complete() error {
-	e.queue.lock.Lock()
-	defer e.queue.lock.Unlock()
-	return e.queue.remove(e.key)
-}
-
 // Queue provides a durable, FIFO interface to the datastore for storing cids
 //
 // Durability just means that cids in the process of being provided when a
@@ -44,41 +38,41 @@ type Queue struct {
 	tail uint64
 	head uint64
 
-	lock      sync.Mutex
-	datastore ds.Datastore
+	enqueueLock sync.Mutex
+	ds          datastore.Datastore // Must be threadsafe
 
 	dequeue  chan *Entry
 	added chan struct{}
 }
 
 // NewQueue creates a queue for cids
-func NewQueue(ctx context.Context, name string, datastore ds.Datastore) (*Queue, error) {
-	namespaced := namespace.Wrap(datastore, ds.NewKey("/"+name+"/queue/"))
+func NewQueue(ctx context.Context, name string, ds datastore.Datastore) (*Queue, error) {
+	namespaced := namespace.Wrap(ds, datastore.NewKey("/"+name+"/queue/"))
 	head, tail, err := getQueueHeadTail(ctx, name, namespaced)
 	if err != nil {
 		return nil, err
 	}
 	q := &Queue{
-		name:      name,
-		ctx:       ctx,
-		head:      head,
-		tail:      tail,
-		lock:      sync.Mutex{},
-		datastore: namespaced,
-		dequeue:   make(chan *Entry),
-		added:  make(chan struct{}),
+		name:        name,
+		ctx:         ctx,
+		head:        head,
+		tail:        tail,
+		enqueueLock: sync.Mutex{},
+		ds:          namespaced,
+		dequeue:     make(chan *Entry),
+		added:       make(chan struct{}),
 	}
 	return q, nil
 }
 
 // Enqueue puts a cid in the queue
 func (q *Queue) Enqueue(cid cid.Cid) error {
-	q.lock.Lock()
-	defer q.lock.Unlock()
+	q.enqueueLock.Lock()
+	defer q.enqueueLock.Unlock()
 
 	nextKey := q.queueKey(q.tail)
 
-	if err := q.datastore.Put(nextKey, cid.Bytes()); err != nil {
+	if err := q.ds.Put(nextKey, cid.Bytes()); err != nil {
 		return err
 	}
 
@@ -126,8 +120,9 @@ func (q *Queue) Run() {
 			case <-q.ctx.Done():
 				return
 			case q.dequeue <- entry:
+				q.head++
+				err = q.ds.Delete(entry.key)
 			}
-
 		}
 	}()
 }
@@ -135,12 +130,7 @@ func (q *Queue) Run() {
 // Find the next item in the queue, crawl forward if an entry is not
 // found in the next spot.
 func (q *Queue) next() (*Entry, error) {
-	q.lock.Lock()
-	defer func() {
-		q.lock.Unlock()
-	}()
-
-	var nextKey ds.Key
+	var key datastore.Key
 	var value []byte
 	var err error
 	for {
@@ -152,9 +142,12 @@ func (q *Queue) next() (*Entry, error) {
 			return nil, nil
 		default:
 		}
-		nextKey = q.queueKey(q.head)
-		value, err = q.datastore.Get(nextKey)
-		if err == ds.ErrNotFound {
+		key = q.queueKey(q.head)
+
+		value, err = q.ds.Get(key)
+
+		value, err = q.ds.Get(key)
+		if err == datastore.ErrNotFound {
 			q.head++
 			continue
 		} else if err != nil {
@@ -171,21 +164,23 @@ func (q *Queue) next() (*Entry, error) {
 
 	entry := &Entry{
 		cid:   id,
-		key:   nextKey,
+		key:   key,
 		queue: q,
 	}
 
-	q.head++
+	if err != nil {
+		return nil, err
+	}
 
 	return entry, nil
 }
 
-func (q *Queue) queueKey(id uint64) ds.Key {
-	return ds.NewKey(strconv.FormatUint(id, 10))
+func (q *Queue) queueKey(id uint64) datastore.Key {
+	return datastore.NewKey(strconv.FormatUint(id, 10))
 }
 
 // crawl over the queue entries to find the head and tail
-func getQueueHeadTail(ctx context.Context, name string, datastore ds.Datastore) (uint64, uint64, error) {
+func getQueueHeadTail(ctx context.Context, name string, datastore datastore.Datastore) (uint64, uint64, error) {
 	q := query.Query{}
 	results, err := datastore.Query(q)
 	if err != nil {
@@ -223,7 +218,3 @@ func getQueueHeadTail(ctx context.Context, name string, datastore ds.Datastore)
 
 	return head, tail, nil
 }
-
-func (q *Queue) remove(key ds.Key) error {
-	return q.datastore.Delete(key)
-}

provider/queue_test.go

+package provider
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	cid "github.com/ipfs/go-cid"
+	datastore "github.com/ipfs/go-datastore"
+	sync "github.com/ipfs/go-datastore/sync"
+)
+
+func makeCids(n int) []cid.Cid {
+	cids := make([]cid.Cid, 0, 10)
+	for i := 0; i < 10; i++ {
+		c := blockGenerator.Next().Cid()
+		cids = append(cids, c)
+	}
+	return cids
+}
+
+func assertOrdered(cids []cid.Cid, q *Queue, t *testing.T) {
+	for _, c := range cids {
+		select {
+		case dequeued := <- q.dequeue:
+			if c != dequeued.cid {
+				t.Fatalf("Error in ordering of CIDs retrieved from queue. Expected: %s, got: %s", c, dequeued.cid)
+			}
+
+		case <-time.After(time.Second * 1):
+			t.Fatal("Timeout waiting for cids to be provided.")
+		}
+	}
+}
+
+func TestBasicOperation(t *testing.T) {
+	ctx := context.Background()
+	defer ctx.Done()
+
+	ds := sync.MutexWrap(datastore.NewMapDatastore())
+	queue, err := NewQueue(ctx, "test", ds)
+	if err != nil {
+		t.Fatal(err)
+	}
+	queue.Run()
+
+	cids := makeCids(10)
+
+	for _, c := range cids {
+		err = queue.Enqueue(c)
+		if err != nil {
+			t.Fatal("Failed to enqueue CID")
+		}
+	}
+
+	assertOrdered(cids, queue, t)
+}
+
+func TestInitialization(t *testing.T) {
+	ctx := context.Background()
+	defer ctx.Done()
+
+	ds := sync.MutexWrap(datastore.NewMapDatastore())
+	queue, err := NewQueue(ctx, "test", ds)
+	if err != nil {
+		t.Fatal(err)
+	}
+	queue.Run()
+
+	cids := makeCids(10)
+
+	for _, c := range cids {
+		err = queue.Enqueue(c)
+		if err != nil {
+			t.Fatal("Failed to enqueue CID")
+		}
+	}
+
+	assertOrdered(cids[:5], queue, t)
+
+	// make a new queue, same data
+	queue, err = NewQueue(ctx, "test", ds)
+	if err != nil {
+		t.Fatal(err)
+	}
+	queue.Run()
+
+	assertOrdered(cids[5:], queue, t)
+}