Submit and SubmitWait return error if called on stopped workerpool

workerpool.go

@@ -2,6 +2,7 @@
 
 import (
 	"context"
+	"errors"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -14,6 +15,8 @@
 	idleTimeout = 2 * time.Second
 )
 
+var ErrStopped = errors.New("workerpool stopped")
+
 // New creates and starts a pool of worker goroutines.
 //
 // The maxWorkers parameter specifies the maximum number of workers that can
@@ -104,23 +107,51 @@
 // period until there are no more idle workers. Since the time to start new
 // goroutines is not significant, there is no need to retain idle workers
 // indefinitely.
-func (p *WorkerPool) Submit(task func()) {
-	if task != nil {
-		p.taskQueue <- task
+func (p *WorkerPool) Submit(task func()) (err error) {
+	if task == nil {
+		return
 	}
+	defer func() {
+		r := recover()
+		if r != nil {
+			if e, ok := r.(error); ok {
+				if e.Error() == "send on closed channel" {
+					err = ErrStopped
+					return
+				}
+			}
+			panic(r)
+		}
+	}()
+	p.taskQueue <- task
+	return
 }
 
 // SubmitWait enqueues the given function and waits for it to be executed.
-func (p *WorkerPool) SubmitWait(task func()) {
+func (p *WorkerPool) SubmitWait(task func()) (err error) {
 	if task == nil {
 		return
 	}
+	defer func() {
+		r := recover()
+		if r != nil {
+			if e, ok := r.(error); ok {
+				if e.Error() == "send on closed channel" {
+					err = ErrStopped
+					return
+				}
+			}
+			panic(r)
+		}
+	}()
+
 	doneChan := make(chan struct{})
 	p.taskQueue <- func() {
 		task()
 		close(doneChan)
 	}
 	<-doneChan
+	return
 }
 
 // WaitingQueueSize returns the count of tasks in the waiting queue.
@@ -149,7 +180,7 @@
 	ready := new(sync.WaitGroup)
 	ready.Add(p.maxWorkers)
 	for i := 0; i < p.maxWorkers; i++ {
 		p.Submit(func() {
 			ready.Done()
 			select {
 			case <-ctx.Done():

workerpool_test.go

@@ -2,6 +2,7 @@
 
 import (
 	"context"
+	"errors"
 	"sync"
 	"testing"
 	"time"
@@ -20,7 +21,7 @@
 	rspChan := make(chan string, len(requests))
 	for _, r := range requests {
 		r := r
 		wp.Submit(func() {
 			rspChan <- r
 		})
 	}
@@ -63,7 +64,7 @@
 
 	// Start workers, and have them all wait on a channel before completing.
 	for i := 0; i < max; i++ {
 		wp.Submit(func() {
 			started <- struct{}{}
 			<-release
 		})
@@ -97,7 +98,7 @@
 
 	// Cause worker to be created, and available for reuse before next task.
 	for i := 0; i < 10; i++ {
 		wp.Submit(func() { <-release })
 		release <- struct{}{}
 		time.Sleep(time.Millisecond)
 	}
@@ -265,7 +266,7 @@
 
 	// Check that these are noop.
 	wp.Submit(nil)
 	wp.SubmitWait(nil)
 
 	done1 := make(chan struct{})
 	wp.Submit(func() {
@@ -279,7 +280,7 @@
 	}
 
 	done2 := make(chan struct{})
 	wp.SubmitWait(func() {
 		time.Sleep(100 * time.Millisecond)
 		close(done2)
 	})
@@ -598,6 +599,44 @@
 	wp.Stop()
 }
 
+func TestPanicSubmit(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	wp := New(1)
+	wp.Stop()
+
+	done := make(chan struct{})
+	err := wp.Submit(func() {
+		close(done)
+	})
+	if err == nil {
+		t.Fatal("expected an error")
+	}
+	if !errors.Is(err, ErrStopped) {
+		t.Fatal("wrong error:", err)
+	}
+	select {
+	case <-done:
+		t.Fatal("task function should not have called")
+	default:
+	}
+
+	err = wp.SubmitWait(func() {
+		close(done)
+	})
+	if err == nil {
+		t.Fatal("expected an error")
+	}
+	if !errors.Is(err, ErrStopped) {
+		t.Fatal("wrong error:", err)
+	}
+	select {
+	case <-done:
+		t.Fatal("task function should not have called")
+	default:
+	}
+}
+
 func anyReady(w *WorkerPool) bool {
 	release := make(chan struct{})
 	wait := func() {