diff --git a/shutdown.go b/shutdown.go index 2d02836..af4130d 100644 --- a/shutdown.go +++ b/shutdown.go @@ -15,7 +15,7 @@ import ( // log.Println("Shutting down") // }) type Shutdowner struct { - initialized bool + initializer sync.Once shutdown bool quit chan os.Signal funcs []ShutdownFunc @@ -29,49 +29,42 @@ type ShutdownFunc func() // Internal method func (f ShutdownFunc) execute(c context.Context) { - done := false - // Execute ShutdownFunc in goroutine and set done = true + done := make(chan interface{}) + // Execute ShutdownFunc in goroutine and close channel when finished go func() { f() - done = true + close(done) }() - for { - // Check if context canceled or ShutdownFunc finished - select { - case <-c.Done(): - // Context canceled, return - return - default: - if done { - // ShutdownFunc finished, return - return - } - // Otherwise continue - } + // Check if context canceled or ShutdownFunc finished + select { + case <-c.Done(): + // Context canceled, return + return + case <-done: + // ShutdownFunc finished, return + return } } // Internal method func (s *Shutdowner) init() { - // Check if already initialized - if s.initialized { - return - } - // Initialize cancel context and signal channel - s.cancelContext, s.cancelFunc = context.WithCancel(context.Background()) - s.quit = make(chan os.Signal, 2) - signal.Notify(s.quit, - os.Interrupt, - syscall.SIGINT, - syscall.SIGTERM, // e.g. Docker stop - ) - go func() { - for range s.quit { - s.Shutdown() - } - }() - // Finished - s.initialized = true + s.initializer.Do(func() { + s.mutex.Lock() + defer s.mutex.Unlock() + // Initialize cancel context and signal channel + s.cancelContext, s.cancelFunc = context.WithCancel(context.Background()) + s.quit = make(chan os.Signal, 2) + signal.Notify(s.quit, + os.Interrupt, + syscall.SIGINT, + syscall.SIGTERM, // e.g. Docker stop + ) + go func() { + for range s.quit { + s.Shutdown() + } + }() + }) } // Add a func, that should be called with s.Shutdown() or when receiving a shutdown signal