graceful shutdown & reload

graceful.go

package main

import (
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"sync"
	"time"
)

type gracefull interface {
	ReloadedChan() chan bool
	DoneChan() chan bool
}

type reloadableService struct {
	name           string
	reloadedChan   chan bool
	shutdownedChan chan bool
}

func (r reloadableService) String() string {
	return r.name
}

func (r *reloadableService) reload() {
	time.Sleep(time.Duration(rand.Intn(100)) * time.Microsecond)
}

func (r *reloadableService) shutdownLoop(shutdownChan chan bool) {
	<-shutdownChan
	time.Sleep(time.Duration(rand.Intn(100)) * time.Microsecond)
	fmt.Printf("%s service shutdowned\n", r)
	r.doneChan <- true
}

func (r *reloadableService) reloadLoop(reloadChanFunc func() chan bool) {
	fmt.Printf("%s service reload loop started\n", r)
	reloadChan := reloadChanFunc()
	for {
		<-reloadChan
		fmt.Printf("%s service reload started\n", r)
		r.reload()
		fmt.Printf("%s service reloaded\n", r)
		reloadChan = reloadChanFunc()
		r.reloadedChan <- true
	}
}

func (r *reloadableService) DoneChan() chan bool {
	return r.doneChan
}

func (r *reloadableService) ReloadedChan() chan bool {
	return r.reloadedChan
}

func NewSevice(name string, shutdownChan chan bool, reloadChanFunc func() chan bool) *reloadableService {
	r := &reloadableService{
		name:         name,
		reloadedChan: make(chan bool),
		doneChan:     make(chan bool),
	}
	go r.reloadLoop(reloadChanFunc)
	go r.shutdownLoop(shutdownChan)
	return r
}

func checkChans(chans []chan bool) bool {
	allDone := true
	idx := 0
	for {
		if idx == len(chans) {
			idx = 0
			if allDone {
				return true
			}
			allDone = true
		}
		timeout := time.After(50 * time.Second)
		select {
		case <-chans[idx]:
			chans[idx] = nil
		case <-timeout:
			return false
		default:
			allDone = allDone
		}
		allDone = allDone && chans[idx] == nil
		idx++
	}
}

func reload(reloadChanFunc func() chan bool, services ...gracefull) {
	fmt.Printf("reload raised\n")
	reloadChan = reloadChanFunc()
	close(reloadChan)
	reloadChan = nil

	chans := []chan bool{}
	for _, s := range services {
		chans = append(chans, s.ReloadedChan())
	}
	if ok := checkChans(chans); ok {
		fmt.Printf("all services reloaded\n")
	} else {
		fmt.Printf("service reload timeout\n")
	}
}

func shutdown(shutdownChan chan bool, services ...gracefull) {
	fmt.Printf("shutdown raised\n")
	close(shutdownChan)

	chans := []chan bool{}
	for _, s := range services {
		chans = append(chans, s.DoneChan())
	}
	if ok := checkChans(chans); ok {
		fmt.Printf("all services shutdowned\n")
	} else {
		fmt.Printf("service shutdown timeout\n")
	}
}

var (
	reloadChan     chan bool
	reloadChanLock *sync.Mutex
)

func main() {
	reloadChanLock := &sync.Mutex{}

	shutdownChan := make(chan bool)

	f := func() chan bool {
		reloadChanLock.Lock()
		defer reloadChanLock.Unlock()
		if reloadChan == nil {
			reloadChan = make(chan bool)
		}
		return reloadChan
	}

	services := []gracefull{
		NewSevice("A", shutdownChan, f),
		NewSevice("B", shutdownChan, f),
		NewSevice("C", shutdownChan, f),
	}

	time.Sleep(time.Duration(rand.Intn(100)) * time.Microsecond)

	reload(f, services...)
	reload(f, services...)
	reload(f, services...)

	time.Sleep(time.Duration(rand.Intn(100)) * time.Microsecond)

	stopChan := make(chan os.Signal, 1)
	signal.Notify(stopChan, os.Interrupt)
	<-stopChan

	shutdown(shutdownChan, services...)
}