Merge pull request #1 from gasparian/refactoring

worker pool manager structure refactored

Author: gasparian

README.md

@@ -4,23 +4,27 @@
 
 Implementation of worker pool with extendable functionality.  
 By default, I've implemented "round-robin" balancing to evenly distribute jobs across workers.  
-You need to implement the `Manager` interface in order to use other balancing strategies:  
+You need to implement the `BalancingStrategy` interface in order to use other balancing strategies:  
 ```go
-type Manager interface {
-	ScheduleJob(f JobFunc) (chan Result, error)
+type BalancingStrategy interface {
+	NextWorkerId(workersStats []Stats) int
 }
 ```  
 
 API:  
 ```go
-New(config Config) *WorkerPool
-(wp WorkerPool) GetCurrentJobsNumber() int64
-(wp *WorkerPool) GetWorkerStats(workerId int) (Stats, error)
-(wp *WorkerPool) TerminateWorker(workerId int) error
-(wp *WorkerPool) ReloadWorker(workerId int) error
+// worker pool struct private api
+(wp *workerPool) getCurrentJobsNumber() int64
+(wp *workerPool) getWorkerStats(workerId int) (Stats, error)
+(wp *workerPool) terminateWorker(workerId int) error
+(wp *workerPool) reloadWorker(workerId int) error
 
-NewRoundRobin(pool *WorkerPool) *RoundRobin
-(rr *RoundRobin) ScheduleJob(f JobFunc) (chan Result, error)
+// key public structs and methods
+New(config Config, balancer BalancingStrategy) *Manager
+(m *Manager) ScheduleJob(f JobFunc) (chan Result, error)
+
+NewRoundRobin() *RoundRobin
+(rr *RoundRobin) NextWorkerId(workerStats []Stats) int
 ```  
 
 Install:  
@@ -37,6 +41,7 @@ import (
     wp "github.com/gasparian/worker-pool-go"
 )
 
