[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
-// worker goroutines have drained the existing items in the queue, they will be
-// instructed to exit.
+// ShutDown will cause q to ignore all new items added to it and
+// immediately instruct the worker goroutines 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