[queue] Implement ShutDownWithDrain allowing the queue to drain when shutting down

Signed-off-by: Alexander Constantinescu <aconstan@redhat.com>

Kubernetes-commit: 5b740f430e0a4892e9db3a1fea9f349a06267755

util/workqueue/queue.go

@@ -29,6 +29,7 @@ type Interface interface {
 	Get() (item interface{}, shutdown bool)
 	Done(item interface{})
 	ShutDown()
+	ShutDownWithDrain()
 	ShuttingDown() bool
 }
 
@@ -86,6 +87,7 @@ type Type struct {
 	cond *sync.Cond
 
 	shuttingDown bool
+	drain        bool
 
 	metrics queueMetrics
 
@@ -110,6 +112,10 @@ func (s set) delete(item t) {
 	delete(s, item)
 }
 
+func (s set) len() int {
+	return len(s)
+}
+
 // Add marks item as needing processing.
 func (q *Type) Add(item interface{}) {
 	q.cond.L.Lock()
@@ -178,13 +184,71 @@ func (q *Type) Done(item interface{}) {
 	if q.dirty.has(item) {
 		q.queue = append(q.queue, item)
 		q.cond.Signal()
+	} else if q.processing.len() == 0 {
+		q.cond.Signal()
 	}
 }
 
-// ShutDown will cause q to ignore all new items added to it. As soon as the
+// ShutDown will cause q to ignore all new items added to it and
-// worker goroutines have drained the existing items in the queue, they will be
+// immediately instruct the worker goroutines to exit.
-// instructed to exit.
 func (q *Type) ShutDown() {
+	q.setDrain(false)
+	q.shutdown()
+}
+
+// ShutDownWithDrain will cause q to ignore all new items added to it. As soon
+// as the worker goroutines have "drained", i.e: finished processing and called
+// Done on all existing items in the queue; they will be instructed to exit and
+// ShutDownWithDrain will return. Hence: a strict requirement for using this is;
+// your workers must ensure that Done is called on all items in the queue once
+// the shut down has been initiated, if that is not the case: this will block
+// indefinitely. It is, however, safe to call ShutDown after having called
+// ShutDownWithDrain, as to force the queue shut down to terminate immediately
+// without waiting for the drainage.
+func (q *Type) ShutDownWithDrain() {
+	q.setDrain(true)
+	q.shutdown()
+	for q.isProcessing() && q.shouldDrain() {
+		q.waitForProcessing()
+	}
+}
+
+// isProcessing indicates if there are still items on the work queue being
+// processed. It's used to drain the work queue on an eventual shutdown.
+func (q *Type) isProcessing() bool {
+	q.cond.L.Lock()
+	defer q.cond.L.Unlock()
+	return q.processing.len() != 0
+}
+
+// waitForProcessing waits for the worker goroutines to finish processing items
+// and call Done on them.
+func (q *Type) waitForProcessing() {
+	q.cond.L.Lock()
+	defer q.cond.L.Unlock()
+	// Ensure that we do not wait on a queue which is already empty, as that
+	// could result in waiting for Done to be called on items in an empty queue
+	// which has already been shut down, which will result in waiting
+	// indefinitely.
+	if q.processing.len() == 0 {
+		return
+	}
+	q.cond.Wait()
+}
+
+func (q *Type) setDrain(shouldDrain bool) {
+	q.cond.L.Lock()
+	defer q.cond.L.Unlock()
+	q.drain = shouldDrain
+}
+
+func (q *Type) shouldDrain() bool {
+	q.cond.L.Lock()
+	defer q.cond.L.Unlock()
+	return q.drain
+}
+
+func (q *Type) shutdown() {
 	q.cond.L.Lock()
 	defer q.cond.L.Unlock()
 	q.shuttingDown = true

util/workqueue/queue_test.go

@@ -25,8 +25,21 @@ import (
 )
 
 func TestBasic(t *testing.T) {
+	tests := []struct {
+		queue         *workqueue.Type
+		queueShutDown func(workqueue.Interface)
+	}{
+		{
+			queue:         workqueue.New(),
+			queueShutDown: workqueue.Interface.ShutDown,
+		},
+		{
+			queue:         workqueue.New(),
+			queueShutDown: workqueue.Interface.ShutDownWithDrain,
+		},
+	}
+	for _, test := range tests {
 		// If something is seriously wrong this test will never complete.
-	q := workqueue.New()
 
 		// Start producers
 		const producers = 50
@@ -36,7 +49,7 @@ func TestBasic(t *testing.T) {
 			go func(i int) {
 				defer producerWG.Done()
 				for j := 0; j < 50; j++ {
-				q.Add(i)
+					test.queue.Add(i)
 					time.Sleep(time.Millisecond)
 				}
 			}(i)
@@ -50,7 +63,7 @@ func TestBasic(t *testing.T) {
 			go func(i int) {
 				defer consumerWG.Done()
 				for {
-				item, quit := q.Get()
+					item, quit := test.queue.Get()
 					if item == "added after shutdown!" {
 						t.Errorf("Got an item added after shutdown.")
 					}
@@ -60,19 +73,36 @@ func TestBasic(t *testing.T) {
 					t.Logf("Worker %v: begin processing %v", i, item)
 					time.Sleep(3 * time.Millisecond)
 					t.Logf("Worker %v: done processing %v", i, item)
-				q.Done(item)
+					test.queue.Done(item)
 				}
 			}(i)
 		}
 
 		producerWG.Wait()
-	q.ShutDown()
+		test.queueShutDown(test.queue)
-	q.Add("added after shutdown!")
+		test.queue.Add("added after shutdown!")
 		consumerWG.Wait()
+		if test.queue.Len() != 0 {
+			t.Errorf("Expected the queue to be empty, had: %v items", test.queue.Len())
+		}
+	}
 }
 
 func TestAddWhileProcessing(t *testing.T) {
-	q := workqueue.New()
+	tests := []struct {
+		queue         *workqueue.Type
+		queueShutDown func(workqueue.Interface)
+	}{
+		{
+			queue:         workqueue.New(),
+			queueShutDown: workqueue.Interface.ShutDown,
+		},
+		{
+			queue:         workqueue.New(),
+			queueShutDown: workqueue.Interface.ShutDownWithDrain,
+		},
+	}
+	for _, test := range tests {
 
 		// Start producers
 		const producers = 50
@@ -81,7 +111,7 @@ func TestAddWhileProcessing(t *testing.T) {
 		for i := 0; i < producers; i++ {
 			go func(i int) {
 				defer producerWG.Done()
-			q.Add(i)
+				test.queue.Add(i)
 			}(i)
 		}
 
@@ -96,22 +126,26 @@ func TestAddWhileProcessing(t *testing.T) {
 				// This tests the dirty-while-processing case.
 				counters := map[interface{}]int{}
 				for {
-				item, quit := q.Get()
+					item, quit := test.queue.Get()
 					if quit {
 						return
 					}
 					counters[item]++
 					if counters[item] < 2 {
-					q.Add(item)
+						test.queue.Add(item)
 					}
-				q.Done(item)
+					test.queue.Done(item)
 				}
 			}(i)
 		}
 
 		producerWG.Wait()
-	q.ShutDown()
+		test.queueShutDown(test.queue)
 		consumerWG.Wait()
+		if test.queue.Len() != 0 {
+			t.Errorf("Expected the queue to be empty, had: %v items", test.queue.Len())
+		}
+	}
 }
 
 func TestLen(t *testing.T) {
@@ -159,3 +193,131 @@ func TestReinsert(t *testing.T) {
 		t.Errorf("Expected queue to be empty. Has %v items", a)
 	}
 }
+
+func TestQueueDrainageUsingShutDownWithDrain(t *testing.T) {
+
+	q := workqueue.New()
+
+	q.Add("foo")
+	q.Add("bar")
+
+	firstItem, _ := q.Get()
+	secondItem, _ := q.Get()
+
+	finishedWG := sync.WaitGroup{}
+	finishedWG.Add(1)
+	go func() {
+		defer finishedWG.Done()
+		q.ShutDownWithDrain()
+	}()
+
+	// This is done as to simulate a sequence of events where ShutDownWithDrain
+	// is called before we start marking all items as done - thus simulating a
+	// drain where we wait for all items to finish processing.
+	shuttingDown := false
+	for !shuttingDown {
+		_, shuttingDown = q.Get()
+	}
+
+	// Mark the first two items as done, as to finish up
+	q.Done(firstItem)
+	q.Done(secondItem)
+
+	finishedWG.Wait()
+}
+
+func TestNoQueueDrainageUsingShutDown(t *testing.T) {
+
+	q := workqueue.New()
+
+	q.Add("foo")
+	q.Add("bar")
+
+	q.Get()
+	q.Get()
+
+	finishedWG := sync.WaitGroup{}
+	finishedWG.Add(1)
+	go func() {
+		defer finishedWG.Done()
+		// Invoke ShutDown: suspending the execution immediately.
+		q.ShutDown()
+	}()
+
+	// We can now do this and not have the test timeout because we didn't call
+	// Done on the first two items before arriving here.
+	finishedWG.Wait()
+}
+
+func TestForceQueueShutdownUsingShutDown(t *testing.T) {
+
+	q := workqueue.New()
+
+	q.Add("foo")
+	q.Add("bar")
+
+	q.Get()
+	q.Get()
+
+	finishedWG := sync.WaitGroup{}
+	finishedWG.Add(1)
+	go func() {
+		defer finishedWG.Done()
+		q.ShutDownWithDrain()
+	}()
+
+	// This is done as to simulate a sequence of events where ShutDownWithDrain
+	// is called before ShutDown
+	shuttingDown := false
+	for !shuttingDown {
+		_, shuttingDown = q.Get()
+	}
+
+	// Use ShutDown to force the queue to shut down (simulating a caller
+	// which can invoke this function on a second SIGTERM/SIGINT)
+	q.ShutDown()
+
+	// We can now do this and not have the test timeout because we didn't call
+	// done on any of the items before arriving here.
+	finishedWG.Wait()
+}
+
+func TestQueueDrainageUsingShutDownWithDrainWithDirtyItem(t *testing.T) {
+	q := workqueue.New()
+
+	q.Add("foo")
+	gotten, _ := q.Get()
+	q.Add("foo")
+
+	finishedWG := sync.WaitGroup{}
+	finishedWG.Add(1)
+	go func() {
+		defer finishedWG.Done()
+		q.ShutDownWithDrain()
+	}()
+
+	// Ensure that ShutDownWithDrain has started and is blocked.
+	shuttingDown := false
+	for !shuttingDown {
+		_, shuttingDown = q.Get()
+	}
+
+	// Finish "working".
+	q.Done(gotten)
+
+	// `shuttingDown` becomes false because Done caused an item to go back into
+	// the queue.
+	again, shuttingDown := q.Get()
+	if shuttingDown {
+		t.Fatalf("should not have been done")
+	}
+	q.Done(again)
+
+	// Now we are really done.
+	_, shuttingDown = q.Get()
+	if !shuttingDown {
+		t.Fatalf("should have been done")
+	}
+
+	finishedWG.Wait()
+}