Merge pull request #102 from ichinaski/master

Add support for duplicated items in PriorityQueue

Author: dustinhiatt-wf

queue/priority_queue.go

@@ -105,12 +105,13 @@ func (items *priorityItems) push(item Item) {
 // items that implement the Item interface and adds them
 // to the queue in priority order.
 type PriorityQueue struct {
-	waiters     waiters
-	items       priorityItems
-	itemMap     map[Item]struct{}
-	lock        sync.Mutex
-	disposeLock sync.Mutex
-	disposed    bool
+	waiters         waiters
+	items           priorityItems
+	itemMap         map[Item]struct{}
+	lock            sync.Mutex
+	disposeLock     sync.Mutex
+	disposed        bool
+	allowDuplicates bool
 }
 
 // Put adds items to the queue.
@@ -127,7 +128,9 @@ func (pq *PriorityQueue) Put(items ...Item) error {
 	}
 
 	for _, item := range items {
-		if _, ok := pq.itemMap[item]; !ok {
+		if pq.allowDuplicates {
+			pq.items.push(item)
+		} else if _, ok := pq.itemMap[item]; !ok {
 			pq.itemMap[item] = struct{}{}
 			pq.items.push(item)
 		}
@@ -186,7 +189,9 @@ func (pq *PriorityQueue) Get(number int) ([]Item, error) {
 		}
 
 		items = pq.items.get(number)
-		deleteItems(items)
+		if !pq.allowDuplicates {
+			deleteItems(items)
+		}
 		sema.response.Done()
 		return items, nil
 	}
@@ -252,9 +257,10 @@ func (pq *PriorityQueue) Dispose() {
 }
 
 // NewPriorityQueue is the constructor for a priority queue.
-func NewPriorityQueue(hint int) *PriorityQueue {
+func NewPriorityQueue(hint int, allowDuplicates bool) *PriorityQueue {
 	return &PriorityQueue{
-		items:   make(priorityItems, 0, hint),
-		itemMap: make(map[Item]struct{}, hint),
+		items:           make(priorityItems, 0, hint),
+		itemMap:         make(map[Item]struct{}, hint),
+		allowDuplicates: allowDuplicates,
 	}
 }

queue/priority_queue_test.go

@@ -24,7 +24,7 @@ import (
 )
 
 func TestPriorityPut(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 
 	q.Put(mockItem(2))
 
@@ -41,7 +41,7 @@ func TestPriorityPut(t *testing.T) {
 }
 
 func TestPriorityGet(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 
 	q.Put(mockItem(2))
 	result, err := q.Get(2)
@@ -84,15 +84,15 @@ func TestPriorityGet(t *testing.T) {
 }
 
 func TestAddEmptyPriorityPut(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 
 	q.Put()
 
 	assert.Len(t, q.items, 0)
 }
 
 func TestPriorityGetNonPositiveNumber(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 
 	q.Put(mockItem(1))
 
@@ -112,7 +112,7 @@ func TestPriorityGetNonPositiveNumber(t *testing.T) {
 }
 
 func TestPriorityEmpty(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 	assert.True(t, q.Empty())
 
 	q.Put(mockItem(1))
@@ -121,7 +121,7 @@ func TestPriorityEmpty(t *testing.T) {
 }
 
 func TestPriorityGetEmpty(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 
 	go func() {
 		q.Put(mockItem(1))
@@ -139,7 +139,7 @@ func TestPriorityGetEmpty(t *testing.T) {
 }
 
 func TestMultiplePriorityGetEmpty(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 	var wg sync.WaitGroup
 	wg.Add(2)
 	results := make([][]Item, 2)
@@ -175,7 +175,7 @@ func TestMultiplePriorityGetEmpty(t *testing.T) {
 }
 
 func TestEmptyPriorityGetWithDispose(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 	var wg sync.WaitGroup
 	wg.Add(1)
 
@@ -197,7 +197,7 @@ func TestEmptyPriorityGetWithDispose(t *testing.T) {
 }
 
 func TestPriorityGetPutDisposed(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 	q.Dispose()
 
 	_, err := q.Get(1)
@@ -208,7 +208,7 @@ func TestPriorityGetPutDisposed(t *testing.T) {
 }
 
 func BenchmarkPriorityQueue(b *testing.B) {
-	q := NewPriorityQueue(b.N)
+	q := NewPriorityQueue(b.N, false)
 	var wg sync.WaitGroup
 	wg.Add(1)
 	i := 0
@@ -232,16 +232,24 @@ func BenchmarkPriorityQueue(b *testing.B) {
 }
 
 func TestPriorityPeek(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 	q.Put(mockItem(1))
 
 	assert.Equal(t, mockItem(1), q.Peek())
 }
 
 func TestInsertDuplicate(t *testing.T) {
-	q := NewPriorityQueue(1)
+	q := NewPriorityQueue(1, false)
 	q.Put(mockItem(1))
 	q.Put(mockItem(1))
 
 	assert.Equal(t, 1, q.Len())
 }
+
+func TestAllowDuplicates(t *testing.T) {
+	q := NewPriorityQueue(2, true)
+	q.Put(mockItem(1))
+	q.Put(mockItem(1))
+
+	assert.Equal(t, 2, q.Len())
+}