feat(pool): add support for dynamically adjusting max concurrency

Author: alitto

README.md

@@ -42,6 +42,7 @@ Some common use cases include:
   - Panics recovery (panics are captured and returned as errors)
   - Subpools with a fraction of the parent pool's maximum number of workers
   - Blocking and non-blocking submission of tasks when the queue is full
+  - Dynamic resizing of the pool
 - [API reference](https://pkg.go.dev/github.com/alitto/pond/v2)
 
 ## Installation
@@ -406,6 +407,40 @@ When a pool defines a queue size (bounded), you can also specify how to handle t
 pool := pond.NewPool(1, pond.WithQueueSize(10), pond.WithNonBlocking(true))
 ```
 
+### Resizing pools (v2)
+
+You can dynamically change the maximum number of workers in a pool using the `Resize` method. This is useful when you need to adjust the pool's capacity based on runtime conditions.
+
+``` go
+// Create a pool with 5 workers
+pool := pond.NewPool(5)
+
+// Submit some tasks
+for i := 0; i < 20; i++ {
+    pool.Submit(func() {
+        // Do some work
+    })
+}
+
+// Increase the pool size to 10 workers
+pool.Resize(10)
+
+// Submit more tasks that will use the increased capacity
+for i := 0; i < 20; i++ {
+    pool.Submit(func() {
+        // Do some work
+    })
+}
+
+// Decrease the pool size back to 5 workers
+pool.Resize(5)
+```
+
+When resizing a pool:
+- The new maximum concurrency must be greater than 0
+- If you increase the size, new workers will be created as needed up to the new maximum
+- If you decrease the size, existing workers will continue running until they complete their current tasks, but no new workers will be created until the number of running workers is below the new maximum
+
 ### Metrics & monitoring
 
 Each worker pool instance exposes useful metrics that can be queried through the following methods:

pool.go

@@ -20,9 +20,10 @@ const (
 )
 
 var (
-	ErrQueueFull   = errors.New("queue is full")
-	ErrQueueEmpty  = errors.New("queue is empty")
-	ErrPoolStopped = errors.New("pool stopped")
+	ErrQueueFull             = errors.New("queue is full")
+	ErrQueueEmpty            = errors.New("queue is empty")
+	ErrPoolStopped           = errors.New("pool stopped")
+	ErrMaxConcurrencyReached = errors.New("max concurrency reached")
 
 	poolStoppedFuture = func() Task {
 		future, resolve := future.NewFuture(context.Background())
@@ -73,6 +74,11 @@ type basePool interface {
 
 	// Returns true if the pool has been stopped or its context has been cancelled.
 	Stopped() bool
+
+	// Resizes the pool by changing the maximum concurrency (number of workers) of the pool.
+	// The new max concurrency must be greater than 0.
+	// If the new max concurrency is less than the current number of running workers, the pool will continue to run with the new max concurrency.
+	Resize(maxConcurrency int)
 }
 
 // Represents a pool of goroutines that can execute tasks concurrently.
@@ -125,9 +131,43 @@ func (p *pool) Stopped() bool {
 }
 
 func (p *pool) MaxConcurrency() int {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
 	return p.maxConcurrency
 }
 
+func (p *pool) Resize(maxConcurrency int) {
+	if maxConcurrency <= 0 {
+		panic(errors.New("maxConcurrency must be greater than 0"))
+	}
+
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
+	delta := maxConcurrency - p.maxConcurrency
+
+	p.maxConcurrency = maxConcurrency
+
+	if delta > 0 {
+		// Increase the number of workers by delta if there are tasks in the queue
+		queuedTasks := int(p.tasks.Len())
+
+		for i := 0; i < delta && i < queuedTasks; i++ {
+			p.workerCount.Add(1)
+
+			if p.parent == nil {
+				// Launch a new worker
+				p.workerWaitGroup.Add(1)
+				go p.worker(nil)
+			} else {
+				// Submit a task to the parent pool
+				p.subpoolSubmit(nil)
+			}
+		}
+	}
+}
+
 func (p *pool) QueueSize() int {
 	return p.queueSize
 }
@@ -165,9 +205,11 @@ func (p *pool) worker(task any) {
 
 	var readTaskErr, err error
 	for {
-		_, err = invokeTask[any](task)
+		if task != nil {
+			_, err = invokeTask[any](task)
 
-		p.updateMetrics(err)
+			p.updateMetrics(err)
+		}
 
 		task, readTaskErr = p.readTask()
 
@@ -304,9 +346,11 @@ func (p *pool) subpoolSubmit(task any) error {
 	return p.parent.submit(func() (output any, err error) {
 		defer p.workerWaitGroup.Done()
 
-		output, err = invokeTask[any](task)
+		if task != nil {
+			output, err = invokeTask[any](task)
 
-		p.updateMetrics(err)
+			p.updateMetrics(err)
+		}
 
 		// Attempt to submit the next task to the parent pool
 		if task, err := p.readTask(); err == nil {
@@ -338,6 +382,15 @@ func (p *pool) readTask() (task any, err error) {
 		return
 	}
 
+	if p.maxConcurrency > 0 && int(p.workerCount.Load()) > p.maxConcurrency {
+		// Max concurrency reached, kill the worker
+		p.workerCount.Add(-1)
+		p.mutex.Unlock()
+
+		err = ErrMaxConcurrencyReached
+		return
+	}
+
 	task, _ = p.tasks.Read()
 
 	p.mutex.Unlock()

pool_test.go

@@ -274,3 +274,68 @@ func TestPoolWithQueueSizeAndNonBlocking(t *testing.T) {
 
 	pool.Stop().Wait()
 }
+
+func TestPoolResize(t *testing.T) {
+
+	pool := NewPool(1, WithQueueSize(10))
+
+	assert.Equal(t, 1, pool.MaxConcurrency())
+
+	taskStarted := make(chan struct{}, 10)
+	taskWait := make(chan struct{}, 10)
+
+	// Submit 10 tasks
+	for i := 0; i < 10; i++ {
+		pool.Submit(func() {
+			<-taskStarted
+			<-taskWait
+		})
+	}
+
+	// Unblock 3 tasks
+	for i := 0; i < 3; i++ {
+		taskStarted <- struct{}{}
+	}
+
+	// Verify only 1 task is running and 9 are waiting
+	time.Sleep(10 * time.Millisecond)
+	assert.Equal(t, uint64(9), pool.WaitingTasks())
+	assert.Equal(t, int64(1), pool.RunningWorkers())
+
+	// Increase max concurrency to 3
+	pool.Resize(3)
+	assert.Equal(t, 3, pool.MaxConcurrency())
+
+	// Unblock 3 more tasks
+	for i := 0; i < 3; i++ {
+		taskStarted <- struct{}{}
+	}
+
+	// Verify 3 tasks are running and 7 are waiting
+	time.Sleep(10 * time.Millisecond)
+	assert.Equal(t, uint64(7), pool.WaitingTasks())
+	assert.Equal(t, int64(3), pool.RunningWorkers())
+
+	// Decrease max concurrency to 1
+	pool.Resize(2)
+	assert.Equal(t, 2, pool.MaxConcurrency())
+
+	// Complete the 3 running tasks
+	for i := 0; i < 3; i++ {
+		taskWait <- struct{}{}
+	}
+
+	// Unblock all remaining tasks
+	for i := 0; i < 4; i++ {
+		taskStarted <- struct{}{}
+	}
+
+	// Ensure 2 tasks are running and 5 are waiting
+	time.Sleep(10 * time.Millisecond)
+	assert.Equal(t, uint64(5), pool.WaitingTasks())
+	assert.Equal(t, int64(2), pool.RunningWorkers())
+
+	close(taskWait)
+
+	pool.Stop().Wait()
+}

subpool_test.go

@@ -334,3 +334,76 @@ func TestSubpoolWithQueueSizeOverride(t *testing.T) {
 	subpool.StopAndWait()
 	pool.StopAndWait()
 }
+
+func TestSubpoolResize(t *testing.T) {
+
+	parentPool := NewPool(10, WithQueueSize(10))
+
+	pool := parentPool.NewSubpool(1)
+
+	assert.Equal(t, 1, pool.MaxConcurrency())
+	assert.Equal(t, 10, parentPool.MaxConcurrency())
+
+	taskStarted := make(chan struct{}, 10)
+	taskWait := make(chan struct{}, 10)
+
+	// Submit 10 tasks
+	for i := 0; i < 10; i++ {
+		pool.Submit(func() {
+			<-taskStarted
+			<-taskWait
+		})
+	}
+
+	// Unblock 3 tasks
+	for i := 0; i < 3; i++ {
+		taskStarted <- struct{}{}
+	}
+
+	// Verify only 1 task is running and 9 are waiting
+	time.Sleep(10 * time.Millisecond)
+	assert.Equal(t, uint64(9), pool.WaitingTasks())
+	assert.Equal(t, int64(1), pool.RunningWorkers())
+	assert.Equal(t, int64(1), parentPool.RunningWorkers())
+
+	// Increase max concurrency to 3
+	pool.Resize(3)
+	assert.Equal(t, 3, pool.MaxConcurrency())
+	assert.Equal(t, 10, parentPool.MaxConcurrency())
+
+	// Unblock 3 more tasks
+	for i := 0; i < 3; i++ {
+		taskStarted <- struct{}{}
+	}
+
+	// Verify 3 tasks are running and 7 are waiting
+	time.Sleep(10 * time.Millisecond)
+	assert.Equal(t, uint64(7), pool.WaitingTasks())
+	assert.Equal(t, int64(3), pool.RunningWorkers())
+	assert.Equal(t, int64(3), parentPool.RunningWorkers())
+
+	// Decrease max concurrency to 1
+	pool.Resize(2)
+	assert.Equal(t, 2, pool.MaxConcurrency())
+	assert.Equal(t, 10, parentPool.MaxConcurrency())
+
+	// Complete the 3 running tasks
+	for i := 0; i < 3; i++ {
+		taskWait <- struct{}{}
+	}
+
+	// Unblock all remaining tasks
+	for i := 0; i < 4; i++ {
+		taskStarted <- struct{}{}
+	}
+
+	// Ensure 2 tasks are running and 5 are waiting
+	time.Sleep(10 * time.Millisecond)
+	assert.Equal(t, uint64(5), pool.WaitingTasks())
+	assert.Equal(t, int64(2), pool.RunningWorkers())
+	assert.Equal(t, int64(2), parentPool.RunningWorkers())
+
+	close(taskWait)
+
+	pool.Stop().Wait()
+}