Go channel based Simple Worker pattern

go_channel_simple_worker_pattern.go

package main

import (
	"fmt"
	"log"
	"sync"
)

/*
This is an example of a worker pattern for channels
*/

type WorkItem struct {
	Num int
}

type WorkManager struct {
	workToDo chan WorkItem
	wgWork   sync.WaitGroup // waitgroup for work items
}

func NewWorkManager(workBufferSize int) *WorkManager {
	return &WorkManager{
		workToDo: make(chan WorkItem, workBufferSize),
		wgWork:   sync.WaitGroup{},
	}
}

func (w *WorkManager) Startup(workerCount int) {
	for i := 0; i < workerCount; i++ {
		go w.newWorker(i)
	}
}

// newWorker is the implementation of the the work you need to do
// get some work to do of the channel and process it, then mark the work wg done
func (w *WorkManager) newWorker(id int) {
	log.Printf("Worker %d started", id)
	for workItem := range w.workToDo {
		log.Printf("%v\n", workItem)
		w.wgWork.Done() // mark work item as done
	}
	log.Printf("Worker %d finished", id)
}

func (w *WorkManager) AddWork(work WorkItem) {
	w.wgWork.Add(1)
	w.workToDo <- work // this will block if the channel becomes full as the workers are not working fast enough
}

func (w *WorkManager) WaitForWorkComplete() {
	fmt.Println("Waiting for work to complete")
	w.wgWork.Wait() // wait for any work to complete
	fmt.Println("Work is complete")
}

func (w *WorkManager) Shutdown() {
	w.WaitForWorkComplete()
	close(w.workToDo) // close the work channel, this will cause each worker go routine to quit out the range exit nicely
}

func main() {
	wm := NewWorkManager(100)
	wm.Startup(10)

	log.Println("create work items block 1")
	for i := 0; i < 10; i++ {
		wm.AddWork(WorkItem{Num: i})
	}
	wm.WaitForWorkComplete()

	log.Println("create work items block 2")
	for i := 0; i < 10; i++ {
		wm.AddWork(WorkItem{Num: 1000 + i})
	}
	wm.Shutdown()
}