diff --git a/pool.go b/pool.go index ad9542e..1ba2092 100644 --- a/pool.go +++ b/pool.go @@ -182,8 +182,11 @@ func (p *pool) Resize(maxConcurrency int) { p.mutex.Lock() - // Calculate the number of new workers to launch to reach the new max concurrency or the number of tasks in the queue, whichever is smaller - newWorkers := int(math.Min(float64(maxConcurrency-p.maxConcurrency), float64(p.tasks.Len()))) + // Workers from an earlier limit can still be running after a shrink. + newWorkers := maxConcurrency - int(p.workerCount.Load()) + if queuedTasks := int(p.tasks.Len()); newWorkers > queuedTasks { + newWorkers = queuedTasks + } p.maxConcurrency = maxConcurrency diff --git a/pool_test.go b/pool_test.go index 0616833..a185d6a 100644 --- a/pool_test.go +++ b/pool_test.go @@ -491,6 +491,141 @@ func TestPoolResize(t *testing.T) { pool.Stop().Wait() } +func TestPoolResizeBeforeWorkersRetire(t *testing.T) { + for _, subpool := range []bool{false, true} { + for _, test := range []struct { + name string + retired int + maxWorkers int + running int64 + startedQueued int + }{ + {name: "still_above_limit", maxWorkers: 3, running: 5}, + {name: "partially_retired", retired: 2, maxWorkers: 4, running: 4, startedQueued: 1}, + {name: "grow_above_original_limit", maxWorkers: 7, running: 7, startedQueued: 2}, + } { + name := test.name + if subpool { + name = "subpool/" + name + } + t.Run(name, func(t *testing.T) { + var parent, pool Pool + if subpool { + parent = NewPool(10) + pool = parent.NewSubpool(5) + } else { + pool = NewPool(5) + } + + started := make(chan struct{}, 5) + queuedStarted := make(chan struct{}, 2) + releaseQueued := make(chan struct{}) + releases := make([]chan struct{}, 5) + for i := range releases { + releases[i] = make(chan struct{}) + } + released := 0 + t.Cleanup(func() { + for _, release := range releases[released:] { + close(release) + } + close(releaseQueued) + pool.StopAndWait() + if parent != nil { + parent.StopAndWait() + } + }) + + for _, release := range releases { + release := release + pool.Submit(func() { + started <- struct{}{} + <-release + }) + } + for range releases { + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for initial workers") + } + } + for i := 0; i < 2; i++ { + pool.Submit(func() { + queuedStarted <- struct{}{} + <-releaseQueued + }) + } + + pool.Resize(1) + for _, release := range releases[:test.retired] { + close(release) + released++ + } + deadline := time.Now().Add(2 * time.Second) + for pool.RunningWorkers() != int64(5-test.retired) { + if time.Now().After(deadline) { + t.Fatal("timed out waiting for workers to retire") + } + time.Sleep(time.Millisecond) + } + + pool.Resize(test.maxWorkers) + assert.Equal(t, test.maxWorkers, pool.MaxConcurrency()) + assert.Equal(t, test.running, pool.RunningWorkers()) + for i := 0; i < test.startedQueued; i++ { + select { + case <-queuedStarted: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for queued tasks") + } + } + assert.Equal(t, test.startedQueued, 2-int(pool.WaitingTasks())) + if parent != nil { + assert.Equal(t, uint64(5+2*test.startedQueued), parent.SubmittedTasks()) + } + }) + } + } +} + +func TestSubpoolResizeBeforeWorkersRetireWithFullParent(t *testing.T) { + parent := NewPool(5, WithQueueSize(0)) + pool := parent.NewSubpool(5, WithQueueSize(2)) + started := make(chan struct{}, 5) + release := make(chan struct{}) + resized := make(chan struct{}) + t.Cleanup(func() { + close(release) + <-resized + pool.StopAndWait() + parent.StopAndWait() + }) + + for i := 0; i < 5; i++ { + pool.Submit(func() { + started <- struct{}{} + <-release + }) + } + for i := 0; i < 5; i++ { + <-started + } + pool.Submit(func() {}) + pool.Submit(func() {}) + pool.Resize(1) + go func() { + pool.Resize(3) + close(resized) + }() + + select { + case <-resized: + case <-time.After(2 * time.Second): + t.Fatal("resize blocked on the full parent while existing workers exceed the new limit") + } +} + func TestPoolResizeWithZeroMaxConcurrency(t *testing.T) { pool := NewPool(10)