restore concurrency limit and re-run benchmarks

Author: rubengp99

README.md

@@ -158,15 +158,15 @@ goos: linux, goarch: amd64, cpu: 13th Gen Intel i9-13900KS
 
 | Name                                      | Iterations | ns/op    | B/op   | allocs/op |
 |-------------------------------------------|-----------:|---------:|-------:|-----------:|
-| **ErrGroup**                               | 6,253,089  | **177.8** | **24** | **1**     |
-| GoPool                              | 4,754,223  | **255.9** | 80     | 2          |
-| ChannelsWithOutputAndErrChannel           | 4,463,799  | 258.3    | **72** | 2          |
-| GoPoolWithDrainer                    | 4,570,286  | 262.8    | 118    | 3          |
-| ChannelsWithWaitGroup                      | 4,499,217  | 270.5    | 80     | 2          |
-| ChannelsWithErrGroup                       | 4,336,857  | 277.6    | 80     | 2          |
-| MutexWithErrGroup                          | 4,380,441  | 368.9    | 127    | 2          |
-
-![Benchmark Comparison](go_async_benchmarks.png)
+| **ErrGroup**                               | 6,211,902  | **180.3** | **24** | **1**     |
+| **GoPool**                                 | 5,020,380  | **214.4** | 80     | 2          |
+| ChannelsWithOutputAndErrChannel           | 4,426,651  | 260.6    | **72** | 2          |
+| AsyncPackageWithDrainer                    | 4,531,092  | 274.5    | 119    | 3          |
+| ChannelsWithWaitGroup                      | 4,480,616  | 271.5    | 80     | 2          |
+| ChannelsWithErrGroup                       | 4,336,473  | 279.1    | 80     | 2          |
+| MutexWithErrGroup                          | 2,842,214  | 420.6    | 135    | 2          |
+
+![Benchmark Comparison](benchmark_chart.png)
 
 Even though `go-pool` adds a small constant overhead compared to `errgroup` (≈100–130 ns per operation),
 it provides type safety, retries, automatic draining, and deterministic cleanup — all while staying within ~1.7× of native concurrency performance.

go_async_benchmarks.png renamed to benchmark_chart.png

@@ -1,7 +1,6 @@
 package gopool_test
 
 import (
-	"os"
 	"sync"
 	"testing"
 
@@ -21,8 +20,6 @@ func SimulatedTask() error {
 
 // BenchmarkAsyncPackage benchmarks the `go-async` package.
 func BenchmarkAsyncPackage(b *testing.B) {
-	//disable internal limit on test
-	os.Setenv("STAGE", "prod")
 	// Create a Drain channel for async operations
 	d := gopool.NewPool()
 
@@ -37,13 +34,11 @@ func BenchmarkAsyncPackage(b *testing.B) {
 			b.Fatal(err)
 		}
 	})
-	os.Setenv("STAGE", "test")
 }
 
 // BenchmarkAsyncPackageWithDrainer benchmarks the `go-async` package + drain logic.
 func BenchmarkAsyncPackageWithDrainer(b *testing.B) {
 	//disable internal limit on test
-	os.Setenv("STAGE", "prod")
 	d := gopool.NewPool()
 
 	b.Run("GoPool", func(b *testing.B) {
@@ -61,7 +56,6 @@ func BenchmarkAsyncPackageWithDrainer(b *testing.B) {
 			b.Fatal(err)
 		}
 	})
-	os.Setenv("STAGE", "test")
 }
 
 // BenchmarkErrGroup benchmarks the `errgroup.Group`.

benchmark_test.go

@@ -1,22 +1,34 @@
 package gopool
 
 import (
+	"context"
+	"fmt"
 	"sync"
 	"time"
 )
 
+type token struct{}
+
+// Pool is a pool for N workers
 type Pool struct {
 	wg       sync.WaitGroup
 	attempts uint
 	sleep    time.Duration
 	errOnce  sync.Once
 	err      error
+	ctx      context.Context
 	cancel   func(error)
+	sem      chan token
 }
 
 // NewPool creates a pool
 func NewPool() *Pool {
-	return &Pool{attempts: 1}
+	return NewPoolWithContext(context.Background())
+}
+
+func NewPoolWithContext(ctx context.Context) *Pool {
+	ctx, cancel := context.WithCancelCause(ctx)
+	return &Pool{attempts: 1, cancel: cancel, ctx: ctx}
 }
 
 // WithRetry sets pool-wide retry config
