github.com/lingyao2333/mo-zero@v1.4.1/core/proc/shutdown.go (about)

     1  //go:build linux || darwin
     2  // +build linux darwin
     3  
     4  package proc
     5  
     6  import (
     7  	"os"
     8  	"os/signal"
     9  	"sync"
    10  	"syscall"
    11  	"time"
    12  
    13  	"github.com/lingyao2333/mo-zero/core/logx"
    14  	"github.com/lingyao2333/mo-zero/core/threading"
    15  )
    16  
    17  const (
    18  	wrapUpTime = time.Second
    19  	// why we use 5500 milliseconds is because most of our queue are blocking mode with 5 seconds
    20  	waitTime = 5500 * time.Millisecond
    21  )
    22  
    23  var (
    24  	wrapUpListeners          = new(listenerManager)
    25  	shutdownListeners        = new(listenerManager)
    26  	delayTimeBeforeForceQuit = waitTime
    27  )
    28  
    29  // AddShutdownListener adds fn as a shutdown listener.
    30  // The returned func can be used to wait for fn getting called.
    31  func AddShutdownListener(fn func()) (waitForCalled func()) {
    32  	return shutdownListeners.addListener(fn)
    33  }
    34  
    35  // AddWrapUpListener adds fn as a wrap up listener.
    36  // The returned func can be used to wait for fn getting called.
    37  func AddWrapUpListener(fn func()) (waitForCalled func()) {
    38  	return wrapUpListeners.addListener(fn)
    39  }
    40  
    41  // SetTimeToForceQuit sets the waiting time before force quitting.
    42  func SetTimeToForceQuit(duration time.Duration) {
    43  	delayTimeBeforeForceQuit = duration
    44  }
    45  
    46  func gracefulStop(signals chan os.Signal) {
    47  	signal.Stop(signals)
    48  
    49  	logx.Info("Got signal SIGTERM, shutting down...")
    50  	go wrapUpListeners.notifyListeners()
    51  
    52  	time.Sleep(wrapUpTime)
    53  	go shutdownListeners.notifyListeners()
    54  
    55  	time.Sleep(delayTimeBeforeForceQuit - wrapUpTime)
    56  	logx.Infof("Still alive after %v, going to force kill the process...", delayTimeBeforeForceQuit)
    57  	syscall.Kill(syscall.Getpid(), syscall.SIGTERM)
    58  }
    59  
    60  type listenerManager struct {
    61  	lock      sync.Mutex
    62  	waitGroup sync.WaitGroup
    63  	listeners []func()
    64  }
    65  
    66  func (lm *listenerManager) addListener(fn func()) (waitForCalled func()) {
    67  	lm.waitGroup.Add(1)
    68  
    69  	lm.lock.Lock()
    70  	lm.listeners = append(lm.listeners, func() {
    71  		defer lm.waitGroup.Done()
    72  		fn()
    73  	})
    74  	lm.lock.Unlock()
    75  
    76  	return func() {
    77  		lm.waitGroup.Wait()
    78  	}
    79  }
    80  
    81  func (lm *listenerManager) notifyListeners() {
    82  	lm.lock.Lock()
    83  	defer lm.lock.Unlock()
    84  
    85  	group := threading.NewRoutineGroup()
    86  	for _, listener := range lm.listeners {
    87  		group.RunSafe(listener)
    88  	}
    89  	group.Wait()
    90  }