Implementa

Author: dustinhiatt-wf

queue/priority_queue_test.go

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

queue/queue_test.go

@@ -189,7 +189,7 @@ func TestEmptyGetWithDispose(t *testing.T) {
 
 	wg.Wait()
 
-	assert.IsType(t, DisposedError{}, err)
+	assert.IsType(t, disposedError, err)
 }
 
 func TestGetPutDisposed(t *testing.T) {
@@ -198,10 +198,10 @@ func TestGetPutDisposed(t *testing.T) {
 	q.Dispose()
 
 	_, err := q.Get(1)
-	assert.IsType(t, DisposedError{}, err)
+	assert.IsType(t, disposedError, err)
 
 	err = q.Put(`a`)
-	assert.IsType(t, DisposedError{}, err)
+	assert.IsType(t, disposedError, err)
 }
 
 func BenchmarkQueue(b *testing.B) {
@@ -289,7 +289,7 @@ func TestTakeUntilOnDisposedQueue(t *testing.T) {
 	})
 
 	assert.Nil(t, result)
-	assert.IsType(t, DisposedError{}, err)
+	assert.IsType(t, disposedError, err)
 }
 
 func TestExecuteInParallel(t *testing.T) {

queue/ring.go

@@ -15,16 +15,128 @@ limitations under the License.
 */
 package queue
 
+import (
+	"errors"
+	"log"
+	"runtime"
+	"sync/atomic"
+)
+
+func init() {
+	log.Printf(`I HATE THIS.`)
+}
+
+// roundUp takes a uint64 greater than 0 and rounds it up to the next
+// power of 2.
+func roundUp(v uint64) uint64 {
+	v--
+	v |= v >> 1
+	v |= v >> 2
+	v |= v >> 4
+	v |= v >> 8
+	v |= v >> 16
+	v |= v >> 32
+	v++
+	return v
+}
+
+type node struct {
+	position uint64
+	data     interface{}
+}
+
+type nodes []*node
+
 type RingBuffer struct {
-	ifs               []interface{}
-	placed, retrieved uint64
-	size              int64
+	nodes                          nodes
+	queue, dequeue, mask, disposed uint64
+}
+
+func (rb *RingBuffer) init(size uint64) {
+	size = roundUp(size)
+	rb.nodes = make(nodes, size)
+	for i := uint64(0); i < size; i++ {
+		rb.nodes[i] = &node{position: i}
+	}
+	rb.mask = size - 1 // so we don't have to do this with every put/get operation
 }
 
-func (rb *RingBuffer) Put(items ...interface{}) error {
-	if len(items) == 0 {
-		return nil
+func (rb *RingBuffer) Put(item interface{}) error {
+	var n *node
+	pos := atomic.LoadUint64(&rb.queue)
+L:
+	for {
+		if atomic.LoadUint64(&rb.disposed) == 1 {
+			return disposedError
+		}
+
+		n = rb.nodes[pos&rb.mask]
+		seq := atomic.LoadUint64(&n.position)
+		switch dif := seq - pos; {
+		case dif == 0:
+			if atomic.CompareAndSwapUint64(&rb.queue, pos, pos+1) {
+				break L
+			}
+		case dif < 0:
+			return errors.New(`Full.`)
+		default:
+			pos = atomic.LoadUint64(&rb.queue)
+		}
+		runtime.Gosched() // free up the cpu before the next iteration
 	}
 
+	n.data = item
+	atomic.StoreUint64(&n.position, pos+1)
 	return nil
 }
+
+func (rb *RingBuffer) Get() (interface{}, error) {
+	var n *node
+	pos := atomic.LoadUint64(&rb.dequeue)
+L:
+	for {
+		if atomic.LoadUint64(&rb.disposed) == 1 {
+			return nil, disposedError
+		}
+
+		n = rb.nodes[pos&rb.mask]
+		seq := atomic.LoadUint64(&n.position)
+		switch dif := seq - (pos + 1); {
+		case dif == 0:
+			if atomic.CompareAndSwapUint64(&rb.dequeue, pos, pos+1) {
+				break L
+			}
+		case dif < 0:
+			return nil, errors.New(`Queue empty.`)
+		default:
+			pos = atomic.LoadUint64(&rb.dequeue)
+		}
+		runtime.Gosched() // free up cpu before next iteration
+	}
+	data := n.data
+	n.data = nil
+	atomic.StoreUint64(&n.position, pos+rb.mask+1)
+	return data, nil
+}
+
+func (rb *RingBuffer) Len() uint64 {
+	return atomic.LoadUint64(&rb.queue) - atomic.LoadUint64(&rb.dequeue)
+}
+
+func (rb *RingBuffer) Cap() uint64 {
+	return uint64(len(rb.nodes))
+}
+
+func (rb *RingBuffer) Dispose() {
+	atomic.CompareAndSwapUint64(&rb.disposed, 0, 1)
+}
+
+func (rb *RingBuffer) IsDisposed() bool {
+	return atomic.LoadUint64(&rb.disposed) == 1
+}
+
+func NewRingBuffer(size uint64) *RingBuffer {
+	rb := &RingBuffer{}
+	rb.init(size)
+	return rb
+}