Close waits on doneSignal. Test Close no-op when called multiple times.

Author: ax-jochoa

internal/queues/callback_queue.go

@@ -21,15 +21,17 @@ var ErrQueueClosed = errors.New("queue is closed")
 type CallBackQueue struct {
 	// The ordered list of callbacks to be processed.
 	list *lists.List[*callBackRequest]
-	// enqueueSignals signals when a new item is added. It must be a buffered
+	// enqueueSignals are sent when a new item is added. It must be a buffered
 	// channel to avoid missing signals.
 	enqueueSignals chan struct{}
 	// running prevents concurrent processCallBacks execution.
 	running sync.Mutex
 	// If false, the queue must not be used; return ErrQueueClosed.
 	opened atomic.Bool
-	// closeSignal is closed as a signal to shut down the queue processing.
+	// closeSignal is closed to signal the queue to begin shutdown.
 	closeSignal chan struct{}
+	// doneSignal is closed to signal the queue is fully shutdown.
+	doneSignal chan struct{}
 }
 
 // newUnopenedCallBackQueue creates a queue in the closed state. Use
@@ -39,6 +41,7 @@ func newUnopenedCallBackQueue() *CallBackQueue {
 		list:           lists.NewList[*callBackRequest](),
 		enqueueSignals: make(chan struct{}, 1),
 		closeSignal:    make(chan struct{}),
+		doneSignal:     make(chan struct{}),
 	}
 }
 