@@ -30,6 +42,9 @@ func (p *Pool) WithRetry(attempts uint, sleep time.Duration) *Pool {
 func (p *Pool) Go(tasks ...Worker) *Pool {
 	for i := range tasks {
 		task := tasks[i] // copy by value internally
+		if p.sem != nil {
+			p.sem <- token{}
+		}
 		p.wg.Add(1)
 		go executeTask(p, task)
 	}
@@ -39,6 +54,10 @@ func (p *Pool) Go(tasks ...Worker) *Pool {
 // executeTask is standalone to avoid closure allocation
 func executeTask(p *Pool, w Worker) {
 	defer p.wg.Done()
+	if p.sem != nil {
+		<-p.sem
+	}
+
 	currentSleep := p.sleep
 	for i := uint(0); i < p.attempts; i++ {
 		err := w.Execute()
@@ -63,5 +82,21 @@ func executeTask(p *Pool, w Worker) {
 // Wait waits for all tasks to finish
 func (p *Pool) Wait() error {
 	p.wg.Wait()
+	if p.cancel != nil {
+		p.cancel(p.err)
+	}
 	return p.err
 }
+
+// WithLimit stablishes limitted concurrency behavior
+func (p *Pool) WithLimit(n int) *Pool {
+	if n < 0 {
+		p.sem = nil
+		return p
+	}
+	if len(p.sem) != 0 {
+		panic(fmt.Errorf("can't modify limit while %v goroutines still active", len(p.sem)))
+	}
+	p.sem = make(chan token, n)
+	return p
+}

pool.go

@@ -2,18 +2,13 @@ package gopool_test
 
 import (
 	"fmt"
-	"os"
 	"testing"
 	"time"
 
 	gopool "github.com/rubengp99/go-pool"
 	"github.com/stretchr/testify/assert"
 )
 
-func init() {
-	os.Setenv("STAGE", "test")
-}
-
 type typeA struct {
 	value string
 }
@@ -36,7 +31,7 @@ func TestConcurrentClient(t *testing.T) {
 		gopool.NewTask(tFunc),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 	err := pool.Go(requests...).Wait()
 
 	t.Run("No errors", func(t *testing.T) {
@@ -65,7 +60,7 @@ func TestConcurrentClientWithError(t *testing.T) {
 		gopool.NewTask(tFunc),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 	err := pool.Go(requests...).Wait()
 
 	t.Run("errors", func(t *testing.T) {
@@ -103,7 +98,7 @@ func TestConcurrentClientWithRetry(t *testing.T) {
 		}).WithRetry(3, 100*time.Millisecond),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 	err := pool.Go(requests...).Wait()
 
 	t.Run("6 requests done", func(t *testing.T) {
@@ -140,7 +135,7 @@ func TestConcurrentClientWithRetryFailure(t *testing.T) {
 		}).WithRetry(3, 100*time.Millisecond),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 
 	err := pool.Go(requests...).Wait()
 	t.Run("6 requests done", func(t *testing.T) {
@@ -168,15 +163,15 @@ func TestConcurrentClientWithAllRetry(t *testing.T) {
 		gopool.NewTask(func(t gopool.Args[any]) error {
 			numInvocations++
 
-			if numInvocations > 1 {
+			if numInvocations > 2 {
 				return nil
 			}
 
 			return fmt.Errorf("bye 2")
 		}),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 	pool.WithRetry(3, 100*time.Millisecond)
 
 	err := pool.Go(requests...).Wait()
@@ -205,7 +200,7 @@ func TestConcurrentClientWithTaskChannel(t *testing.T) {
 		gopool.NewTask(tFunc).DrainTo(output),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 
 	// Run the Task(s)
 	err := pool.Go(requests...).Wait()
@@ -250,7 +245,7 @@ func TestConcurrentClientWith2WorkersameChannel(t *testing.T) {
 		gopool.NewTask(tFunc2).DrainTo(output),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 
 	// Run the Task(s)
 	err := pool.Go(requests...).Wait()
@@ -268,8 +263,8 @@ func TestConcurrentClientWith2WorkersameChannel(t *testing.T) {
 
 	t.Run("results drained", func(t *testing.T) {
 		if assert.Equal(t, 2, len(results)) {
-			assert.Equal(t, "hello-world!", results[1].value)
-			assert.Equal(t, "hello-world!2", results[0].value)
+			assert.Equal(t, "hello-world!", results[0].value)
+			assert.Equal(t, "hello-world!2", results[1].value)
 		}
 	})
 }
@@ -297,7 +292,7 @@ func TestConcurrentClientWith2TaskDiffTypes(t *testing.T) {
 		gopool.NewTask(tFunc2).DrainTo(output2),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 
 	// Run the Task(s)
 	err := pool.Go(requests...).Wait()
@@ -346,7 +341,7 @@ func TestConcurrentClientWith2TaskDiffTypes1Output(t *testing.T) {
 		gopool.NewTask(tFunc2),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 
 	// Run the Task(s)
 	err := pool.Go(requests...).Wait()
@@ -395,7 +390,7 @@ func TestConcurrentClientWith2TaskDiffTypes1Output1Input(t *testing.T) {
 		gopool.NewTask(tFunc2).WithInput(&initial),
 	}
 
-	pool := gopool.NewPool()
+	pool := gopool.NewPool().WithLimit(1)
 
 	// Run the Task(s)
 	err := pool.Go(requests...).Wait()