Allow to add/remove workers during pool lifecycle

Author: vardius

.travis.yml

@@ -4,6 +4,7 @@ go:
   - tip
 script:
   - go build
-  - go test -v -cover -race -coverprofile=coverage.txt -covermode=atomic
+  - go test -v -race -cover -coverprofile=coverage.txt -covermode=atomic
+  - go test -bench=. -cpu=4 -benchmem
 after_script:
   - bash <(curl -s https://codecov.io/bash)

README.md

@@ -40,14 +40,15 @@ For **GoDoc** reference, **visit [pkg.go.dev](https://pkg.go.dev/github.com/vard
 ➜  worker-pool git:(master) ✗ go test -bench=. -cpu=4 -benchmem
 goos: darwin
 goarch: amd64
-BenchmarkWorker1-4                	 3000000	       453 ns/op	      56 B/op	       3 allocs/op
-BenchmarkWorker1Parallel-4        	 3000000	       506 ns/op	      48 B/op	       2 allocs/op
-BenchmarkWorker100-4              	 3000000	       485 ns/op	      56 B/op	       3 allocs/op
-BenchmarkWorker100Parallel-4      	 3000000	       444 ns/op	      48 B/op	       2 allocs/op
-BenchmarkWorkerNumCPU-4           	 3000000	       467 ns/op	      56 B/op	       3 allocs/op
-BenchmarkWorkerNumCPUParallel-4   	 3000000	       431 ns/op	      48 B/op	       2 allocs/op
+pkg: github.com/vardius/worker-pool/v2
+BenchmarkWorker1-4                	 3435846	       358 ns/op	     152 B/op	       4 allocs/op
+BenchmarkWorker1Parallel-4        	 2993271	       403 ns/op	     144 B/op	       3 allocs/op
+BenchmarkWorker100-4              	 2140670	       543 ns/op	     152 B/op	       4 allocs/op
+BenchmarkWorker100Parallel-4      	 3379311	       332 ns/op	     144 B/op	       3 allocs/op
+BenchmarkWorkerNumCPU-4           	 2536502	       438 ns/op	     152 B/op	       4 allocs/op
+BenchmarkWorkerNumCPUParallel-4   	 3061671	       353 ns/op	     144 B/op	       3 allocs/op
 PASS
-ok  	worker-pool	11.570s
+ok  	github.com/vardius/worker-pool/v2	9.590s
 ```
 
 ## 🏫 Basic example
@@ -58,7 +59,7 @@ import (
     "fmt"
     "sync"
 
-    "github.com/vardius/worker-pool"
+    "github.com/vardius/worker-pool/v2"
 )
 
 func main() {
@@ -71,11 +72,16 @@ func main() {
 	// create new pool
 	pool := workerpool.New(poolSize)
 	out := make(chan int, jobsAmount)
-
-	pool.Start(workersAmount, func(i int) {
-		defer wg.Done()
-		out <- i
-	})
+	worker := func(i int) {
+        defer wg.Done()
+        out <- i
+    }
+
+	for i := 1; i <= workersAmount; i++ {
+		if err := pool.AddWorker(worker); err != nil {
+			panic(err)
+		}
+	}
 
 	wg.Add(jobsAmount)
 

benchmark_test.go

@@ -26,7 +26,13 @@ func runBenchmark(b *testing.B, workersAmount int, runInParallel bool) {
 	pool := New(b.N)
 	defer pool.Stop()
 
-	pool.Start(workersAmount, func(i int, out chan<- int) { out <- i })
+	worker := func(i int, out chan<- int) { out <- i }
+
+	for i := 1; i <= workersAmount; i++ {
+		if err := pool.AddWorker(worker); err != nil {
+			b.Fatal(err)
+		}
+	}
 
 	go func() {
 		if runInParallel {

example_test.go

@@ -5,7 +5,7 @@ import (
 	"strconv"
 	"sync"
 
-	workerpool "github.com/vardius/worker-pool"
+	workerpool "github.com/vardius/worker-pool/v2"
 )
 
 func Example() {
@@ -18,11 +18,16 @@ func Example() {
 	// create new pool
 	pool := workerpool.New(poolSize)
 	out := make(chan int, jobsAmount)
-
-	pool.Start(workersAmount, func(i int) {
+	worker := func(i int) {
 		defer wg.Done()
 		out <- i
-	})
+	}
+
+	for i := 1; i <= workersAmount; i++ {
+		if err := pool.AddWorker(worker); err != nil {
+			panic(err)
+		}
+	}
 
 	wg.Add(jobsAmount)
 
@@ -59,7 +64,13 @@ func Example_second() {
 	pool := workerpool.New(poolSize)
 	defer pool.Stop()
 
-	pool.Start(workersAmount, func(i int, out chan<- int) { out <- i })
+	worker := func(i int, out chan<- int) { out <- i }
+
+	for i := 1; i <= workersAmount; i++ {
+		if err := pool.AddWorker(worker); err != nil {
+			panic(err)
+		}
+	}
 
 	go func() {
 		for n := 0; n < jobsAmount; n++ {
@@ -104,12 +115,18 @@ func Example_third() {
 		}(i)
 	}
 
-	// start worker
-	pool.Start(workersAmount, func(s string) {
+	worker := func(s string) {
 		defer wg.Done()
 		defer fmt.Println("job " + s + " is done !")
 		fmt.Println("job " + s + " is running ..")
-	})
+	}
+
+	// start workers
+	for i := 1; i <= workersAmount; i++ {
+		if err := pool.AddWorker(worker); err != nil {
+			panic(err)
+		}
+	}
 
 	// clean up
 	wg.Wait()

go.mod

@@ -1,3 +1,3 @@
-module github.com/vardius/worker-pool
+module github.com/vardius/worker-pool/v2
 
-go 1.12
+go 1.13

pool.go

@@ -3,6 +3,7 @@ package workerpool
 import (
 	"fmt"
 	"reflect"
+	"sync"
 )
 
 // Pool implements worker pool
@@ -11,52 +12,90 @@ type Pool interface {
 	// will block if channel is full, you might want to wrap it with goroutine to avoid it
 	// will panic if called after Stop()
 	Delegate(args ...interface{})
-	// Start given number of workers that will take jobs from a queue
-	Start(maxWorkers int, fn interface{}) error
+
+	// AddWorker adds worker to the pool
+	AddWorker(fn interface{}) error
+	// RemoveWorker removes worker from the pool
+	RemoveWorker(fn interface{}) error
+
+	// WorkersNum returns number of workers in the pool
+	WorkersNum() int
 	// Stop all workers
 	Stop()
 }
 
 type pool struct {
 	queue         chan []reflect.Value
 	isQueueClosed bool
+	workers       []reflect.Value
+	mtx           sync.RWMutex
 }
 
 func (p *pool) Delegate(args ...interface{}) {
 	p.queue <- buildQueueValue(args)
 }
 
-func (p *pool) Start(maxWorkers int, fn interface{}) error {
-	if maxWorkers < 1 {
-		return fmt.Errorf("Invalid number of workers: %d", maxWorkers)
+func (p *pool) AddWorker(fn interface{}) error {
+	if err := isValidHandler(fn); err != nil {
+		return err
 	}
 
-	if reflect.TypeOf(fn).Kind() != reflect.Func {
-		return fmt.Errorf("%s is not a reflect.Func", reflect.TypeOf(fn))
+	if p.isQueueClosed {
+		return fmt.Errorf("can not add new worker to already stopped pool")
 	}
 
-	if p.isQueueClosed {
-		return fmt.Errorf("Can not start already stopped worker")
+	worker := reflect.ValueOf(fn)
+
+	go func() {
+		for args := range p.queue {
+			worker.Call(args)
+		}
+	}()
+
+	p.mtx.Lock()
+	defer p.mtx.Unlock()
+
+	p.workers = append(p.workers, worker)
+
+	return nil
+}
+
+func (p *pool) RemoveWorker(fn interface{}) error {
+	if err := isValidHandler(fn); err != nil {
+		return err
 	}
 
-	task := reflect.ValueOf(fn)
+	rv := reflect.ValueOf(fn)
+
+	p.mtx.Lock()
+	defer p.mtx.Unlock()
 
-	for i := 1; i <= maxWorkers; i++ {
-		go func() {
-			for args := range p.queue {
-				task.Call(args)
-			}
-		}()
+	for i, worker := range p.workers {
+		if worker == rv {
+			p.workers = append(p.workers[:i], p.workers[i+1:]...)
+		}
 	}
 
 	return nil
 }
 
+func (p *pool) WorkersNum() int {
+	return len(p.workers)
+}
+
 func (p *pool) Stop() {
 	close(p.queue)
 	p.isQueueClosed = true
 }
 
+func isValidHandler(fn interface{}) error {
+	if reflect.TypeOf(fn).Kind() != reflect.Func {
+		return fmt.Errorf("%s is not a reflect.Func", reflect.TypeOf(fn))
+	}
+
+	return nil
+}
+
 func buildQueueValue(args []interface{}) []reflect.Value {
 	reflectedArgs := make([]reflect.Value, 0)
 
@@ -70,6 +109,7 @@ func buildQueueValue(args []interface{}) []reflect.Value {
 // New creates new worker pool with a given job queue length
 func New(queueLength int) Pool {
 	return &pool{
-		queue: make(chan []reflect.Value, queueLength),
+		queue:   make(chan []reflect.Value, queueLength),
+		workers: make([]reflect.Value, 0),
 	}
 }

pool_test.go

@@ -14,20 +14,12 @@ func TestNew(t *testing.T) {
 		t.Fail()
 	}
 }
-func TestInvalidWorkersNumber(t *testing.T) {
-	pool := New(2)
-	defer pool.Stop()
-
-	if pool.Start(0, func() {}) == nil {
-		t.Fail()
-	}
-}
 
 func TestInvalidWorker(t *testing.T) {
 	pool := New(2)
 	defer pool.Stop()
 
-	if pool.Start(1, "worker") == nil {
+	if pool.AddWorker("worker") == nil {
 		t.Fail()
 	}
 }
@@ -43,11 +35,16 @@ func delegateWorkToWorkers(t *testing.T, poolSize int, jobsAmount int, workersAm
 
 	pool := New(poolSize)
 	out := make(chan int, jobsAmount)
-
-	pool.Start(workersAmount, func(i int) {
+	worker := func(i int) {
 		defer wg.Done()
 		out <- i
-	})
+	}
+
+	for i := 1; i <= workersAmount; i++ {
+		if err := pool.AddWorker(worker); err != nil {
+			t.Fatal(err)
+		}
+	}
 
 	wg.Add(jobsAmount)