Improvements

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