github.com/gogf/gf@v1.16.9/os/gproc/gproc_signal.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/gogf/gf.
     6  
     7  package gproc
     8  
     9  import (
    10  	"context"
    11  	"github.com/gogf/gf/internal/intlog"
    12  	"os"
    13  	"os/signal"
    14  	"sync"
    15  	"syscall"
    16  )
    17  
    18  // SigHandler defines a function type for signal handling.
    19  type SigHandler func(sig os.Signal)
    20  
    21  var (
    22  	signalHandlerMap  = make(map[os.Signal][]SigHandler)
    23  	shutdownSignalMap = map[os.Signal]struct{}{
    24  		syscall.SIGINT:  {},
    25  		syscall.SIGQUIT: {},
    26  		syscall.SIGKILL: {},
    27  		syscall.SIGTERM: {},
    28  		syscall.SIGABRT: {},
    29  	}
    30  )
    31  
    32  func init() {
    33  	for sig, _ := range shutdownSignalMap {
    34  		signalHandlerMap[sig] = make([]SigHandler, 0)
    35  	}
    36  }
    37  
    38  // AddSigHandler adds custom signal handler for custom one or more signals.
    39  func AddSigHandler(handler SigHandler, signals ...os.Signal) {
    40  	for _, sig := range signals {
    41  		signalHandlerMap[sig] = append(signalHandlerMap[sig], handler)
    42  	}
    43  }
    44  
    45  // AddSigHandlerShutdown adds custom signal handler for shutdown signals:
    46  // syscall.SIGINT,
    47  // syscall.SIGQUIT,
    48  // syscall.SIGKILL,
    49  // syscall.SIGTERM,
    50  // syscall.SIGABRT.
    51  func AddSigHandlerShutdown(handler ...SigHandler) {
    52  	for _, h := range handler {
    53  		for sig, _ := range shutdownSignalMap {
    54  			signalHandlerMap[sig] = append(signalHandlerMap[sig], h)
    55  		}
    56  	}
    57  }
    58  
    59  // Listen blocks and does signal listening and handling.
    60  func Listen() {
    61  	signals := make([]os.Signal, 0)
    62  	for sig, _ := range signalHandlerMap {
    63  		signals = append(signals, sig)
    64  	}
    65  	sigChan := make(chan os.Signal, 1)
    66  	signal.Notify(sigChan, signals...)
    67  	var sig os.Signal
    68  	for {
    69  		wg := sync.WaitGroup{}
    70  		sig = <-sigChan
    71  		intlog.Printf(context.TODO(), `signal received: %s`, sig.String())
    72  		if handlers, ok := signalHandlerMap[sig]; ok {
    73  			for _, handler := range handlers {
    74  				wg.Add(1)
    75  				go func(handler SigHandler, sig os.Signal) {
    76  					defer wg.Done()
    77  					handler(sig)
    78  				}(handler, sig)
    79  			}
    80  		}
    81  		// If it is shutdown signal, it exits this signal listening.
    82  		if _, ok := shutdownSignalMap[sig]; ok {
    83  			// Wait until signal handlers done.
    84  			wg.Wait()
    85  			return
    86  		}
    87  	}
    88  }