Go (lang) context patterns

context_cancel.go

package main

import (
	"context"
	"flag"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// nextFunc executes next
type nextFunc func(next func())

// processor runs using the given context
type processor interface {
	run(ctx context.Context)
}

// processorFunc encapsulates the behaviour of a processor
//
// it is modeled after http.HandlerFunc
type processorFunc func(context.Context)

// run calls f(ctx)
func (f processorFunc) run(ctx context.Context) {
	f(ctx)
}

// sWorker is an arbitrary func that runs forever
func sWorker(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			log.Println("context done!")
			return
		default:
			log.Println()
			// random sleep for [500, 2500]ms
			time.Sleep(time.Duration(rand.Intn(2001)+500) * time.Millisecond)
		}
	}
}

// pWorker is an sWorker that runs concurrently
func pWorker(wg *sync.WaitGroup) func(context.Context) {
	return func(ctx context.Context) {
		defer wg.Done()
		sWorker(ctx)
	}
}

// worker is a single worker
var worker = sWorker

// workers returns a func that will run n workers
func workers(n int) func(context.Context) {
	return func(ctx context.Context) {
		wg := sync.WaitGroup{}
		for ; n > 0; n-- {
			wg.Add(1)
			go pWorker(&wg)(ctx)
		}
		// defer the wait to ensures all pWorkers complete
		defer wg.Wait()
		// wait on context
		<-ctx.Done()
	}
}

func workerFactory(n int) func(context.Context) {
	// one
	if n <= 1 {
		return worker
	}
	// many
	return workers(n)
}

var n = flag.Int("n", 1, "number of workers to spawn")

func init() {
	log.SetFlags(log.Flags() | log.Lmicroseconds)

	log.Println("init")
	if !flag.Parsed() {
		flag.Parse()
	}
}

func main() {
	ctx := context.Background()
	// setup context with values etc.

	// processor
	proc := processorFunc(workerFactory(*n))
	// wait on interrupt signal
	onSignal(ctx, proc, func(next func()) {
		log.Println("on signal func")
		log.Println("calling next")
		next()
	})
}

// signal names map
var signalM = map[os.Signal]string{
	syscall.SIGINT:  "SIGINT",
	syscall.SIGKILL: "SIGKILL",
	syscall.SIGQUIT: "SIGQUIT",
	syscall.SIGTERM: "SIGTERM",
}

// signal slice
var signalS = []os.Signal{
	syscall.SIGINT,
	syscall.SIGKILL,
	syscall.SIGQUIT,
	syscall.SIGTERM,
}

// waitOnSignal waits for the process to receive a signal before calling
// the process Canceler
func onSignal(ctx context.Context, proc processor, onSignalFunc nextFunc) {
	// make context cancellable
	ctx, cancel := context.WithCancel(ctx)
	// set interrupt listerner
	interrupt := make(chan os.Signal)
	signal.Notify(interrupt, signalS...)
	// exec interrupt listener in go-routine
	go func() {
		sig := <-interrupt
		log.Printf("caught signal: %v %v\n", sig, signalM[sig])
		log.Println("calling onSignalFunc with cancel")
		onSignalFunc(cancel)
	}()
	// run our processor
	proc.run(ctx)
}