Allow to add same worker multiple times, delegate returns now error

Author: vardius

README.md

@@ -41,14 +41,14 @@ For **GoDoc** reference, **visit [pkg.go.dev](https://pkg.go.dev/github.com/vard
 goos: darwin
 goarch: amd64
 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
+BenchmarkWorker1-4                	 3944299	       284 ns/op	      56 B/op	       3 allocs/op
+BenchmarkWorker1Parallel-4        	 7394715	       138 ns/op	      48 B/op	       2 allocs/op
+BenchmarkWorker100-4              	 1657569	       693 ns/op	      56 B/op	       3 allocs/op
+BenchmarkWorker100Parallel-4      	 3673483	       368 ns/op	      48 B/op	       2 allocs/op
+BenchmarkWorkerNumCPU-4           	 2590293	       445 ns/op	      56 B/op	       3 allocs/op
+BenchmarkWorkerNumCPUParallel-4   	 3591553	       298 ns/op	      48 B/op	       2 allocs/op
 PASS
-ok  	github.com/vardius/worker-pool/v2	9.590s
+ok  	github.com/vardius/worker-pool/v2	9.511s
 ```
 
 ## 🏫 Basic example
@@ -93,7 +93,7 @@ func main() {
 		// stop all workers after jobs are done
 		wg.Wait()
 		close(out)
-		pool.Stop()
+		pool.Stop() // stop removes all workers from pool, to resume work add them again
 	}()
 
 	sum := 0

benchmark_test.go

@@ -5,49 +5,38 @@ import (
 	"testing"
 )
 
-func poolDelegate(b *testing.B, pool Pool, out chan<- int) {
+func poolDelegate(b *testing.B, pool Pool) {
 	for n := 0; n < b.N; n++ {
-		pool.Delegate(n, out)
+		pool.Delegate(n)
 	}
 }
 
-func poolDelegateParallel(b *testing.B, pool Pool, out chan<- int) {
+func poolDelegateParallel(b *testing.B, pool Pool) {
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
-			pool.Delegate(1, out)
+			pool.Delegate(1)
 		}
 	})
 }
 
 func runBenchmark(b *testing.B, workersAmount int, runInParallel bool) {
-	ch := make(chan int, b.N)
-	defer close(ch)
-
 	pool := New(b.N)
 	defer pool.Stop()
 
-	worker := func(i int, out chan<- int) { out <- i }
+	worker := func(i int) {}
 
 	for i := 1; i <= workersAmount; i++ {
 		if err := pool.AddWorker(worker); err != nil {
 			b.Fatal(err)
 		}
 	}
 
-	go func() {
-		if runInParallel {
-			poolDelegateParallel(b, pool, ch)
-		} else {
-			poolDelegate(b, pool, ch)
-		}
-	}()
+	b.ResetTimer()
 
-	var i = 0
-	for i < b.N {
-		select {
-		case <-ch:
-			i++
-		}
+	if runInParallel {
+		poolDelegateParallel(b, pool)
+	} else {
+		poolDelegate(b, pool)
 	}
 }
 

pool.go

@@ -11,7 +11,7 @@ type Pool interface {
 	// Delegate job to a workers
 	// 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{})
+	Delegate(args ...interface{}) error
 
 	// AddWorker adds worker to the pool
 	AddWorker(fn interface{}) error
@@ -20,42 +20,59 @@ type Pool interface {
 
 	// WorkersNum returns number of workers in the pool
 	WorkersNum() int
-	// Stop all workers
+
+	// Stop removes all workers workers
+	// to resume work add them again
 	Stop()
 }
 
+type quitCh chan interface{}
+type workers map[reflect.Value][]quitCh
+
 type pool struct {
-	queue         chan []reflect.Value
-	isQueueClosed bool
-	workers       []reflect.Value
-	mtx           sync.RWMutex
+	queue   chan []reflect.Value
+	workers workers
+	mtx     sync.RWMutex
 }
 
-func (p *pool) Delegate(args ...interface{}) {
+func (p *pool) Delegate(args ...interface{}) error {
+	if len(p.workers) == 0 {
+		return fmt.Errorf("there is no workers in pool")
+	}
+
 	p.queue <- buildQueueValue(args)
+
+	return nil
 }
 
 func (p *pool) AddWorker(fn interface{}) error {
 	if err := isValidHandler(fn); err != nil {
 		return err
 	}
 
-	if p.isQueueClosed {
-		return fmt.Errorf("can not add new worker to already stopped pool")
-	}
-
 	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)
+	q := make(quitCh)
+
+	if _, ok := p.workers[worker]; !ok {
+		p.workers[worker] = []quitCh{q}
+	} else {
+		p.workers[worker] = append(p.workers[worker], q)
+	}
+
+	go func() {
+		for {
+			select {
+			case args := <-p.queue:
+				worker.Call(args)
+			case <-q:
+				return
+			}
+		}
+	}()
 
 	return nil
 }
@@ -65,27 +82,37 @@ func (p *pool) RemoveWorker(fn interface{}) error {
 		return err
 	}
 
-	rv := reflect.ValueOf(fn)
+	worker := reflect.ValueOf(fn)
 
 	p.mtx.Lock()
 	defer p.mtx.Unlock()
 
-	for i, worker := range p.workers {
-		if worker == rv {
-			p.workers = append(p.workers[:i], p.workers[i+1:]...)
-		}
+	if len(p.workers[worker]) > 0 {
+		close(p.workers[worker][len(p.workers[worker])-1])
+
+		p.workers[worker] = p.workers[worker][:len(p.workers[worker])-1]
+	} else {
+		delete(p.workers, worker)
 	}
 
 	return nil
 }
 
 func (p *pool) WorkersNum() int {
-	return len(p.workers)
+	sum := 0
+	for _, qChs := range p.workers {
+		sum += len(qChs)
+	}
+
+	return sum
 }
 
 func (p *pool) Stop() {
-	close(p.queue)
-	p.isQueueClosed = true
+	for _, qChs := range p.workers {
+		for _, ch := range qChs {
+			close(ch)
+		}
+	}
 }
 
 func isValidHandler(fn interface{}) error {
@@ -110,6 +137,6 @@ func buildQueueValue(args []interface{}) []reflect.Value {
 func New(queueLength int) Pool {
 	return &pool{
 		queue:   make(chan []reflect.Value, queueLength),
-		workers: make([]reflect.Value, 0),
+		workers: make(workers),
 	}
 }

pool_test.go

@@ -15,6 +15,31 @@ func TestNew(t *testing.T) {
 	}
 }
 
+func TestRemoveWorker(t *testing.T) {
+	pool := New(2)
+	defer pool.Stop()
+
+	worker := func(i int) {}
+
+	for i := 1; i <= 2; i++ {
+		if err := pool.AddWorker(worker); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	if err := pool.RemoveWorker(worker); err != nil {
+		t.Fatal(err)
+	}
+
+	if pool.WorkersNum() != 1 {
+		t.Fatal("should have one worker left")
+	}
+
+	if err := pool.Delegate(1); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestInvalidWorker(t *testing.T) {
 	pool := New(2)
 	defer pool.Stop()