Better work pool pattern

Author: ardan-bkennedy

07-concurrency_patterns/work/main/main.go

@@ -1,14 +1,13 @@
 // All material is licensed under the GNU Free Documentation License
 // https://github.com/ArdanStudios/gotraining/blob/master/LICENSE
 
-// This example is provided with help by Jason Waldrip.
-
 // This sample program demostrates how to use the work package
 // to use a pool of goroutines to get work done.
 package main
 
 import (
 	"fmt"
+	"log"
 	"sync"
 	"time"
 
@@ -30,21 +29,26 @@ type namePrinter struct {
 }
 
 // Work implements the Worker interface.
-func (m *namePrinter) Work() {
+func (m *namePrinter) Work(id int) {
 	fmt.Println(m.name)
 	time.Sleep(time.Second)
 }
 
+// logger is called by the work pool when
+// it has things to log.
+func logger(message string) {
+	log.Println(message)
+}
+
 // main is the entry point for all Go programs.
 func main() {
 	// Create a work value with 2 goroutines.
-	w := work.New(2)
-	w.LogStats(100 * time.Millisecond)
+	w, _ := work.New(2, time.Second, logger)
 
 	var wg sync.WaitGroup
-	wg.Add(10 * len(names))
+	wg.Add(100 * len(names))
 
-	for i := 0; i < 10; i++ {
+	for i := 0; i < 100; i++ {
 		// Iterate over the slice of names.
 		for _, name := range names {
 			// Create a namePrinter and provide the
@@ -63,6 +67,8 @@ func main() {
 	}
 
 	for {
+		// Enter a number and hit enter to change the size
+		// of the work pool.
 		var c int
 		fmt.Scanf("%d", &c)
 		if c == 0 {

07-concurrency_patterns/work/work.go

@@ -1,151 +1,216 @@
 // All material is licensed under the GNU Free Documentation License
 // https://github.com/ArdanStudios/gotraining/blob/master/LICENSE
 
-// Package work manages a pool of goroutines to perform work.
+// Package work manages a pool of routines to perform work.
 package work
 
 import (
+	"errors"
 	"fmt"
 	"sync"
 	"sync/atomic"
 	"time"
 )
 
+const (
+	addRoutine = 1
+	rmvRoutine = 2
+)
+
+// ErrorInvalidMinRoutines is the error for the invalid minRoutine parameter.
+var ErrorInvalidMinRoutines = errors.New("Invalid minimum number of routines")
+
+// ErrorInvalidStatTime is the error for the invalid stat time parameter.
+var ErrorInvalidStatTime = errors.New("Invalid duration for stat time")
+
 // Worker must be implemented by types that want to use
 // this worker processes.
 type Worker interface {
-	Work()
+	Work(id int)
 }
 
-// Stats contains information about the work pool.
-type Stats struct {
-	Goroutines int64
-	Pending    int64
-	Active     int64
-}
-
-// Work provides a pool of goroutines that can execute any Worker
+// Work provides a pool of routines that can execute any Worker
 // tasks that are submitted.
 type Work struct {
-	tasks    chan Worker    // Unbuffered channel that work is sent into.
-	kill     chan struct{}  // Unbuffered channel to signal for a goroutine to die.
-	shutdown chan struct{}  // Closed when the Work pool is being shutdown.
-	wg       sync.WaitGroup // Manages the number of goroutines for shutdown.
-	mutex    sync.Mutex     // Provides synchronization for managing the pool.
-	stats    Stats          // Stats about the health of the pool.
+	minRoutines int            // Minumum number of routines always in the pool.
+	statTime    time.Duration  // Time to display stats.
+	counter     int            // Maintains a running total number of routines ever created.
+	tasks       chan Worker    // Unbuffered channel that work is sent into.
+	control     chan int       // Unbuffered channel that work for the manager is send into.
+	kill        chan struct{}  // Unbuffered channel to signal for a goroutine to die.
+	shutdown    chan struct{}  // Closed when the Work pool is being shutdown.
+	wg          sync.WaitGroup // Manages the number of routines for shutdown.
+	routines    int64          // Number of routines
+	active      int64          // Active number of routines in the work pool.
+	pending     int64          // Pending number of routines waiting to submit work.
+
+	logFunc func(message string) // Function called to providing logging support
 }
 
 // New creates a new Worker.
-func New(goroutines int) *Work {
-	w := Work{
-		tasks:    make(chan Worker),
-		kill:     make(chan struct{}),
-		shutdown: make(chan struct{}),
+func New(minRoutines int, statTime time.Duration, logFunc func(message string)) (*Work, error) {
+	if minRoutines <= 0 {
+		return nil, ErrorInvalidMinRoutines
 	}
 
-	w.Add(goroutines)
-
-	return &w
-}
-
-// LogStats display work pool stats on the specified duration.
-func (w *Work) LogStats(d time.Duration) {
-	w.wg.Add(1)
+	if statTime < time.Millisecond {
+		return nil, ErrorInvalidStatTime
+	}
 
-	go func() {
-		for {
-			select {
-			case <-w.shutdown:
-				w.wg.Done()
-				return
+	w := Work{
+		minRoutines: minRoutines,
+		statTime:    statTime,
+		tasks:       make(chan Worker),
+		control:     make(chan int),
+		kill:        make(chan struct{}),
+		shutdown:    make(chan struct{}),
+		logFunc:     logFunc,
+	}
 
-			case <-time.After(d):
-				s := w.Stats()
-				fmt.Printf("G[%d] P[%d] A[%d]\n", s.Goroutines, s.Pending, s.Active)
-			}
-		}
-	}()
-}
+	// Start the manager.
+	w.manager()
 
-// Stats returns the current status for the work pool.
-func (w *Work) Stats() Stats {
-	var s Stats
-	s.Goroutines = atomic.LoadInt64(&w.stats.Goroutines)
-	s.Pending = atomic.LoadInt64(&w.stats.Pending)
-	s.Active = atomic.LoadInt64(&w.stats.Active)
+	// Add the routines.
+	w.Add(minRoutines)
 
-	return s
+	return &w, nil
 }
 
-// Add creates goroutines to process work or sets a count for
-// goroutines to terminate.
-func (w *Work) Add(goroutines int) {
-	if goroutines == 0 {
+// Add creates routines to process work or sets a count for
+// routines to terminate.
+func (w *Work) Add(routines int) {
+	if routines == 0 {
 		return
 	}
 
-	w.mutex.Lock()
-	{
-		if goroutines > 0 {
-			// We are adding goroutines to the pool.
-			w.wg.Add(goroutines)
-			atomic.AddInt64(&w.stats.Goroutines, int64(goroutines))
-
-			for i := 0; i < goroutines; i++ {
-				go w.work()
-			}
-		} else {
-			// We are removing goroutines from the pool.
-			goroutines = goroutines * -1
-			current := int(atomic.LoadInt64(&w.stats.Goroutines))
-			if goroutines > current {
-				goroutines = current
-			}
+	cmd := addRoutine
+	if routines < 0 {
+		routines = routines * -1
+		cmd = rmvRoutine
+	}
 
-			// Send the kill signal and wait for these goroutines
-			// to get the signal to die.
-			for i := 0; i < goroutines; i++ {
-				w.kill <- struct{}{}
-			}
-		}
+	for i := 0; i < routines; i++ {
+		w.control <- cmd
 	}
-	w.mutex.Unlock()
 }
 
 // work performs the users work and keeps stats.
-func (w *Work) work() {
+func (w *Work) work(id int) {
 done:
 	for {
 		select {
-		case t, ok := <-w.tasks:
-			if !ok {
-				break done
+		case t := <-w.tasks:
+			atomic.AddInt64(&w.active, 1)
+			{
+				// Perform the work.
+				t.Work(id)
 			}
-
-			atomic.AddInt64(&w.stats.Active, 1)
-			t.Work()
-			atomic.AddInt64(&w.stats.Active, -1)
+			atomic.AddInt64(&w.active, -1)
 
 		case <-w.kill:
 			break done
 		}
 	}
 
-	atomic.AddInt64(&w.stats.Goroutines, -1)
+	// Decrement the counts.
+	atomic.AddInt64(&w.routines, -1)
 	w.wg.Done()
+
+	w.log("Worker : Shutting Down")
 }
 
 // Run wait for the goroutine pool to take the work
 // to be executed.
 func (w *Work) Run(work Worker) {
-	atomic.AddInt64(&w.stats.Pending, 1)
-	w.tasks <- work
-	atomic.AddInt64(&w.stats.Pending, -1)
+	atomic.AddInt64(&w.pending, 1)
+	{
+		w.tasks <- work
+	}
+	atomic.AddInt64(&w.pending, -1)
 }
 
 // Shutdown waits for all the workers to finish.
 func (w *Work) Shutdown() {
-	close(w.tasks)
 	close(w.shutdown)
 	w.wg.Wait()
 }
+
+// manager controls changes to the work pool including stats
+// and shutting down.
+func (w *Work) manager() {
+	w.wg.Add(1)
+
+	go func() {
+		w.log("Work Manager : Started")
+
+		// Create a timer to run stats.
+		timer := time.NewTimer(w.statTime)
+
+		for {
+			select {
+			case <-w.shutdown:
+				// Capture the current number of routines.
+				routines := int(atomic.LoadInt64(&w.routines))
+
+				// Send a kill to all the existing routines.
+				for i := 0; i < routines; i++ {
+					w.kill <- struct{}{}
+				}
+
+				// Decrement the waitgroup and kill the manager.
+				w.wg.Done()
+				return
+
+			case c := <-w.control:
+				switch c {
+				case addRoutine:
+					w.log("Work Manager : Add Routine")
+
+					// Capture a unique id.
+					w.counter++
+
+					// Add to the counts.
+					w.wg.Add(1)
+					atomic.AddInt64(&w.routines, 1)
+
+					// Create the routine.
+					go w.work(w.counter)
+
+				case rmvRoutine:
+					w.log("Work Manager : Remove Routine")
+
+					// Capture the number of routines.
+					routines := int(atomic.LoadInt64(&w.routines))
+
+					// Are there routines to remove.
+					if routines <= w.minRoutines {
+						w.log("Work Manager : Reached Minimum Can't Remove")
+						break
+					}
+
+					// Send a kill signal to remove a routine.
+					w.kill <- struct{}{}
+				}
+
+			case <-timer.C:
+				// Capture the stats.
+				routines := atomic.LoadInt64(&w.routines)
+				pending := atomic.LoadInt64(&w.pending)
+				active := atomic.LoadInt64(&w.active)
+
+				// Display the stats.
+				w.log(fmt.Sprintf("Work Manager : Stats : G[%d] P[%d] A[%d]", routines, pending, active))
+
+				// Reset the clock.
+				timer.Reset(w.statTime)
+			}
+		}
+	}()
+}
+
+// log sending logging messages back to the client.
+func (w *Work) log(message string) {
+	if w.logFunc != nil {
+		w.logFunc(message)
+	}
+}