add GetN method to retrieve batches of available items

dsqueue.go

@@ -44,6 +44,7 @@ type DSQueue struct {
 	enqueue      chan string
 	clear        chan chan<- int
 	closeTimeout time.Duration
+	getn         chan getRequest
 	name         string
 }
 
@@ -61,6 +62,7 @@ func New(ds datastore.Batching, name string, options ...Option) *DSQueue {
 		enqueue:      make(chan string),
 		clear:        make(chan chan<- int),
 		closeTimeout: cfg.closeTimeout,
+		getn:         make(chan getRequest),
 		name:         name,
 	}
 
@@ -107,6 +109,36 @@ func (q *DSQueue) Put(item []byte) (err error) {
 	return
 }
 
+type getRequest struct {
+	n   int
+	rsp chan getResponse
+}
+
+type getResponse struct {
+	items [][]byte
+	err   error
+}
+
+// GetN retrieves up to n items that are currently available in the queue. If
+// there are no items currently available, then none are returned and GetN does
+// not wait for any.
+//
+// GetN is used to poll the DSQueue for items and return batches of those
+// items. This is the most efficient way of fetching currently available items.
+//
+// GetN and Out can both be used to read items from the DSQueue, but they
+// should not be used concurrently as items will be returned by one or the
+// other indeterminately.
+func (q *DSQueue) GetN(n int) ([][]byte, error) {
+	rsp := make(chan getResponse)
+	q.getn <- getRequest{
+		n:   n,
+		rsp: rsp,
+	}
+	getRsp := <-rsp
+	return getRsp.items, getRsp.err
+}
+
 // Out returns a channel that for reading entries from the queue,
 func (q *DSQueue) Out() <-chan []byte {
 	return q.dequeue
@@ -264,9 +296,43 @@ func (q *DSQueue) worker(ctx context.Context, bufferSize, dedupCacheSize int, id
 			if bufferSize != 0 && inBuf.Len() >= bufferSize {
 				commit = true
 			}
+		case getRequest := <-q.getn:
+			n := getRequest.n
+			rspChan := getRequest.rsp
+			var outItems [][]byte
+
+			if item != "" {
+				outItems = append(outItems, []byte(item))
+
+				if !dsEmpty {
+					outItems, err = q.readDatastore(ctx, n-len(outItems), outItems)
+					if err != nil {
+						rspChan <- getResponse{
+							err: err,
+						}
+						continue
+					}
+				}
+
+				item = ""
+				idle = false
+			}
+			if len(outItems) < n {
+				for itm := range inBuf.IterPopFront() {
+					outItems = append(outItems, []byte(itm))
+					if len(outItems) == n {
+						break
+					}
+				}
+			}
+			rspChan <- getResponse{
+				items: outItems,
+			}
+
 		case dequeue <- []byte(item):
 			item = ""
 			idle = false
+
 		case <-batchTimer.C:
 			if idle {
 				if inBuf.Len() != 0 {
@@ -414,3 +480,67 @@ func (q *DSQueue) commitInput(ctx context.Context, counter uint64, items *deque.
 
 	return nil
 }
+
+// readDatastore reads at most n items from the data store queue, in order, and
+// appends them to items slice. Items are batch-deleted from the datastore as
+// they are read. The modified items slice is returned.
+func (q *DSQueue) readDatastore(ctx context.Context, n int, items [][]byte) ([][]byte, error) {
+	qry := query.Query{
+		KeysOnly: true,
+		Orders:   []query.Order{query.OrderByKey{}},
+		Limit:    n,
+	}
+	results, err := q.ds.Query(ctx, qry)
+	if err != nil {
+		return nil, err
+	}
+	defer results.Close()
+
+	batch, err := q.ds.Batch(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("cannot create datastore batch: %w", err)
+	}
+	var delCount int
+
+	for result := range results.Next() {
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
+		if result.Error != nil {
+			return nil, result.Error
+		}
+
+		if err = batch.Delete(ctx, datastore.NewKey(result.Key)); err != nil {
+			return nil, fmt.Errorf("error deleting queue item: %w", err)
+		}
+		delCount++
+
+		if delCount >= DefaultBufferSize {
+			delCount = 0
+			if err = batch.Commit(ctx); err != nil {
+				return nil, fmt.Errorf("cannot commit datastore updates: %w", err)
+			}
+		}
+
+		parts := strings.SplitN(strings.TrimPrefix(result.Key, "/"), "/", 2)
+		if len(parts) != 2 {
+			log.Errorw("malformed queued item, removing it from queue", "err", err, "key", result.Key, "qname", q.name)
+			continue
+		}
+		itemBin, err := base64.RawURLEncoding.DecodeString(parts[1])
+		if err != nil {
+			log.Errorw("error decoding queued item, removing it from queue", "err", err, "key", result.Key, "qname", q.name)
+			continue
+		}
+		items = append(items, itemBin)
+	}
+
+	if err = batch.Commit(ctx); err != nil {
+		return nil, fmt.Errorf("cannot commit datastore updated: %w", err)
+	}
+	if err = q.ds.Sync(ctx, datastore.NewKey("")); err != nil {
+		return nil, fmt.Errorf("cannot sync datastore: %w", err)
+	}
+
+	return items, nil
+}

dsqueue_test.go

@@ -18,7 +18,7 @@ import (
 
 const dsqName = "testq"
 
-func assertOrdered(cids []cid.Cid, q *dsqueue.DSQueue, t *testing.T) {
+func assertOrdered(t *testing.T, cids []cid.Cid, q *dsqueue.DSQueue) {
 	t.Helper()
 
 	var count int
@@ -60,41 +60,33 @@ func TestBasicOperation(t *testing.T) {
 	case <-time.After(time.Millisecond):
 	}
 
+	items := []string{"apple", "banana", "cherry"}
+
 	out := make(chan []string)
 	go func() {
 		var outStrs []string
-		for {
-			select {
-			case dq, open := <-queue.Out():
-				if !open {
-					out <- outStrs
-					return
-				}
-				dqItem := string(dq)
-				t.Log("got:", dqItem)
-				outStrs = append(outStrs, dqItem)
+		for dq := range queue.Out() {
+			dqItem := string(dq)
+			outStrs = append(outStrs, dqItem)
+			if len(outStrs) == len(items) {
+				break
 			}
 		}
+		out <- outStrs
 	}()
 
-	items := []string{"apple", "banana", "cherry"}
 	for _, item := range items {
 		queue.Put([]byte(item))
 	}
 
-	time.Sleep(time.Second)
-	err := queue.Close()
-	if err != nil {
-		t.Fatal(err)
-	}
-
 	qout := <-out
 
 	if len(qout) != len(items) {
-		t.Fatalf("dequeued wrond number of items, expected %d, got %d", len(items), len(qout))
+		t.Fatalf("dequeued wrong number of items, expected %d, got %d", len(items), len(qout))
 	}
 
-	if err = queue.Close(); err != nil {
+	err := queue.Close()
+	if err != nil {
 		t.Fatal(err)
 	}
 
@@ -104,6 +96,45 @@ func TestBasicOperation(t *testing.T) {
 	}
 }
 
+func TestGetN(t *testing.T) {
+	ds := sync.MutexWrap(datastore.NewMapDatastore())
+	queue := dsqueue.New(ds, dsqName, dsqueue.WithBufferSize(5), dsqueue.WithDedupCacheSize(0))
+	defer queue.Close()
+
+	cids := random.Cids(29)
+	for _, c := range cids {
+		queue.Put(c.Bytes())
+	}
+
+	outItems, err := queue.GetN(50)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(outItems) != len(cids) {
+		t.Fatalf("dequeued wrond number of items, expected %d, got %d", len(cids), len(outItems))
+	}
+
+	for i := range outItems {
+		outCid, err := cid.Parse(outItems[i])
+		if err != nil {
+			t.Fatal(err)
+		}
+		if outCid != cids[i] {
+			t.Fatal("retrieved items out of order")
+		}
+	}
+
+	outItems, err = queue.GetN(10)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(outItems) != 0 {
+		t.Fatal("shoul not get anymore items from queue")
+	}
+}
+
 func TestMangledData(t *testing.T) {
 	ds := sync.MutexWrap(datastore.NewMapDatastore())
 
@@ -125,8 +156,7 @@ func TestMangledData(t *testing.T) {
 	}
 
 	// expect to only see the valid cids we entered
-	expected := cids
-	assertOrdered(expected, queue, t)
+	assertOrdered(t, cids, queue)
 }
 
 func TestInitialization(t *testing.T) {
@@ -139,7 +169,7 @@ func TestInitialization(t *testing.T) {
 		queue.Put(c.Bytes())
 	}
 
-	assertOrdered(cids[:5], queue, t)
+	assertOrdered(t, cids[:5], queue)
 
 	err := queue.Close()
 	if err != nil {
@@ -150,7 +180,7 @@ func TestInitialization(t *testing.T) {
 	queue = dsqueue.New(ds, dsqName)
 	defer queue.Close()
 
-	assertOrdered(cids[5:], queue, t)
+	assertOrdered(t, cids[5:], queue)
 }
 
 func TestIdleFlush(t *testing.T) {
@@ -230,7 +260,7 @@ func TestPersistManyCids(t *testing.T) {
 	queue = dsqueue.New(ds, dsqName)
 	defer queue.Close()
 
-	assertOrdered(cids, queue, t)
+	assertOrdered(t, cids, queue)
 }
 
 func TestPersistOneCid(t *testing.T) {
@@ -250,7 +280,7 @@ func TestPersistOneCid(t *testing.T) {
 	queue = dsqueue.New(ds, dsqName)
 	defer queue.Close()
 
-	assertOrdered(cids, queue, t)
+	assertOrdered(t, cids, queue)
 }
 
 func TestDeduplicateCids(t *testing.T) {
@@ -268,7 +298,7 @@ func TestDeduplicateCids(t *testing.T) {
 	queue.Put(cids[0].Bytes())
 	queue.Put(cids[4].Bytes())
 
-	assertOrdered(cids, queue, t)
+	assertOrdered(t, cids, queue)
 
 	// Test with dedup cache disabled.
 	queue = dsqueue.New(ds, dsqName, dsqueue.WithDedupCacheSize(-1))
@@ -278,7 +308,7 @@ func TestDeduplicateCids(t *testing.T) {
 	for _, c := range cids {
 		queue.Put(c.Bytes())
 	}
-	assertOrdered(cids, queue, t)
+	assertOrdered(t, cids, queue)
 }
 
 func TestClear(t *testing.T) {