1
0

closer.go 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. package closer
  2. import (
  3. "context"
  4. "os/signal"
  5. "sync"
  6. "syscall"
  7. "time"
  8. )
  9. var closer = &Closer{}
  10. type Closer struct {
  11. initOnce sync.Once
  12. waitOnce sync.Once
  13. mu sync.Mutex
  14. functions []func(ctx context.Context) error
  15. timeout *time.Duration
  16. logger Logger
  17. }
  18. func Init(c Config) error {
  19. closer.initOnce.Do(func() {
  20. closer.timeout = c.timeout
  21. closer.logger = c.logger
  22. })
  23. return nil
  24. }
  25. func Add(f func(ctx context.Context) error) {
  26. closer.mu.Lock()
  27. defer closer.mu.Unlock()
  28. closer.functions = append(closer.functions, f)
  29. }
  30. func Wait(ctx context.Context) {
  31. closer.waitOnce.Do(func() {
  32. ctx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
  33. defer stop()
  34. <-ctx.Done()
  35. // start graceful shutdown
  36. closeCtx := context.Background()
  37. var cancel context.CancelFunc
  38. if closer.timeout != nil {
  39. closeCtx, cancel = context.WithTimeout(closeCtx, *closer.timeout)
  40. defer cancel()
  41. }
  42. complete := make(chan struct{}, 1)
  43. go func() {
  44. for _, f := range closer.functions {
  45. if err := f(closeCtx); err != nil {
  46. if closer.logger != nil {
  47. closer.logger.Errorf(closeCtx, "closer: %s", err.Error())
  48. }
  49. }
  50. }
  51. complete <- struct{}{}
  52. }()
  53. select {
  54. case <-complete:
  55. if closer.logger != nil {
  56. closer.logger.Info(closeCtx, "closer: all processes successfully completed")
  57. }
  58. break
  59. case <-closeCtx.Done():
  60. if closer.logger != nil {
  61. closer.logger.Info(closeCtx, "closer: completed by timeout")
  62. }
  63. }
  64. })
  65. }