sonet/pkg/utils/shutdown/signal.go

package shutdown

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

var (
	AwaitSeconds  = 3
	shutdownHooks [3][]func() // front,middle,back hooks
	lock          = &sync.Mutex{}
	sigChan       = make(chan os.Signal, 1)
)

func AddHook(hook func(), options ...Option) {
	opts := defaultOptions
	for _, opt := range options {
		opt(&opts)
	}

	lock.Lock()
	defer lock.Unlock()
	shutdownHooks[opts.Order] = append(shutdownHooks[opts.Order], hook)
}

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

	awaitSeconds := AwaitSeconds
	var success, fail int32
	for i, hooks := range shutdownHooks {
		for _, hook := range hooks {
			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()
			}()
		}

		// 间隔1s再执行
		if i < 2 && len(hooks) > 0 {
			time.Sleep(time.Second)
			awaitSeconds -= 1
		}
	}

	if awaitSeconds < 1 {
		awaitSeconds = 1
	}

	logger.Infof("execute %d shutdown hook %d ok, %d failed, exit in %d seconds...\n", success+fail, success, fail, awaitSeconds)
	time.Sleep(time.Second * time.Duration(awaitSeconds))
}

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