Merge pull request #147 from jim-slattery-rs/fix_queue_poll_memory_leak

fix queue.Poll() memory leak

Author: dustinhiatt-wf

queue/queue.go

@@ -76,6 +76,21 @@ func (w *waiters) put(sema *sema) {
 	*w = append(*w, sema)
 }
 
+func (w *waiters) remove(sema *sema) {
+	if len(*w) == 0 {
+		return
+	}
+	// build new slice, copy all except sema
+	ws := *w
+	newWs := make(waiters, 0, len(*w))
+	for i := range ws {
+		if ws[i] != sema {
+			newWs = append(newWs, ws[i])
+		}
+	}
+	*w = newWs
+}
+
 type items []interface{}
 
 func (items *items) get(number int64) []interface{} {
@@ -237,6 +252,10 @@ func (q *Queue) Poll(number int64, timeout time.Duration) ([]interface{}, error)
 			select {
 			case sema.ready <- true:
 				// we called this before Put() could
+				// Remove sema from waiters.
+				q.lock.Lock()
+				q.waiters.remove(sema)
+				q.lock.Unlock()
 			default:
 				// Put() got it already, we need to call Done() so Put() can move on
 				sema.response.Done()

queue/queue_test.go

@@ -125,6 +125,18 @@ func TestPoll(t *testing.T) {
 	assert.Equal(t, ErrTimeout, err)
 }
 
+func TestPollNoMemoryLeak(t *testing.T) {
+	q := New(0)
+
+	assert.Len(t, q.waiters, 0)
+
+	for i := 0; i < 10; i++ {
+		// Poll() should cleanup waiters after timeout
+		q.Poll(1, time.Nanosecond)
+		assert.Len(t, q.waiters, 0)
+	}
+}
+
 func TestAddEmptyPut(t *testing.T) {
 	q := New(10)
 
@@ -419,6 +431,68 @@ func TestTakeUntilOnDisposedQueue(t *testing.T) {
 	assert.IsType(t, ErrDisposed, err)
 }
 
+func TestWaiters(t *testing.T) {
+	s1, s2, s3, s4 := newSema(), newSema(), newSema(), newSema()
+
+	w := waiters{}
+	assert.Len(t, w, 0)
+
+	//
+	// test put()
+	w.put(s1)
+	assert.Equal(t, waiters{s1}, w)
+
+	w.put(s2)
+	w.put(s3)
+	w.put(s4)
+	assert.Equal(t, waiters{s1, s2, s3, s4}, w)
+
+	//
+	// test remove()
+	//
+	// remove from middle
+	w.remove(s2)
+	assert.Equal(t, waiters{s1, s3, s4}, w)
+
+	// remove non-existing element
+	w.remove(s2)
+	assert.Equal(t, waiters{s1, s3, s4}, w)
+
+	// remove from beginning
+	w.remove(s1)
+	assert.Equal(t, waiters{s3, s4}, w)
+
+	// remove from end
+	w.remove(s4)
+	assert.Equal(t, waiters{s3}, w)
+
+	// remove last element
+	w.remove(s3)
+	assert.Empty(t, w)
+
+	// remove non-existing element
+	w.remove(s3)
+	assert.Empty(t, w)
+
+	//
+	// test get()
+	//
+	// start with 3 elements in list
+	w.put(s1)
+	w.put(s2)
+	w.put(s3)
+	assert.Equal(t, waiters{s1, s2, s3}, w)
+
+	// get() returns each item in insertion order
+	assert.Equal(t, s1, w.get())
+	assert.Equal(t, s2, w.get())
+	w.put(s4) // interleave a put(), item should go to the end
+	assert.Equal(t, s3, w.get())
+	assert.Equal(t, s4, w.get())
+	assert.Empty(t, w)
+	assert.Nil(t, w.get())
+}
+
 func TestExecuteInParallel(t *testing.T) {
 	q := New(10)
 	for i := 0; i < 10; i++ {