diff --git a/deferred_executor.go b/deferred_executor.go index ee00b57..a6f46c7 100644 --- a/deferred_executor.go +++ b/deferred_executor.go @@ -33,8 +33,7 @@ import ( // their order. type DeferredExecutor struct { queue chan func() - closeChan chan bool - closed chan bool + closed chan struct{} closeOnce sync.Once } @@ -43,28 +42,24 @@ func NewDeferredExecutor(queueSize int) *DeferredExecutor { queueSize = 0 } result := &DeferredExecutor{ - queue: make(chan func(), queueSize), - closeChan: make(chan bool, 1), - closed: make(chan bool, 1), + queue: make(chan func(), queueSize), + closed: make(chan struct{}), } go result.run() return result } func (e *DeferredExecutor) run() { -loop: + defer close(e.closed) + for { - select { - case f := <-e.queue: - if f == nil { - break loop - } - f() - case <-e.closeChan: - break loop + f := <-e.queue + if f == nil { + break } + + f() } - e.closed <- true } func getFunctionName(i interface{}) string { @@ -83,14 +78,9 @@ func (e *DeferredExecutor) Execute(f func()) { } func (e *DeferredExecutor) Close() { - select { - case e.closeChan <- true: - e.closeOnce.Do(func() { - close(e.queue) - }) - default: - // Already closed. - } + e.closeOnce.Do(func() { + close(e.queue) + }) } func (e *DeferredExecutor) waitForStop() { diff --git a/deferred_executor_test.go b/deferred_executor_test.go index 5aa8c08..6e1b12c 100644 --- a/deferred_executor_test.go +++ b/deferred_executor_test.go @@ -109,3 +109,12 @@ func TestDeferredExecutor_DeferAfterClose(t *testing.T) { t.Error("method should not have been called") }) } + +func TestDeferredExecutor_WaitForStopTwice(t *testing.T) { + e := NewDeferredExecutor(64) + defer e.waitForStop() + + e.Close() + + e.waitForStop() +}