Merge pull request #91 from tylertreat/queue_poll

Queue poll

Author: alexandercampbell-wk

queue/error.go

@@ -18,4 +18,11 @@ package queue
 
 import "errors"
 
-var disposedError = errors.New(`Queue has been disposed.`)
+var (
+	// ErrDisposed is returned when an operation is performed on a disposed
+	// queue.
+	ErrDisposed = errors.New(`queue: disposed`)
+
+	// ErrTimeout is returned when an applicable queue operation times out.
+	ErrTimeout = errors.New(`queue: poll timed out`)
+)

queue/priority_queue.go

@@ -108,7 +108,7 @@ func (pq *PriorityQueue) Put(items ...Item) error {
 	pq.lock.Lock()
 	if pq.disposed {
 		pq.lock.Unlock()
-		return disposedError
+		return ErrDisposed
 	}
 
 	for _, item := range items {
@@ -122,7 +122,7 @@ func (pq *PriorityQueue) Put(items ...Item) error {
 		}
 
 		sema.response.Add(1)
-		sema.wg.Done()
+		sema.ready <- true
 		sema.response.Wait()
 		if len(pq.items) == 0 {
 			break
@@ -145,22 +145,21 @@ func (pq *PriorityQueue) Get(number int) ([]Item, error) {
 
 	if pq.disposed {
 		pq.lock.Unlock()
-		return nil, disposedError
+		return nil, ErrDisposed
 	}
 
 	var items []Item
 
 	if len(pq.items) == 0 {
 		sema := newSema()
 		pq.waiters.put(sema)
-		sema.wg.Add(1)
 		pq.lock.Unlock()
 
-		sema.wg.Wait()
+		<-sema.ready
 		pq.disposeLock.Lock()
 		if pq.disposed {
 			pq.disposeLock.Unlock()
-			return nil, disposedError
+			return nil, ErrDisposed
 		}
 		pq.disposeLock.Unlock()
 
@@ -221,7 +220,7 @@ func (pq *PriorityQueue) Dispose() {
 	pq.disposed = true
 	for _, waiter := range pq.waiters {
 		waiter.response.Add(1)
-		waiter.wg.Done()
+		waiter.ready <- true
 	}
 
 	pq.items = nil

queue/priority_queue_test.go

@@ -193,18 +193,18 @@ func TestEmptyPriorityGetWithDispose(t *testing.T) {
 
 	wg.Wait()
 
-	assert.IsType(t, disposedError, err)
+	assert.IsType(t, ErrDisposed, err)
 }
 
 func TestPriorityGetPutDisposed(t *testing.T) {
 	q := NewPriorityQueue(1)
 	q.Dispose()
 
 	_, err := q.Get(1)
-	assert.IsType(t, disposedError, err)
+	assert.IsType(t, ErrDisposed, err)
 
 	err = q.Put(mockItem(1))
-	assert.IsType(t, disposedError, err)
+	assert.IsType(t, ErrDisposed, err)
 }
 
 func BenchmarkPriorityQueue(b *testing.B) {

queue/queue.go

@@ -55,6 +55,7 @@ import (
 	"runtime"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
 type waiters []*sema
@@ -119,13 +120,13 @@ func (items *items) getUntil(checker func(item interface{}) bool) []interface{}
 }
 
 type sema struct {
-	wg       *sync.WaitGroup
+	ready    chan bool
 	response *sync.WaitGroup
 }
 
 func newSema() *sema {
 	return &sema{
-		wg:       &sync.WaitGroup{},
+		ready:    make(chan bool, 1),
 		response: &sync.WaitGroup{},
 	}
 }
@@ -149,7 +150,7 @@ func (q *Queue) Put(items ...interface{}) error {
 
 	if q.disposed {
 		q.lock.Unlock()
-		return disposedError
+		return ErrDisposed
 	}
 
 	q.items = append(q.items, items...)
@@ -159,8 +160,12 @@ func (q *Queue) Put(items ...interface{}) error {
 			break
 		}
 		sema.response.Add(1)
-		sema.wg.Done()
-		sema.response.Wait()
+		select {
+		case sema.ready <- true:
+			sema.response.Wait()
+		default:
+			// This semaphore timed out.
+		}
 		if len(q.items) == 0 {
 			break
 		}
@@ -175,6 +180,15 @@ func (q *Queue) Put(items ...interface{}) error {
 // parameter.  If no items are in the queue, this method will pause
 // until items are added to the queue.
 func (q *Queue) Get(number int64) ([]interface{}, error) {
+	return q.Poll(number, 0)
+}
+
+// Poll retrieves items from the queue.  If there are some items in the queue,
+// Poll will return a number UP TO the number passed in as a parameter.  If no
+// items are in the queue, this method will pause until items are added to the
+// queue or the provided timeout is reached.  A non-positive timeout will block
+// until items are added.  If a timeout occurs, ErrTimeout is returned.
+func (q *Queue) Poll(number int64, timeout time.Duration) ([]interface{}, error) {
 	if number < 1 {
 		// thanks again go
 		return []interface{}{}, nil
@@ -184,25 +198,32 @@ func (q *Queue) Get(number int64) ([]interface{}, error) {
 
 	if q.disposed {
 		q.lock.Unlock()
-		return nil, disposedError
+		return nil, ErrDisposed
 	}
 
 	var items []interface{}
 
 	if len(q.items) == 0 {
 		sema := newSema()
 		q.waiters.put(sema)
-		sema.wg.Add(1)
 		q.lock.Unlock()
 
-		sema.wg.Wait()
-		// we are now inside the put's lock
-		if q.disposed {
-			return nil, disposedError
+		var timeoutC <-chan time.Time
+		if timeout > 0 {
+			timeoutC = time.After(timeout)
+		}
+		select {
+		case <-sema.ready:
+			// we are now inside the put's lock
+			if q.disposed {
+				return nil, ErrDisposed
+			}
+			items = q.items.get(number)
+			sema.response.Done()
+			return items, nil
+		case <-timeoutC:
+			return nil, ErrTimeout
 		}
-		items = q.items.get(number)
-		sema.response.Done()
-		return items, nil
 	}
 
 	items = q.items.get(number)
@@ -222,7 +243,7 @@ func (q *Queue) TakeUntil(checker func(item interface{}) bool) ([]interface{}, e
 
 	if q.disposed {
 		q.lock.Unlock()
-		return nil, disposedError
+		return nil, ErrDisposed
 	}
 
 	result := q.items.getUntil(checker)
@@ -264,7 +285,7 @@ func (q *Queue) Dispose() {
 	q.disposed = true
 	for _, waiter := range q.waiters {
 		waiter.response.Add(1)
-		waiter.wg.Done()
+		waiter.ready <- true
 	}
 
 	q.items = nil

queue/queue_test.go

@@ -20,6 +20,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -81,6 +82,44 @@ func TestGet(t *testing.T) {
 	assert.Equal(t, `2`, result[0])
 }
 
+func TestPoll(t *testing.T) {
+	q := New(10)
+
+	q.Put(`test`)
+	result, err := q.Poll(2, 0)
+	if !assert.Nil(t, err) {
+		return
+	}
+
+	assert.Len(t, result, 1)
+	assert.Equal(t, `test`, result[0])
+	assert.Equal(t, int64(0), q.Len())
+
+	q.Put(`1`)
+	q.Put(`2`)
+
+	result, err = q.Poll(1, time.Millisecond)
+	if !assert.Nil(t, err) {
+		return
+	}
+
+	assert.Len(t, result, 1)
+	assert.Equal(t, `1`, result[0])
+	assert.Equal(t, int64(1), q.Len())
+
+	result, err = q.Poll(2, time.Millisecond)
+	if !assert.Nil(t, err) {
+		return
+	}
+
+	assert.Equal(t, `2`, result[0])
+
+	before := time.Now()
+	_, err = q.Poll(1, 5*time.Millisecond)
+	assert.InDelta(t, 5, time.Since(before).Seconds()*1000, 2)
+	assert.Equal(t, ErrTimeout, err)
+}
+
 func TestAddEmptyPut(t *testing.T) {
 	q := New(10)
 
@@ -189,7 +228,7 @@ func TestEmptyGetWithDispose(t *testing.T) {
 
 	wg.Wait()
 
-	assert.IsType(t, disposedError, err)
+	assert.IsType(t, ErrDisposed, err)
 }
 
 func TestGetPutDisposed(t *testing.T) {
@@ -198,10 +237,10 @@ func TestGetPutDisposed(t *testing.T) {
 	q.Dispose()
 
 	_, err := q.Get(1)
-	assert.IsType(t, disposedError, err)
+	assert.IsType(t, ErrDisposed, err)
 
 	err = q.Put(`a`)
-	assert.IsType(t, disposedError, err)
+	assert.IsType(t, ErrDisposed, err)
 }
 
 func BenchmarkQueue(b *testing.B) {
@@ -289,7 +328,7 @@ func TestTakeUntilOnDisposedQueue(t *testing.T) {
 	})
 
 	assert.Nil(t, result)
-	assert.IsType(t, disposedError, err)
+	assert.IsType(t, ErrDisposed, err)
 }
 
 func TestExecuteInParallel(t *testing.T) {
@@ -360,6 +399,28 @@ func BenchmarkQueueGet(b *testing.B) {
 	}
 }
 
+func BenchmarkQueuePoll(b *testing.B) {
+	numItems := int64(1000)
+
+	qs := make([]*Queue, 0, b.N)
+
+	for i := 0; i < b.N; i++ {
+		q := New(numItems)
+		for j := int64(0); j < numItems; j++ {
+			q.Put(j)
+		}
+		qs = append(qs, q)
+	}
+
+	b.ResetTimer()
+
+	for _, q := range qs {
+		for j := int64(0); j < numItems; j++ {
+			q.Poll(1, time.Millisecond)
+		}
+	}
+}
+
 func BenchmarkExecuteInParallel(b *testing.B) {
 	numItems := int64(1000)
 

queue/ring.go

@@ -78,7 +78,7 @@ func (rb *RingBuffer) Put(item interface{}) error {
 L:
 	for {
 		if atomic.LoadUint64(&rb.disposed) == 1 {
-			return disposedError
+			return ErrDisposed
 		}
 
 		n = rb.nodes[pos&rb.mask]
@@ -118,7 +118,7 @@ func (rb *RingBuffer) Get() (interface{}, error) {
 L:
 	for {
 		if atomic.LoadUint64(&rb.disposed) == 1 {
-			return nil, disposedError
+			return nil, ErrDisposed
 		}
 
 		n = rb.nodes[pos&rb.mask]