kit/core/proc/shutdown.go

package proc

import (
	"git.diulo.com/mogfee/kit/core/threading"
	"git.diulo.com/mogfee/kit/logx"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

const (
	wrapUpTime = time.Second
	waitTime   = 5500 * time.Millisecond
)

var (
	wrapUpListeners          = new(listenerManager)
	shutdownListeners        = new(listenerManager)
	delayTimeBeforeForceQuit = waitTime
)

func AddShutdownListener(fn func()) (waitForCalled func()) {
	return shutdownListeners.addListener(fn)
}
func AddWrapUpListener(fn func()) (waitForCalled func()) {
	return wrapUpListeners.addListener(fn)
}
func SetTimeToForceQuit(duration time.Duration) {
	delayTimeBeforeForceQuit = duration
}
func Shutdown() {
	shutdownListeners.notifyListeners()
}
func WrapUp() {
	wrapUpListeners.notifyListeners()
}
func gracefulStop(signals chan os.Signal) {
	signal.Stop(signals)
	logx.Info("Got signal SIGTERM, shutting down...")
	time.Sleep(wrapUpTime)
	go shutdownListeners.notifyListeners()
	time.Sleep(delayTimeBeforeForceQuit - wrapUpTime)
	logx.Infof("Still alive after %v, going to force kill the process...", delayTimeBeforeForceQuit)
	syscall.Kill(syscall.Getpid(), syscall.SIGTERM)
}

type listenerManager struct {
	lock      sync.Mutex
	waitGroup sync.WaitGroup
	listeners []func()
}

func (l *listenerManager) addListener(fn func()) (waitForCalled func()) {
	l.waitGroup.Add(1)

	l.lock.Lock()
	l.listeners = append(l.listeners, func() {
		defer l.waitGroup.Done()
		fn()
	})
	l.lock.Unlock()
	return func() {
		l.waitGroup.Wait()
	}
}
func (l *listenerManager) notifyListeners() {
	l.lock.Lock()
	defer l.lock.Unlock()

	group := threading.NewRoutineGroup()
	for _, listener := range l.listeners {
		group.RunSafe(listener)
	}
	group.Wait()
}