sonet/pkg/utils/shutdown/signal.go

package shutdown

import (
	"log"
	"os"
	"os/signal"
	"sonet/pkg/utils/logger"
	"sync"
	"sync/atomic"
	"syscall"
)

var (
	shutdownHooks []func()
	lock          = &sync.Mutex{}
	sigChan       = make(chan os.Signal, 1)
)

func AddShutdownHook(hook func()) {
	lock.Lock()
	defer lock.Unlock()
	shutdownHooks = append(shutdownHooks, hook)
}

func Await() {
	// 监听两个信号: TERM信号(kill + 进程号)触发, 中断信号(ctrl + c)触发
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
	s := <-sigChan
	// 监听到关闭信号
	logger.Info("catch exit signal: ", s)

	var success, fail int32
	for _, hook := range shutdownHooks {
		func() {
			defer func() {
				if err := recover(); err != nil {
					log.Println("exec shutdown hook panic: ", err)
					atomic.AddInt32(&fail, 1)
					return
				}
				atomic.AddInt32(&success, 1)
			}()

			hook()
		}()
	}

	logger.Infof("execute %d shutdown hook %d ok, %d failed\n", success+fail, success, fail)
}

//func Shutdown() {
//	sigChan <- syscall.SIGQUIT
//}