Merge branch 'main' into iter-err

Author: camdencheek

pool/pool.go

@@ -47,20 +47,18 @@ func (p *Pool) Go(f func()) {
 		default:
 			// No goroutine was available to handle the task.
 			// Spawn a new one and send it the task.
-			p.handle.Go(p.worker)
-			p.tasks <- f
+			p.handle.Go(func() {
+				p.worker(f)
+			})
 		}
 	} else {
 		select {
 		case p.limiter <- struct{}{}:
 			// If we are below our limit, spawn a new worker rather
 			// than waiting for one to become available.
-			p.handle.Go(p.worker)
-
-			// We know there is at least one worker running, so wait
-			// for it to become available. This ensures we never spawn
-			// more workers than the number of tasks.
-			p.tasks <- f
+			p.handle.Go(func() {
+				p.worker(f)
+			})
 		case p.tasks <- f:
 			// A worker is available and has accepted the task.
 			return
@@ -76,6 +74,10 @@ func (p *Pool) Wait() {
 
 	close(p.tasks)
 
+	// After Wait() returns, reset the struct so tasks will be reinitialized on
+	// next use. This better matches the behavior of sync.WaitGroup
+	defer func() { p.initOnce = sync.Once{} }()
+
 	p.handle.Wait()
 }
 
@@ -145,11 +147,15 @@ func (p *Pool) WithContext(ctx context.Context) *ContextPool {
 	}
 }
 
-func (p *Pool) worker() {
+func (p *Pool) worker(initialFunc func()) {
 	// The only time this matters is if the task panics.
 	// This makes it possible to spin up new workers in that case.
 	defer p.limiter.release()
 
+	if initialFunc != nil {
+		initialFunc()
+	}
+
 	for f := range p.tasks {
 		f()
 	}

pool/pool_test.go

@@ -122,6 +122,26 @@ func TestPool(t *testing.T) {
 		p := New().WithMaxGoroutines(42)
 		require.Equal(t, 42, p.MaxGoroutines())
 	})
+
+	t.Run("is reusable", func(t *testing.T) {
+		t.Parallel()
+		var count atomic.Int64
+		p := New()
+		for i := 0; i < 10; i++ {
+			p.Go(func() {
+				count.Add(1)
+			})
+		}
+		p.Wait()
+		require.Equal(t, int64(10), count.Load())
+		for i := 0; i < 10; i++ {
+			p.Go(func() {
+				count.Add(1)
+			})
+		}
+		p.Wait()
+		require.Equal(t, int64(20), count.Load())
+	})
 }
 
 func BenchmarkPool(b *testing.B) {