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

Author: alitto

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,46 @@ 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
+		for i := 0; i < delta; i++ {
+			if poppedTask, _ := p.tasks.Read(); poppedTask != nil {
+
+				p.workerCount.Add(1)
+
+				if p.parent == nil {
+					p.workerWaitGroup.Add(1)
+					// Launch a new worker
+					go p.worker(poppedTask)
+				} else {
+					// Submit task to the parent pool
+					p.subpoolSubmit(poppedTask)
+				}
+			} else {
+				return
+			}
+		}
+	}
+}
+
 func (p *pool) QueueSize() int {
 	return p.queueSize
 }
@@ -338,6 +381,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()
+}