feat: implement ReleaseContext for all pool types (#387)

Fixes #384

Author: panjf2000

ants.go

@@ -139,6 +139,14 @@ func ReleaseTimeout(timeout time.Duration) error {
 	return defaultAntsPool.ReleaseTimeout(timeout)
 }
 
+// ReleaseContext is like Release but with a context, it waits all workers to exit before the context is done.
+//
+// Note that if the context is nil, it is the same as Release,
+// just return immediately without waiting for all workers to exit.
+func ReleaseContext(ctx context.Context) error {
+	return defaultAntsPool.ReleaseContext(ctx)
+}
+
 // Reboot reboots the default pool.
 func Reboot() {
 	defaultAntsPool.Reboot()
@@ -398,12 +406,31 @@ func (p *poolCommon) Release() {
 
 // ReleaseTimeout is like Release but with a timeout, it waits all workers to exit before timing out.
 func (p *poolCommon) ReleaseTimeout(timeout time.Duration) error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	err := p.ReleaseContext(ctx)
+	if errors.Is(err, context.DeadlineExceeded) {
+		return ErrTimeout
+	}
+	return err
+}
+
+// ReleaseContext is like Release but with a context, it waits all workers to exit before the context is done.
+//
+// Note that if the context is nil, it is the same as Release,
+// just return immediately without waiting for all workers to exit.
+func (p *poolCommon) ReleaseContext(ctx context.Context) error {
 	if p.IsClosed() || (!p.options.DisablePurge && p.stopPurge == nil) || p.stopTicktock == nil {
 		return ErrPoolClosed
 	}
 
 	p.Release()
 
+	// Don't wait for all workers to exit, just return immediately if the context is nil.
+	if ctx == nil {
+		return nil
+	}
+
 	var purgeCh <-chan struct{}
 	if !p.options.DisablePurge {
 		purgeCh = p.purgeCtx.Done()
@@ -417,12 +444,10 @@ func (p *poolCommon) ReleaseTimeout(timeout time.Duration) error {
 		})
 	}
 
-	timer := time.NewTimer(timeout)
-	defer timer.Stop()
 	for {
 		select {
-		case <-timer.C:
-			return ErrTimeout
+		case <-ctx.Done():
+			return ctx.Err()
 		case <-p.allDone:
 			<-purgeCh
 			<-p.ticktockCtx.Done()

ants_test.go

@@ -23,6 +23,7 @@
 package ants_test
 
 import (
+	"context"
 	"log"
 	"os"
 	"runtime"
@@ -1351,6 +1352,34 @@ func TestDefaultPoolReleaseTimeout(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestDefaultPoolReleaseContext(t *testing.T) {
+	ants.Reboot()
+	for i := 0; i < 5; i++ {
+		_ = ants.Submit(func() {
+			time.Sleep(time.Second)
+		})
+	}
+	require.NotZero(t, ants.Running())
+	err := ants.ReleaseContext(context.Background())
+	require.NoError(t, err)
+}
+
+func TestReleaseContextWithNil(t *testing.T) {
+	p, err := ants.NewPool(10)
+	require.NoError(t, err)
+	for i := 0; i < 5; i++ {
+		_ = p.Submit(func() {
+			time.Sleep(time.Second)
+		})
+	}
+	require.NotZero(t, p.Running())
+
+	// Passing nil context should release immediately without waiting for workers to exit.
+	err = p.ReleaseContext(nil) //nolint:staticcheck
+	require.NoError(t, err)
+	require.True(t, p.IsClosed())
+}
+
 func TestMultiPool(t *testing.T) {
 	_, err := ants.NewMultiPool(-1, 10, 8)
 	require.ErrorIs(t, err, ants.ErrInvalidMultiPoolSize)
@@ -1536,6 +1565,189 @@ func TestMultiPoolWithFuncGeneric(t *testing.T) {
 	mp.Tune(10)
 }
 
+func TestMultiPoolReleaseContext(t *testing.T) {
+	mp, err := ants.NewMultiPool(10, 5, ants.RoundRobin)
+	require.NoError(t, err)
+
+	for i := 0; i < 50; i++ {
+		err = mp.Submit(longRunningFunc)
+		require.NoError(t, err)
+	}
+	require.EqualValues(t, 50, mp.Running())
+
+	// Signal workers to stop, then release with a background context.
+	atomic.StoreInt32(&stopLongRunningFunc, 1)
+	err = mp.ReleaseContext(context.Background())
+	require.NoError(t, err)
+	require.Zero(t, mp.Running())
+	require.True(t, mp.IsClosed())
+	atomic.StoreInt32(&stopLongRunningFunc, 0)
+
+	// Calling ReleaseContext on a closed pool should return ErrPoolClosed.
+	require.ErrorIs(t, mp.ReleaseContext(context.Background()), ants.ErrPoolClosed)
+
+	// Test with LeastTasks strategy.
+	mp, err = ants.NewMultiPool(10, 5, ants.LeastTasks)
+	require.NoError(t, err)
+	for i := 0; i < 50; i++ {
+		err = mp.Submit(longRunningFunc)
+		require.NoError(t, err)
+	}
+	require.EqualValues(t, 50, mp.Running())
+
+	atomic.StoreInt32(&stopLongRunningFunc, 1)
+	err = mp.ReleaseContext(context.Background())
+	require.NoError(t, err)
+	require.Zero(t, mp.Running())
+	require.True(t, mp.IsClosed())
+	atomic.StoreInt32(&stopLongRunningFunc, 0)
+
+	// Test that a cancelled context returns an error.
+	mp, err = ants.NewMultiPool(10, 5, ants.RoundRobin)
+	require.NoError(t, err)
+	for i := 0; i < 50; i++ {
+		err = mp.Submit(longRunningFunc)
+		require.NoError(t, err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel immediately
+	err = mp.ReleaseContext(ctx)
+	require.Error(t, err)
+	atomic.StoreInt32(&stopLongRunningFunc, 1)
+	require.Eventually(t, func() bool {
+		return mp.Running() == 0
+	}, 3*time.Second, 100*time.Millisecond)
+	atomic.StoreInt32(&stopLongRunningFunc, 0)
+
+	// Test reboot after ReleaseContext.
+	mp, err = ants.NewMultiPool(10, 5, ants.RoundRobin)
+	require.NoError(t, err)
+	for i := 0; i < 50; i++ {
+		err = mp.Submit(longRunningFunc)
+		require.NoError(t, err)
+	}
+	atomic.StoreInt32(&stopLongRunningFunc, 1)
+	err = mp.ReleaseContext(context.Background())
+	require.NoError(t, err)
+	atomic.StoreInt32(&stopLongRunningFunc, 0)
+
+	mp.Reboot()
+	require.False(t, mp.IsClosed())
+	for i := 0; i < 50; i++ {
+		err = mp.Submit(longRunningFunc)
+		require.NoError(t, err)
+	}
+	require.EqualValues(t, 50, mp.Running())
+	atomic.StoreInt32(&stopLongRunningFunc, 1)
+	err = mp.ReleaseContext(context.Background())
+	require.NoError(t, err)
+	atomic.StoreInt32(&stopLongRunningFunc, 0)
+}
+
+func TestMultiPoolWithFuncReleaseContext(t *testing.T) {
+	ch := make(chan struct{})
+	mp, err := ants.NewMultiPoolWithFunc(10, 5, longRunningPoolFunc, ants.RoundRobin)
+	require.NoError(t, err)
+
+	for i := 0; i < 50; i++ {
+		err = mp.Invoke(ch)
+		require.NoError(t, err)
+	}
+	require.EqualValues(t, 50, mp.Running())
+
+	close(ch)
+	err = mp.ReleaseContext(context.Background())
+	require.NoError(t, err)
+	require.Zero(t, mp.Running())
+	require.True(t, mp.IsClosed())
+
+	// Calling ReleaseContext on a closed pool should return ErrPoolClosed.
+	require.ErrorIs(t, mp.ReleaseContext(context.Background()), ants.ErrPoolClosed)
+
+	// Test with LeastTasks strategy.
+	ch = make(chan struct{})
+	mp, err = ants.NewMultiPoolWithFunc(10, 5, longRunningPoolFunc, ants.LeastTasks)
+	require.NoError(t, err)
+	for i := 0; i < 50; i++ {
+		err = mp.Invoke(ch)
+		require.NoError(t, err)
+	}
+	close(ch)
+	err = mp.ReleaseContext(context.Background())
+	require.NoError(t, err)
+	require.Zero(t, mp.Running())
+	require.True(t, mp.IsClosed())
+
+	// Test that a cancelled context returns an error.
+	ch = make(chan struct{})
+	mp, err = ants.NewMultiPoolWithFunc(10, 5, longRunningPoolFunc, ants.RoundRobin)
+	require.NoError(t, err)
+	for i := 0; i < 50; i++ {
+		err = mp.Invoke(ch)
+		require.NoError(t, err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err = mp.ReleaseContext(ctx)
+	require.Error(t, err)
+	close(ch)
+	require.Eventually(t, func() bool {
+		return mp.Running() == 0
+	}, 3*time.Second, 100*time.Millisecond)
+}
+
+func TestMultiPoolWithFuncGenericReleaseContext(t *testing.T) {
+	ch := make(chan struct{})
+	mp, err := ants.NewMultiPoolWithFuncGeneric(10, 5, longRunningPoolFuncCh, ants.RoundRobin)
+	require.NoError(t, err)
+
+	for i := 0; i < 50; i++ {
+		err = mp.Invoke(ch)
+		require.NoError(t, err)
+	}
+	require.EqualValues(t, 50, mp.Running())
+
+	close(ch)
+	err = mp.ReleaseContext(context.Background())
+	require.NoError(t, err)
+	require.Zero(t, mp.Running())
+	require.True(t, mp.IsClosed())
+
+	// Calling ReleaseContext on a closed pool should return ErrPoolClosed.
+	require.ErrorIs(t, mp.ReleaseContext(context.Background()), ants.ErrPoolClosed)
+
+	// Test with LeastTasks strategy.
+	ch = make(chan struct{})
+	mp, err = ants.NewMultiPoolWithFuncGeneric(10, 5, longRunningPoolFuncCh, ants.LeastTasks)
+	require.NoError(t, err)
+	for i := 0; i < 50; i++ {
+		err = mp.Invoke(ch)
+		require.NoError(t, err)
+	}
+	close(ch)
+	err = mp.ReleaseContext(context.Background())
+	require.NoError(t, err)
+	require.Zero(t, mp.Running())
+	require.True(t, mp.IsClosed())
+
+	// Test that a cancelled context returns an error.
+	ch = make(chan struct{})
+	mp, err = ants.NewMultiPoolWithFuncGeneric(10, 5, longRunningPoolFuncCh, ants.RoundRobin)
+	require.NoError(t, err)
+	for i := 0; i < 50; i++ {
+		err = mp.Invoke(ch)
+		require.NoError(t, err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err = mp.ReleaseContext(ctx)
+	require.Error(t, err)
+	close(ch)
+	require.Eventually(t, func() bool {
+		return mp.Running() == 0
+	}, 3*time.Second, 100*time.Millisecond)
+}
+
 func TestRebootNewPoolCalc(t *testing.T) {
 	atomic.StoreInt32(&sum, 0)
 	runTimes := 1000

multipool.go

@@ -23,6 +23,7 @@
 package ants
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"math"
@@ -218,6 +219,45 @@ func (mp *MultiPool) ReleaseTimeout(timeout time.Duration) error {
 	return errors.New(strings.TrimSuffix(errStr.String(), " | "))
 }
 
+// ReleaseContext closes the multi-pool with a context,
+// it waits all pools to be closed before the context is done.
+func (mp *MultiPool) ReleaseContext(ctx context.Context) error {
+	if !atomic.CompareAndSwapInt32(&mp.state, OPENED, CLOSED) {
+		return ErrPoolClosed
+	}
+
+	errCh := make(chan error, len(mp.pools))
+	var wg errgroup.Group
+	for i, pool := range mp.pools {
+		func(p *Pool, idx int) {
+			wg.Go(func() error {
+				err := p.ReleaseContext(ctx)
+				if err != nil {
+					err = fmt.Errorf("pool %d: %v", idx, err)
+				}
+				errCh <- err
+				return err
+			})
+		}(pool, i)
+	}
+
+	_ = wg.Wait()
+
+	var errStr strings.Builder
+	for i := 0; i < len(mp.pools); i++ {
+		if err := <-errCh; err != nil {
+			errStr.WriteString(err.Error())
+			errStr.WriteString(" | ")
+		}
+	}
+
+	if errStr.Len() == 0 {
+		return nil
+	}
+
+	return errors.New(strings.TrimSuffix(errStr.String(), " | "))
+}
+
 // Reboot reboots a released multi-pool.
 func (mp *MultiPool) Reboot() {
 	if atomic.CompareAndSwapInt32(&mp.state, CLOSED, OPENED) {