shutdown.go 809 B

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  1. package shutdown
  2. import (
  3. "os"
  4. "os/signal"
  5. "syscall"
  6. )
  7. var _ Hook = (*hook)(nil)
  8. // Hook a graceful shutdown hook, default with signals of SIGINT and SIGTERM
  9. type Hook interface {
  10. // WithSignals add more signals into hook
  11. WithSignals(signals ...syscall.Signal) Hook
  12. // Close register shutdown handles
  13. Close(funcs ...func())
  14. }
  15. type hook struct {
  16. ctx chan os.Signal
  17. }
  18. // NewHook create a Hook instance
  19. func NewHook() Hook {
  20. hook := &hook{
  21. ctx: make(chan os.Signal, 1),
  22. }
  23. return hook.WithSignals(syscall.SIGINT, syscall.SIGTERM)
  24. }
  25. func (h *hook) WithSignals(signals ...syscall.Signal) Hook {
  26. for _, s := range signals {
  27. signal.Notify(h.ctx, s)
  28. }
  29. return h
  30. }
  31. func (h *hook) Close(funcs ...func()) {
  32. select {
  33. case <-h.ctx:
  34. }
  35. signal.Stop(h.ctx)
  36. for _, f := range funcs {
  37. f()
  38. }
  39. }