initial import

core/proc/shutdown.go

@@ -0,0 +1,81 @@
+// +build linux darwin
+
+package proc
+
+import (
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+	"time"
+
+	"zero/core/logx"
+)
+
+const (
+	wrapUpTime = time.Second
+	// why we use 5500 milliseconds is because most of our queue are blocking mode with 5 seconds
+	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 SetTimeoutToForceQuit(duration time.Duration) {
+	delayTimeBeforeForceQuit = duration
+}
+
+func gracefulStop(signals chan os.Signal) {
+	signal.Stop(signals)
+
+	logx.Info("Got signal SIGTERM, shutting down...")
+	wrapUpListeners.notifyListeners()
+
+	time.Sleep(wrapUpTime)
+	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 (lm *listenerManager) addListener(fn func()) (waitForCalled func()) {
+	lm.waitGroup.Add(1)
+
+	lm.lock.Lock()
+	lm.listeners = append(lm.listeners, func() {
+		defer lm.waitGroup.Done()
+		fn()
+	})
+	lm.lock.Unlock()
+
+	return func() {
+		lm.waitGroup.Wait()
+	}
+}
+
+func (lm *listenerManager) notifyListeners() {
+	lm.lock.Lock()
+	defer lm.lock.Unlock()
+
+	for _, listener := range lm.listeners {
+		listener()
+	}
+}