workers.go

package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

const numWorkers = 10
const numMessages = 100

func newWorker() *workers {
	return &workers{
		producerShutdown: make(chan bool),
		q:                make(chan int, numWorkers),
	}
}

type workers struct {
	consumerWg       sync.WaitGroup
	producerShutdown chan bool
	producerWg       sync.WaitGroup
	q                chan int
}

func (w *workers) Start() {
	for i := 0; i < numWorkers; i++ {
		w.consumerWg.Add(1)
		go w.work(i + 1)
	}

	go func() {
		w.producerWg.Add(1)
		defer w.producerWg.Done()

		log.Println("starting publisher")

		for i := 0; i < numMessages; i++ {
			select {
			case <-w.producerShutdown:
				return
			default:
				w.q <- i
				log.Printf("Published %d", i+1)
			}
		}
	}()
}

func (w *workers) Stop() {
	close(w.producerShutdown)
	w.producerWg.Wait()

	close(w.q)
	w.consumerWg.Wait()
}

func (w *workers) work(workerNum int) {
	defer w.consumerWg.Done()

	log.Printf("Starting worker %d", workerNum)

	for {
		select {
		case i, ok := <-w.q:
			if !ok {
				log.Printf("Stopping worker %d", workerNum)
				return
			}

			log.Printf("Got: %d", i)
			time.Sleep(time.Second)
		}
	}
}

func main() {
	w := newWorker()
	w.Start()

	// shutdown hooks
	gracefulStop := make(chan os.Signal)
	signal.Notify(gracefulStop, syscall.SIGTERM, syscall.SIGINT)

	go func() {
		sig := <-gracefulStop

		log.Printf("caught signal: %s", sig)

		log.Println("shutting down workers")

		w.Stop()

		os.Exit(0)
	}()

	select {}
}