Golang Worker Pool

wpool.go

package wpool

import (
	"sync"
)

// New returns new WorkerPool instance
func New(numOfWorkers int) *WorkerPool {
	wp := &WorkerPool{
		jobs:         make(chan func()),
		numOfWorkers: numOfWorkers,
		wg:           sync.WaitGroup{},
	}

	for i := 0; i < wp.numOfWorkers; i++ {
		go func() {
			for job := range wp.jobs {
				job()
				wp.wg.Done()
			}
		}()
	}

	return wp
}

// WorkerPool is a struct for creating workers
type WorkerPool struct {
	jobs         chan func()
	numOfWorkers int
	wg           sync.WaitGroup
}

// Add adds new job to worker pool
func (wp *WorkerPool) Add(job func()) {
	wp.wg.Add(1)
	wp.jobs <- job
}

// Wait waits until all jobs are done
func (wp *WorkerPool) Wait() {
	wp.wg.Wait()
	close(wp.jobs)
}

wpool_test.go

package wpool

import (
	"sync"
	"testing"
)

func TestWpoolWorkers(t *testing.T) {
	wp := New(10)
	mux := sync.Mutex{}
	p := 0

	for i := 0; i < 1000; i++ {
		wp.Add(func() {
			mux.Lock()
			p++
			mux.Unlock()
		})
	}

	wp.Wait()
	if p != 1000 {
		t.Fatalf("expected %d, got %d", 1000, p)
	}
}