+// Use closure to create job function
 func exampleJob(inp int) wp.JobFunc {
     return func() wp.Result {
         res := inp * inp
@@ -50,33 +55,25 @@ func exampleJob(inp int) wp.JobFunc {
 func main() {
     config := wp.Config{
         NWorkers: 3,  // Number of workers to spawn
-        MaxJobs:  10, // Max jobs in a queue, new jobs will be rejected with error
+        MaxJobs:  10, // Max jobs in a queue, new jobs will be rejected with an error
     }
-    pool := wp.New(config)
-    roundRobinPool := wp.NewRoundRobin(pool)
+    balancer := wp.NewRoundRobin()
+    pool := wp.New(config, balancer)
 
     nJobs := 50
-    jobs := make([]chan wp.Result, 0)
+    results := make([]chan wp.Result, 0)
     for i := 0; i < nJobs; i++ {
-        ch, err := roundRobinPool.ScheduleJob(exampleJob(i))
+        ch, err := pool.ScheduleJob(exampleJob(i))
         if err != nil {
             break
         }
-        jobs = append(jobs, ch)
+        results = append(results, ch)
     }
 
-    for _, j := range jobs {
-        res := <-j
+    for _, r := range results {
+        res := <-r
         fmt.Println(res)
     }
-
-    for w := 0; w < config.NWorkers; w++ {
-        s, err := roundRobinPool.GetWorkerStats(w)
-        if err != nil {
-            panic(err)
-        }
-        fmt.Printf("Worker %v stats: %v\n", w, s)
-    }
 }
 ```  
 

workerpool.go

@@ -28,13 +28,8 @@ type Job struct {
 
 // Stats holds basic stats for the worker
 type Stats struct {
-	ProcessedJobs    int64
-	TotalElapsedTime int64 // NOTE: nanoseconds
-}
-
-type workerStats struct {
-	Stats
-	mx sync.RWMutex
+	ProcessedJobs          int64
+	TotalJobsExecutionTime int64 // NOTE: nanoseconds
 }
 
 // Config holds needed data to start worker pool
@@ -43,8 +38,14 @@ type Config struct {
 	MaxJobs  int64
 }
 
-// WorkerPool holds data needed for pool operation
-type WorkerPool struct {
+type workerStats struct {
+	Stats
+	mx sync.RWMutex
+}
+
+// workerPool holds data needed for pool operation
+type workerPool struct {
+	mx            sync.RWMutex
 	config        Config
 	workersChan   []chan Job
 	workersStats  []*workerStats
@@ -60,7 +61,7 @@ func worker(jobs chan Job, terminate chan bool, s *workerStats) {
 			elapsedTime := time.Since(start).Nanoseconds()
 			s.mx.Lock()
 			s.ProcessedJobs++
-			s.TotalElapsedTime += elapsedTime
+			s.TotalJobsExecutionTime += elapsedTime
 			s.mx.Unlock()
 			j.ResultChan <- res
 			close(j.ResultChan)
@@ -72,69 +73,58 @@ func worker(jobs chan Job, terminate chan bool, s *workerStats) {
 	}
 }
 
-// New created new worker pool
-func New(config Config) *WorkerPool {
-	pool := &WorkerPool{
-		config:        config,
-		workersChan:   make([]chan Job, config.NWorkers),
-		workersStats:  make([]*workerStats, config.NWorkers),
-		terminateChan: make([]chan bool, config.NWorkers),
-	}
-	for w := 0; w < config.NWorkers; w++ {
-		pool.workersChan[w] = make(chan Job, config.MaxJobs)
-		pool.terminateChan[w] = make(chan bool)
-		pool.workersStats[w] = &workerStats{}
-		go worker(
-			pool.workersChan[w],
-			pool.terminateChan[w],
-			pool.workersStats[w],
-		)
-	}
-	return pool
-}
+// getCurrentJobsAmount calculates amount of jobs across all workers queues
+func (wp *workerPool) getCurrentJobsAmount() int64 {
+	wp.mx.RLock()
+	defer wp.mx.RUnlock()
 
-// GetCurrentJobsNumber calculates amount of jobs across all workers queues
-func (wp WorkerPool) GetCurrentJobsNumber() int64 {
-	var currentJobsN int64 = 0
+	var currentJobsAmount int64 = 0
 	for _, ch := range wp.workersChan {
-		currentJobsN += int64(len(ch))
+		currentJobsAmount += int64(len(ch))
 	}
-	return currentJobsN
+	return currentJobsAmount
 }
 
-// GetWorkerStats  by specified worker id
-func (wp *WorkerPool) GetWorkerStats(workerId int) (Stats, error) {
+// getWorkerStats by specified worker id
+func (wp *workerPool) getWorkerStats(workerId int) (Stats, error) {
 	resStats := Stats{}
 	if workerId > len(wp.workersChan) || workerId < 0 {
 		return resStats, workerIndexIsOutOfBoundsErr
 	}
+	wp.mx.RLock()
 	s := wp.workersStats[workerId]
+	wp.mx.RUnlock()
+
 	s.mx.RLock()
 	defer s.mx.RUnlock()
 	resStats.ProcessedJobs = s.ProcessedJobs
-	resStats.TotalElapsedTime = s.TotalElapsedTime
+	resStats.TotalJobsExecutionTime = s.TotalJobsExecutionTime
 	return resStats, nil
 }
 
-// TerminateWorker sends termination signal to the specified worker
-func (wp *WorkerPool) TerminateWorker(workerId int) error {
+// terminateWorker sends termination signal to the specified worker
+func (wp *workerPool) terminateWorker(workerId int) error {
 	if workerId > len(wp.terminateChan) || workerId < 0 {
 		return workerIndexIsOutOfBoundsErr
 	}
+	wp.mx.RLock()
 	wp.terminateChan[workerId] <- true
+	wp.mx.RUnlock()
 	return nil
 }
 
-// ReloadWorker terminates worker by id, and spawns new one
-func (wp *WorkerPool) ReloadWorker(workerId int) error {
+// reloadWorker terminates worker by id, and spawns new one
+func (wp *workerPool) reloadWorker(workerId int) error {
 	if workerId > len(wp.workersChan) || workerId < 0 {
 		return workerIndexIsOutOfBoundsErr
 	}
-	err := wp.TerminateWorker(workerId)
+	err := wp.terminateWorker(workerId)
 	if err != nil {
 		return err
 	}
+	wp.mx.Lock()
 	wp.workersStats[workerId] = &workerStats{}
+	wp.mx.Unlock()
 	go worker(
 		wp.workersChan[workerId],
 		wp.terminateChan[workerId],
@@ -143,45 +133,87 @@ func (wp *WorkerPool) ReloadWorker(workerId int) error {
 	return nil
 }
 
-// Manafer interface that holds implementation of balancing strategy
-type Manager interface {
-	ScheduleJob(f JobFunc) (chan Result, error)
+// Strategy interface that holds implementation of balancing strategy
+type BalancingStrategy interface {
+	NextWorkerId(workersStats []Stats) int
 }
 
-/*------------------------------------------------------------------------*/
-// NOTE: default implmementation of the Manager
-
-// RoundRobin evenly distribute jobs across workers
-type RoundRobin struct {
-	*WorkerPool
-	mx           sync.RWMutex
-	nextWorkerId int
+// Manager spawns workers and schedule jobs
+type Manager struct {
+	wp       *workerPool
+	balancer BalancingStrategy
 }
 
-// NewRoundRobin ...
-func NewRoundRobin(pool *WorkerPool) *RoundRobin {
-	return &RoundRobin{
-		WorkerPool:   pool,
-		nextWorkerId: 0,
+// New creates new worker pool
+func New(config Config, balancer BalancingStrategy) *Manager {
+	pool := &workerPool{
+		config:        config,
+		workersChan:   make([]chan Job, config.NWorkers),
+		workersStats:  make([]*workerStats, config.NWorkers),
+		terminateChan: make([]chan bool, config.NWorkers),
 	}
+	for w := 0; w < config.NWorkers; w++ {
+		pool.workersChan[w] = make(chan Job, config.MaxJobs)
+		pool.terminateChan[w] = make(chan bool)
+		pool.workersStats[w] = &workerStats{}
+		go worker(
+			pool.workersChan[w],
+			pool.terminateChan[w],
+			pool.workersStats[w],
+		)
+	}
+	manager := &Manager{
+		wp:       pool,
+		balancer: balancer,
+	}
+	// TODO: add routine that reloads workers
+	return manager
 }
 
 // ScheduleJob puts job in a queue
-func (rr *RoundRobin) ScheduleJob(f JobFunc) (chan Result, error) {
-	rr.mx.RLock()
-	config := rr.config
-	nextWorkerId := rr.nextWorkerId
-	rr.mx.RUnlock()
-
-	currentJobsN := rr.GetCurrentJobsNumber()
-	if currentJobsN >= config.MaxJobs {
+func (m *Manager) ScheduleJob(f JobFunc) (chan Result, error) {
+	m.wp.mx.RLock()
+	config := m.wp.config
+	workersStats := make([]Stats, config.NWorkers)
+	for i, s := range m.wp.workersStats {
+		s.mx.RLock()
+		workersStats[i] = s.Stats
+		s.mx.RUnlock()
+	}
+	m.wp.mx.RUnlock()
+
+	currentJobsAmount := m.wp.getCurrentJobsAmount()
+	if currentJobsAmount >= config.MaxJobs {
 		return nil, maxJobsLimitReachedErr
 	}
 	ch := make(chan Result)
-	rr.workersChan[nextWorkerId] <- Job{f, ch}
 
-	rr.mx.Lock()
-	rr.nextWorkerId = (nextWorkerId + 1) % config.NWorkers
-	rr.mx.Unlock()
+	nextWorkerId := m.balancer.NextWorkerId(workersStats)
+	m.wp.workersChan[nextWorkerId] <- Job{f, ch}
 	return ch, nil
 }
+
+/*------------------------------------------------------------------------*/
+
+// NOTE: default balancing strategy
+
+// RoundRobin evenly distributes jobs across workers
+type RoundRobin struct {
+	mx           sync.RWMutex
+	nextWorkerId int
+}
+
+// NewRoundRobin ...
+func NewRoundRobin() *RoundRobin {
+	return &RoundRobin{nextWorkerId: 0}
+}
+
+// NextWorkerId generates worker id to schedule the job
+func (rr *RoundRobin) NextWorkerId(workerStats []Stats) int {
+	rr.mx.Lock()
+	defer rr.mx.Unlock()
+
+	workerId := rr.nextWorkerId
+	rr.nextWorkerId = (workerId + 1) % len(workerStats)
+	return workerId
+}