Multiple Arguments Pool Function

README.md

@@ -67,7 +67,7 @@ go get -u github.com/panjf2000/ants/v2
 ## 🛠 How to use
 Just take a imagination that your program starts a massive number of goroutines, resulting in a huge consumption of memory. To mitigate that kind of situation, all you need to do is to import `ants` package and submit all your tasks to a default pool with fixed capacity, activated when package `ants` is imported:
 
-``` go
+```go
 package main
 
 import (
@@ -81,12 +81,13 @@ import (
 
 var sum int32
 
-func myFunc(i interface{}) {
-	n := i.(int32)
+func myFunc(i int32, s string) {
+	n := i
 	atomic.AddInt32(&sum, n)
-	fmt.Printf("run with %d\n", n)
+	fmt.Printf("run with %d: %s\n", n, s)
 }
 
+
 func demoFunc() {
 	time.Sleep(10 * time.Millisecond)
 	fmt.Println("Hello World!")
@@ -113,15 +114,15 @@ func main() {
 
 	// Use the pool with a function,
 	// set 10 to the capacity of goroutine pool and 1 second for expired duration.
-	p, _ := ants.NewPoolWithFunc(10, func(i interface{}) {
-		myFunc(i)
+	p, _ := ants.NewPoolWithFunc(10, func(args ...interface{}) {
+		myFunc(args[0].(int32), args[1].(string))
 		wg.Done()
 	})
 	defer p.Release()
 	// Submit tasks one by one.
 	for i := 0; i < runTimes; i++ {
 		wg.Add(1)
-		_ = p.Invoke(int32(i))
+		_ = p.Invoke(int32(i), fmt.Sprintf("hello %d", i))
 	}
 	wg.Wait()
 	fmt.Printf("running goroutines: %d\n", p.Running())

ants_benchmark_test.go

@@ -40,19 +40,25 @@ func demoFunc() {
 	time.Sleep(time.Duration(BenchParam) * time.Millisecond)
 }
 
-func demoPoolFunc(args interface{}) {
-	n := args.(int)
+func demoPoolFunc(args ...interface{}) {
+	n := args[0].(int)
 	time.Sleep(time.Duration(n) * time.Millisecond)
 }
 
+func demoPoolMultipleArgsFunc(arg1 int, arg2 string) string {
+	n := arg1
+	time.Sleep(time.Duration(n) * time.Millisecond)
+	return arg2
+}
+
 func longRunningFunc() {
 	for {
 		runtime.Gosched()
 	}
 }
 
-func longRunningPoolFunc(arg interface{}) {
-	if ch, ok := arg.(chan struct{}); ok {
+func longRunningPoolFunc(args ...interface{}) {
+	if ch, ok := args[0].(chan struct{}); ok {
 		<-ch
 		return
 	}

ants_test.go

@@ -47,10 +47,11 @@ const (
 )
 
 const (
-	Param    = 100
-	AntsSize = 1000
-	TestSize = 10000
-	n        = 100000
+	Param       = 100
+	AntsSize    = 1000
+	TestSize    = 10000
+	n           = 100000
+	StringParam = "test_string"
 )
 
 var curMem uint64
@@ -99,8 +100,8 @@ func TestAntsPoolWaitToGetWorkerPreMalloc(t *testing.T) {
 // TestAntsPoolWithFuncWaitToGetWorker is used to test waiting to get worker.
 func TestAntsPoolWithFuncWaitToGetWorker(t *testing.T) {
 	var wg sync.WaitGroup
-	p, _ := NewPoolWithFunc(AntsSize, func(i interface{}) {
-		demoPoolFunc(i)
+	p, _ := NewPoolWithFunc(AntsSize, func(args ...interface{}) {
+		demoPoolFunc(args[0])
 		wg.Done()
 	})
 	defer p.Release()
@@ -117,10 +118,31 @@ func TestAntsPoolWithFuncWaitToGetWorker(t *testing.T) {
 	t.Logf("memory usage:%d MB", curMem)
 }
 
+func TestAntsPoolWithFuncMultipleArgsWaitToGetWorker(t *testing.T) {
+	var wg sync.WaitGroup
+	p, _ := NewPoolWithFunc(AntsSize, func(args ...interface{}) {
+		str := demoPoolMultipleArgsFunc(args[0].(int), args[1].(string))
+		if str != StringParam {
+			t.Errorf("return value %s is not correct, expected %s", str, StringParam)
+		}
+		wg.Done()
+	})
+	defer p.Release()
+
+	wg.Add(1)
+	_ = p.Invoke(4, StringParam)
+	wg.Wait()
+	t.Logf("pool with func, running workers number:%d", p.Running())
+	mem := runtime.MemStats{}
+	runtime.ReadMemStats(&mem)
+	curMem = mem.TotalAlloc/MiB - curMem
+	t.Logf("memory usage:%d MB", curMem)
+}
+
 func TestAntsPoolWithFuncWaitToGetWorkerPreMalloc(t *testing.T) {
 	var wg sync.WaitGroup
-	p, _ := NewPoolWithFunc(AntsSize, func(i interface{}) {
-		demoPoolFunc(i)
+	p, _ := NewPoolWithFunc(AntsSize, func(args ...interface{}) {
+		demoPoolFunc(args[0])
 		wg.Done()
 	}, WithPreAlloc(true))
 	defer p.Release()
@@ -253,7 +275,7 @@ func TestPanicHandler(t *testing.T) {
 	c := atomic.LoadInt64(&panicCounter)
 	assert.EqualValuesf(t, 1, c, "panic handler didn't work, panicCounter: %d", c)
 	assert.EqualValues(t, 0, p0.Running(), "pool should be empty after panic")
-	p1, err := NewPoolWithFunc(10, func(p interface{}) { panic(p) }, WithPanicHandler(func(p interface{}) {
+	p1, err := NewPoolWithFunc(10, func(args ...interface{}) { panic(args[0]) }, WithPanicHandler(func(p interface{}) {
 		defer wg.Done()
 		atomic.AddInt64(&panicCounter, 1)
 	}))
@@ -285,7 +307,7 @@ func TestPanicHandlerPreMalloc(t *testing.T) {
 	c := atomic.LoadInt64(&panicCounter)
 	assert.EqualValuesf(t, 1, c, "panic handler didn't work, panicCounter: %d", c)
 	assert.EqualValues(t, 0, p0.Running(), "pool should be empty after panic")
-	p1, err := NewPoolWithFunc(10, func(p interface{}) { panic(p) }, WithPanicHandler(func(p interface{}) {
+	p1, err := NewPoolWithFunc(10, func(args ...interface{}) { panic(args[0]) }, WithPanicHandler(func(p interface{}) {
 		defer wg.Done()
 		atomic.AddInt64(&panicCounter, 1)
 	}))
@@ -307,8 +329,8 @@ func TestPoolPanicWithoutHandler(t *testing.T) {
 		panic("Oops!")
 	})
 
-	p1, err := NewPoolWithFunc(10, func(p interface{}) {
-		panic(p)
+	p1, err := NewPoolWithFunc(10, func(p ...interface{}) {
+		panic(p[0])
 	})
 	assert.NoErrorf(t, err, "create new pool with func failed: %v", err)
 	defer p1.Release()
@@ -323,8 +345,8 @@ func TestPoolPanicWithoutHandlerPreMalloc(t *testing.T) {
 		panic("Oops!")
 	})
 
-	p1, err := NewPoolWithFunc(10, func(p interface{}) {
-		panic(p)
+	p1, err := NewPoolWithFunc(10, func(p ...interface{}) {
+		panic(p[0])
 	})
 
 	assert.NoErrorf(t, err, "create new pool with func failed: %v", err)
@@ -428,8 +450,8 @@ func TestMaxBlockingSubmit(t *testing.T) {
 func TestNonblockingSubmitWithFunc(t *testing.T) {
 	poolSize := 10
 	var wg sync.WaitGroup
-	p, err := NewPoolWithFunc(poolSize, func(i interface{}) {
-		longRunningPoolFunc(i)
+	p, err := NewPoolWithFunc(poolSize, func(args ...interface{}) {
+		longRunningPoolFunc(args[0])
 		wg.Done()
 	}, WithNonblocking(true))
 	assert.NoError(t, err, "create TimingPool failed: %v", err)
@@ -520,8 +542,8 @@ func TestRebootNewPool(t *testing.T) {
 	assert.NoError(t, p.Submit(func() { wg.Done() }), "pool should be rebooted")
 	wg.Wait()
 
-	p1, err := NewPoolWithFunc(10, func(i interface{}) {
-		demoPoolFunc(i)
+	p1, err := NewPoolWithFunc(10, func(args ...interface{}) {
+		demoPoolFunc(args[0])
 		wg.Done()
 	})
 	assert.NoErrorf(t, err, "create TimingPoolWithFunc failed: %v", err)
@@ -652,7 +674,7 @@ func TestWithDisablePurgePoolFunc(t *testing.T) {
 	var wg1, wg2 sync.WaitGroup
 	wg1.Add(numWorker)
 	wg2.Add(numWorker)
-	p, _ := NewPoolWithFunc(numWorker, func(i interface{}) {
+	p, _ := NewPoolWithFunc(numWorker, func(args ...interface{}) {
 		wg1.Done()
 		<-sig
 		wg2.Done()
@@ -667,7 +689,7 @@ func TestWithDisablePurgeAndWithExpirationPoolFunc(t *testing.T) {
 	wg1.Add(numWorker)
 	wg2.Add(numWorker)
 	expiredDuration := time.Millisecond * 100
-	p, _ := NewPoolWithFunc(numWorker, func(i interface{}) {
+	p, _ := NewPoolWithFunc(numWorker, func(args ...interface{}) {
 		wg1.Done()
 		<-sig
 		wg2.Done()
@@ -677,8 +699,8 @@ func TestWithDisablePurgeAndWithExpirationPoolFunc(t *testing.T) {
 
 func TestInfinitePoolWithFunc(t *testing.T) {
 	c := make(chan struct{})
-	p, _ := NewPoolWithFunc(-1, func(i interface{}) {
-		demoPoolFunc(i)
+	p, _ := NewPoolWithFunc(-1, func(args ...interface{}) {
+		demoPoolFunc(args[0])
 		<-c
 	})
 	_ = p.Invoke(10)
@@ -744,8 +766,8 @@ func TestReleaseWhenRunningPool(t *testing.T) {
 
 func TestReleaseWhenRunningPoolWithFunc(t *testing.T) {
 	var wg sync.WaitGroup
-	p, _ := NewPoolWithFunc(1, func(i interface{}) {
-		t.Log("do task", i)
+	p, _ := NewPoolWithFunc(1, func(args ...interface{}) {
+		t.Log("do task", args[0])
 		time.Sleep(1 * time.Second)
 	})
 	wg.Add(2)
@@ -899,7 +921,7 @@ func TestPoolTuneScaleUp(t *testing.T) {
 	p.Release()
 
 	// test PoolWithFunc
-	pf, _ := NewPoolWithFunc(2, func(i interface{}) {
+	pf, _ := NewPoolWithFunc(2, func(args ...interface{}) {
 		<-c
 	})
 	for i := 0; i < 2; i++ {
@@ -947,8 +969,8 @@ func TestReleaseTimeout(t *testing.T) {
 	assert.NoError(t, err)
 
 	var pf *PoolWithFunc
-	pf, _ = NewPoolWithFunc(10, func(i interface{}) {
-		dur := i.(time.Duration)
+	pf, _ = NewPoolWithFunc(10, func(args ...interface{}) {
+		dur := args[0].(time.Duration)
 		time.Sleep(dur)
 	})
 	for i := 0; i < 5; i++ {

examples/main.go

@@ -33,10 +33,10 @@ import (
 
 var sum int32
 
-func myFunc(i interface{}) {
-	n := i.(int32)
+func myFunc(i int32, s string) {
+	n := i
 	atomic.AddInt32(&sum, n)
-	fmt.Printf("run with %d\n", n)
+	fmt.Printf("run with %d: %s\n", n, s)
 }
 
 func demoFunc() {
@@ -65,15 +65,15 @@ func main() {
 
 	// Use the pool with a method,
 	// set 10 to the capacity of goroutine pool and 1 second for expired duration.
-	p, _ := ants.NewPoolWithFunc(10, func(i interface{}) {
-		myFunc(i)
+	p, _ := ants.NewPoolWithFunc(10, func(args ...interface{}) {
+		myFunc(args[0].(int32), args[1].(string))
 		wg.Done()
 	})
 	defer p.Release()
 	// Submit tasks one by one.
 	for i := 0; i < runTimes; i++ {
 		wg.Add(1)
-		_ = p.Invoke(int32(i))
+		_ = p.Invoke(int32(i), fmt.Sprintf("hello %d", i))
 	}
 	wg.Wait()
 	fmt.Printf("running goroutines: %d\n", p.Running())

pool_func.go

@@ -53,7 +53,7 @@ type PoolWithFunc struct {
 	cond *sync.Cond
 
 	// poolFunc is the function for processing tasks.
-	poolFunc func(interface{})
+	poolFunc func(...interface{})
 
 	// workerCache speeds up the obtainment of a usable worker in function:retrieveWorker.
 	workerCache sync.Pool
@@ -124,7 +124,7 @@ func (p *PoolWithFunc) purgePeriodically(ctx context.Context) {
 }
 
 // NewPoolWithFunc generates an instance of ants pool with a specific function.
-func NewPoolWithFunc(size int, pf func(interface{}), options ...Option) (*PoolWithFunc, error) {
+func NewPoolWithFunc(size int, pf func(...interface{}), options ...Option) (*PoolWithFunc, error) {
 	if size <= 0 {
 		size = -1
 	}
@@ -184,7 +184,7 @@ func NewPoolWithFunc(size int, pf func(interface{}), options ...Option) (*PoolWi
 // but what calls for special attention is that you will get blocked with the latest
 // Pool.Invoke() call once the current Pool runs out of its capacity, and to avoid this,
 // you should instantiate a PoolWithFunc with ants.WithNonblocking(true).
-func (p *PoolWithFunc) Invoke(args interface{}) error {
+func (p *PoolWithFunc) Invoke(args ...interface{}) error {
 	if p.IsClosed() {
 		return ErrPoolClosed
 	}

worker_func.go

@@ -67,7 +67,7 @@ func (w *goWorkerWithFunc) run() {
 			if args == nil {
 				return
 			}
-			w.pool.poolFunc(args)
+			w.pool.poolFunc(args.([]interface{})...)
 			if ok := w.pool.revertWorker(w); !ok {
 				return
 			}