Add configurable submit buffer

pool_test.go

@@ -21,6 +21,9 @@ func TestNewValidationAndNilTask(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	if cap(p.submit) != defaultSubmitBuffer {
+		t.Fatalf("default submit capacity=%d, want %d", cap(p.submit), defaultSubmitBuffer)
+	}
 	if err := p.Submit(nil); !errors.Is(err, ErrNilTask) {
 		t.Fatalf("Submit(nil) error=%v, want ErrNilTask", err)
 	}
@@ -30,6 +33,37 @@ func TestNewValidationAndNilTask(t *testing.T) {
 	}
 }
 
+func TestSubmitBufferOptions(t *testing.T) {
+	for _, test := range []struct {
+		name    string
+		options []Option
+		want    int
+		invalid bool
+	}{
+		{name: "default", want: 128},
+		{name: "custom", options: []Option{WithSubmitBuffer(7)}, want: 7},
+		{name: "zero", options: []Option{WithSubmitBuffer(0)}, want: 0},
+		{name: "negative", options: []Option{WithSubmitBuffer(-1)}, invalid: true},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			p, err := New(0, 2, test.options...)
+			if test.invalid {
+				if !errors.Is(err, ErrInvalidConfig) {
+					t.Fatalf("New error=%v, want ErrInvalidConfig", err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			if got := cap(p.submit); got != test.want {
+				t.Fatalf("submit buffer=%d, want %d", got, test.want)
+			}
+			p.Close()
+		})
+	}
+}
+
 func TestSchedulerStartsMinimumBeforeRun(t *testing.T) {
 	p := newTestPool(3, 5)
 	s := newScheduler(p)
@@ -91,7 +125,7 @@ func TestConcurrentSubmitRunsExactlyOnceAndNeverExceedsMaximum(t *testing.T) {
 }
 
 func TestGrowsToMaximumBeforeUsingBoundedQueue(t *testing.T) {
-	p, err := New(0, 3) // one cached task
+	p, err := New(0, 3, WithSubmitBuffer(0)) // one cached task, no ingress buffer
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -121,7 +155,7 @@ func TestGrowsToMaximumBeforeUsingBoundedQueue(t *testing.T) {
 }
 
 func TestQueueCapacityAndMultipleBlockedSubmittersWakeIncrementally(t *testing.T) {
-	p, err := New(0, 2) // queue capacity one
+	p, err := New(0, 2, WithSubmitBuffer(0)) // queue capacity one; zero ingress buffer
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -188,7 +222,7 @@ func TestQueueCapacityAndMultipleBlockedSubmittersWakeIncrementally(t *testing.T
 }
 
 func TestZeroQueueMinZeroAcceptsAfterWorkerFinishes(t *testing.T) {
-	p, err := New(0, 1)
+	p, err := New(0, 1, WithSubmitBuffer(0))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -217,7 +251,7 @@ func TestZeroQueueMinZeroAcceptsAfterWorkerFinishes(t *testing.T) {
 }
 
 func TestQueueTasksAreDispatchedFIFO(t *testing.T) {
-	p, err := New(4, 4) // two cached tasks
+	p, err := New(4, 4, WithSubmitBuffer(0)) // two cached tasks
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -643,7 +677,7 @@ func TestAcceptedPendingHandoffSurvivesDelayedExitAndClose(t *testing.T) {
 }
 
 func TestCloseDrainsAcceptedTasksAndWakesBlockedSubmitters(t *testing.T) {
-	p, err := New(2, 2) // queue capacity one
+	p, err := New(2, 2, WithSubmitBuffer(0)) // queue capacity one; zero ingress buffer
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -696,6 +730,98 @@ func TestCloseDrainsAcceptedTasksAndWakesBlockedSubmitters(t *testing.T) {
 	}
 }
 
+func TestCloseDrainsTasksAcceptedIntoSubmitBuffer(t *testing.T) {
+	const bufferedTasks = 2
+	p, err := New(1, 1, WithSubmitBuffer(bufferedTasks))
+	if err != nil {
+		t.Fatal(err)
+	}
+	gate := make(chan struct{})
+	started := make(chan int, bufferedTasks+1)
+	var calls [bufferedTasks + 1]atomic.Int32
+	if err := p.Submit(func() { started <- 0; <-gate; calls[0].Add(1) }); err != nil {
+		t.Fatal(err)
+	}
+	if got := waitSignal(t, started); got != 0 {
+		t.Fatalf("first task index=%d", got)
+	}
+	for i := 1; i <= bufferedTasks; i++ {
+		index := i
+		if err := p.Submit(func() { started <- index; calls[index].Add(1) }); err != nil {
+			t.Fatalf("buffered Submit(%d): %v", i, err)
+		}
+	}
+	closed := make(chan struct{})
+	go func() { p.Close(); close(closed) }()
+	waitSignal(t, p.stopping)
+	select {
+	case <-closed:
+		t.Fatal("Close returned while accepted work remained")
+	default:
+	}
+	close(gate)
+	waitSignal(t, closed)
+	for i := range calls {
+		if got := calls[i].Load(); got != 1 {
+			t.Errorf("task %d ran %d times, want once", i, got)
+		}
+	}
+}
+
+func TestFullSubmitBufferBlocksUntilSchedulerConsumes(t *testing.T) {
+	p, err := New(1, 1, WithSubmitBuffer(1))
+	if err != nil {
+		t.Fatal(err)
+	}
+	release := make(chan struct{})
+	started := make(chan int, 3)
+	var calls [3]atomic.Int32
+	if err := p.Submit(func() { started <- 0; <-release; calls[0].Add(1) }); err != nil {
+		t.Fatal(err)
+	}
+	waitSignal(t, started)
+	if err := p.Submit(func() { started <- 1; calls[1].Add(1) }); err != nil {
+		t.Fatal(err)
+	}
+	third := make(chan error, 1)
+	go func() { third <- p.Submit(func() { started <- 2; calls[2].Add(1) }) }()
+	assertBlocked(t, third, "Submit returned while its ingress buffer was full and worker was busy")
+	close(release)
+	if err := waitSignal(t, third); err != nil {
+		t.Fatalf("Submit after scheduler consumed ingress=%v", err)
+	}
+	p.Close()
+	for i := range calls {
+		if got := calls[i].Load(); got != 1 {
+			t.Errorf("task %d ran %d times, want once", i, got)
+		}
+	}
+}
+
+func TestCloseDrainsBufferedSubmissionsAfterWorkerPanic(t *testing.T) {
+	p := newTestPool(0, 1)
+	p.submit = make(chan func(), 3)
+	s := newScheduler(p)
+	w := s.startWorker(nil, false)
+	close(w.inbox)
+	var calls [3]atomic.Int32
+	for i := range calls {
+		index := i
+		if err := p.Submit(func() { calls[index].Add(1) }); err != nil {
+			t.Fatalf("buffered Submit(%d): %v", i, err)
+		}
+	}
+	go s.run()
+	closed := make(chan struct{})
+	go func() { p.Close(); close(closed) }()
+	waitSignal(t, closed)
+	for i := range calls {
+		if got := calls[i].Load(); got != 1 {
+			t.Errorf("buffered task %d ran %d times, want once", i, got)
+		}
+	}
+}
+
 func TestConcurrentSubmitAndCloseMatchesEveryResult(t *testing.T) {
 	const count = 100
 	p, err := New(2, 4)