golang - graceful shutdown sample

graceful-shutdown-sample.go

package main

import (
	"time"
	"syscall"
	"sync"
	"fmt"
	"os"
	"os/signal"
	"context"
)

func main() {

	go initSelfShutdown() // yes indeed just for simulation purpose

	// spawn some workers
	wg := &sync.WaitGroup{}
	for i := 0; i < 20; i++ {
		wg.Add(1)

		go Work(i, wg)

	}
	wg.Wait() // synchronized by a simple WaitGroup
}

func Work(workerNo int, wg *sync.WaitGroup) {

	gracefulShutdown := NewSignalContext().Done()
	for {
		select {
		case <-gracefulShutdown:
			fmt.Printf("Graceful shutdown worker #%v\n", workerNo)
			wg.Done()
			return
		default:
			fmt.Printf("#%v working\n", workerNo)
			time.Sleep(1000 * time.Millisecond)
		}
	}
}

// NewSignalContext creates a context that will cancel for Interrupt
func NewSignalContext() context.Context {
	return NewSignalContextWithSignals(os.Interrupt)
}

// NewSignalContextWithSignals creates a context that will cancel for the given signals
func NewSignalContextWithSignals(signals ...os.Signal) context.Context {

	ctx, cancelFunc := context.WithCancel(context.Background())

	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, signals...)

		<-c
		cancelFunc()
	}()

	return ctx
}

func initSelfShutdown() {
	selfShutdownTimer := time.NewTimer(time.Second * 4)
	<-selfShutdownTimer.C
	syscall.Kill(syscall.Getpid(), syscall.SIGINT)
}