@@ -50,8 +53,6 @@ func OpenCallBackQueue() *CallBackQueue {
 	q.running.Lock()
 	processCallBacksIsRunning := make(chan struct{})
 	go func() {
-		defer q.running.Unlock()
-		defer q.Close()
 		defer func() {
 			if v := recover(); v != nil {
 				// This should not happen, but if it does, this defer only adds
@@ -60,6 +61,8 @@ func OpenCallBackQueue() *CallBackQueue {
 				panic(fmt.Sprintf("callback queue panicked: %v", v))
 			}
 		}()
+		defer q.running.Unlock()
+		defer close(q.doneSignal)
 		q.opened.Store(true)
 		close(processCallBacksIsRunning)
 		q.processCallBacks()
@@ -78,12 +81,24 @@ func (q *CallBackQueue) Close() {
 	if !q.opened.Swap(false) {
 		return // The queue is already closed.
 	}
-	close(q.closeSignal)
-	// Obtain the running lock to ensure the queue is finished processing before
-	// returning.
-	q.running.Lock()
-	defer q.running.Unlock()
 	q.list.Clear()
+	close(q.closeSignal)
+	<-q.doneSignal
+}
+
+// Push adds a callback to the queue. It panics if cb is nil. It returns
+// ErrQueueClosed if the queue is closed.
+func (q *CallBackQueue) Push(cb CallBackFunc) error {
+	if cb == nil {
+		panic("nil callback function")
+	}
+	if !q.opened.Load() {
+		return ErrQueueClosed
+	}
+	// Preserve order.
+	q.list.PushBack(&callBackRequest{fn: cb, tm: time.Now()})
+	q.signalEnqueue()
+	return nil
 }
 
 // processCallBacks is responsible for invoking callbacks from the list when it
@@ -146,21 +161,6 @@ func (q *CallBackQueue) invokeOneCallBack() {
 	curr.fn(callbackCtx, time.Since(curr.tm))
 }
 
-// Push adds a callback to the queue. It panics if cb is nil. It returns
-// ErrQueueClosed if the queue is closed.
-func (q *CallBackQueue) Push(cb CallBackFunc) error {
-	if cb == nil {
-		panic("nil callback function")
-	}
-	if !q.opened.Load() {
-		return ErrQueueClosed
-	}
-	// Preserve order.
-	q.list.PushBack(&callBackRequest{fn: cb, tm: time.Now()})
-	q.signalEnqueue()
-	return nil
-}
-
 // CallBackFunc is a function type that represents a callback to be executed.
 // The ctx is canceled if the queue is closed while the callback is executing.
 // The delay is the time since the callback was added to the queue.

internal/queues/callback_queue_test.go

@@ -38,9 +38,9 @@ func assertErrorIs(t *testing.T, err error, target error, msg string) {
 func TestCallbackQueue_newUnopenedCallBackQueue(t *testing.T) {
 	q := newUnopenedCallBackQueue()
 	assertTrue(t, q != nil, "newUnopenedCallBackQueue should return a non-nil queue")
-	assertTrue(t, !q.opened.Load(), "newUnopenedCallBackQueue should return a closed queue")
+	assertTrue(t, !q.opened.Load(), "newUnopenedCallBackQueue should return an opened queue")
 	assertTrue(t, q.list.Len() == 0, "newUnopenedCallBackQueue should return an empty queue")
-	// check that the enqueueSignals channel is buffered
+	// Check that the enqueueSignals channel is buffered.
 	select {
 	case q.enqueueSignals <- struct{}{}:
 	default:
@@ -52,16 +52,19 @@ func TestCallbackQueue_OpenCallBackQueue(t *testing.T) {
 	q := OpenCallBackQueue()
 	assertTrue(t, q != nil, "OpenCallBackQueue should return a non-nil queue")
 	assertTrue(t, q.opened.Load(), "OpenCallBackQueue should return an opened queue")
+	assertTrue(t, !q.running.TryLock(), "OpenCallBackQueue should not have the running lock locked")
 	q.Close()
 	assertTrue(t, !q.opened.Load(), "OpenCallBackQueue should close the queue after Close() is called")
-	// check that the queue was closed by acquiring the running lock.
+	// Check that doneSignal is closed.
+	<-q.doneSignal
+	// Check that the running lock was released.
 	q.running.Lock()
 	defer q.running.Unlock()
 }
 
 func TestCallbackQueue_processCallBacks_starts_and_cancels(t *testing.T) {
 	q := newUnopenedCallBackQueue()
-	// stage a callback to be processed
+	// Stage a callback to be processed
 	cbStarted := make(chan struct{})
 	cbFinished := make(chan struct{})
 	cb := func(ctx context.Context, d time.Duration) {
@@ -70,19 +73,19 @@ func TestCallbackQueue_processCallBacks_starts_and_cancels(t *testing.T) {
 		close(cbFinished)
 	}
 	q.list.PushBack(&callBackRequest{fn: cb, tm: time.Now()})
-	// run the processCallBacks method
+	// Run the processCallBacks method
 	go q.processCallBacks()
 	<-cbStarted // wait for the callback to start processing
 	assertTrue(t, q.list.Len() == 0, "Callback queue should be empty after processing")
 	close(q.closeSignal)
-	// the context was canceled when closeSignal was closed.
+	// The context was canceled when closeSignal was closed.
 	<-cbFinished
 }
 
 func TestCallbackQueue_processCallBacks_auto_dequeues(t *testing.T) {
 	q := newUnopenedCallBackQueue()
 	go q.processCallBacks()
-	// stage a callback to be processed
+	// Stage a callback to be processed.
 	n := 10
 	var wg sync.WaitGroup
 	wg.Add(n)
@@ -92,7 +95,7 @@ func TestCallbackQueue_processCallBacks_auto_dequeues(t *testing.T) {
 		}
 		q.list.PushBack(&callBackRequest{fn: cb, tm: time.Now()})
 	}
-	// only send one signal, the processCallBacks should dequeue the rest.
+	// Only send one signal; processCallBacks should dequeue the rest.
 	q.signalEnqueue()
 	wg.Wait()
 	assertTrue(t, q.list.Len() == 0, "Callback queue should be empty after processing all callbacks")
@@ -106,7 +109,7 @@ func TestCallbackQueue_Push_does_not_block(t *testing.T) {
 	neverProcessed := func(ctx context.Context, d time.Duration) {}
 	n := 100
 	for range n {
-		// push callbacks while there is nothing to dequeue.
+		// Push callbacks while there is nothing to dequeue.
 		err := q.Push(neverProcessed)
 		assertNoError(t, err, "Push should not return an error")
 	}
@@ -146,7 +149,7 @@ func TestCallbackQueue_Push_order_preserved(t *testing.T) {
 	// Process callbacks.
 	results := make([]string, len(expectedResults))
 	for i, v := range expectedResults {
-		i, v := i, v // Capture loop variables
+		i, v := i, v // Capture loop variables.
 		err := q.Push(func(_ context.Context, _ time.Duration) {
 			defer wg.Done()
 			results[i] = v
@@ -175,15 +178,18 @@ func TestCallbackQueue_Close(t *testing.T) {
 	q.closeSignal <- struct{}{} // This should panic.
 }
 
+func TestCallbackQueue_Close_multiple_calls_no_ops(t *testing.T) {
+	q := OpenCallBackQueue()
+	for range 10 {
+		q.Close()
+	}
+}
+
 func TestCallbackQueue_Push_after_close_returns_ErrQueueClosed(t *testing.T) {
 	q := OpenCallBackQueue()
 	q.Close()
-	var executed bool
-	err := q.Push(func(_ context.Context, _ time.Duration) {
-		executed = true
-	})
+	err := q.Push(func(_ context.Context, _ time.Duration) {})
 	assertErrorIs(t, err, ErrQueueClosed, "Push should return an error after queue close")
-	assertTrue(t, !executed, "Callback should not be executed after queue close")
 }
 
 func TestCallbackQueue_Push_unopened_returns_ErrClosed(t *testing.T) {