1 files changed, 76 insertions, 64 deletions
diff --git a/modules/graceful/manager.go b/modules/graceful/manager.go
index b9a56ca9c6..eec675e297 100644
--- a/modules/graceful/manager.go
+++ b/modules/graceful/manager.go
@@ -6,9 +6,9 @@ package graceful
 
 import (
 	"context"
+	"sync"
 	"time"
 
-	"code.gitea.io/gitea/modules/git"
 	"code.gitea.io/gitea/modules/log"
 	"code.gitea.io/gitea/modules/process"
 	"code.gitea.io/gitea/modules/setting"
@@ -34,14 +34,24 @@ const (
 const numberOfServersToCreate = 3
 
 // Manager represents the graceful server manager interface
-var Manager *gracefulManager
-
-func init() {
-	Manager = newGracefulManager(context.Background())
-	// Set the git default context to the HammerContext
-	git.DefaultContext = Manager.HammerContext()
-	// Set the process default context to the HammerContext
-	process.DefaultContext = Manager.HammerContext()
+var manager *Manager
+
+var initOnce = sync.Once{}
+
+// GetManager returns the Manager
+func GetManager() *Manager {
+	InitManager(context.Background())
+	return manager
+}
+
+// InitManager creates the graceful manager in the provided context
+func InitManager(ctx context.Context) {
+	initOnce.Do(func() {
+		manager = newGracefulManager(ctx)
+
+		// Set the process default context to the HammerContext
+		process.DefaultContext = manager.HammerContext()
+	})
 }
 
 // CallbackWithContext is combined runnable and context to watch to see if the caller has finished
@@ -61,7 +71,7 @@ type RunnableWithShutdownFns func(atShutdown, atTerminate func(context.Context,
 // Please note that use of the atShutdown and atTerminate callbacks will create go-routines that will wait till their respective signals
 // - users must therefore be careful to only call these as necessary.
 // If run is not expected to run indefinitely RunWithShutdownChan is likely to be more appropriate.
-func (g *gracefulManager) RunWithShutdownFns(run RunnableWithShutdownFns) {
+func (g *Manager) RunWithShutdownFns(run RunnableWithShutdownFns) {
 	g.runningServerWaitGroup.Add(1)
 	defer g.runningServerWaitGroup.Done()
 	run(func(ctx context.Context, atShutdown func()) {
@@ -90,7 +100,7 @@ type RunnableWithShutdownChan func(atShutdown <-chan struct{}, atTerminate Callb
 // (Optionally IsHammer may be waited for instead however, this should be avoided if possible.)
 // The callback function provided to atTerminate must return once termination is complete.
 // Please note that use of the atTerminate function will create a go-routine that will wait till terminate - users must therefore be careful to only call this as necessary.
-func (g *gracefulManager) RunWithShutdownChan(run RunnableWithShutdownChan) {
+func (g *Manager) RunWithShutdownChan(run RunnableWithShutdownChan) {
 	g.runningServerWaitGroup.Add(1)
 	defer g.runningServerWaitGroup.Done()
 	run(g.IsShutdown(), func(ctx context.Context, atTerminate func()) {
@@ -101,14 +111,14 @@ func (g *gracefulManager) RunWithShutdownChan(run RunnableWithShutdownChan) {
 // RunWithShutdownContext takes a function that has a context to watch for shutdown.
 // After the provided context is Done(), the main function must return once shutdown is complete.
 // (Optionally the HammerContext may be obtained and waited for however, this should be avoided if possible.)
-func (g *gracefulManager) RunWithShutdownContext(run func(context.Context)) {
+func (g *Manager) RunWithShutdownContext(run func(context.Context)) {
 	g.runningServerWaitGroup.Add(1)
 	defer g.runningServerWaitGroup.Done()
 	run(g.ShutdownContext())
 }
 
 // RunAtTerminate adds to the terminate wait group and creates a go-routine to run the provided function at termination
-func (g *gracefulManager) RunAtTerminate(ctx context.Context, terminate func()) {
+func (g *Manager) RunAtTerminate(ctx context.Context, terminate func()) {
 	g.terminateWaitGroup.Add(1)
 	go func() {
 		select {
@@ -121,7 +131,7 @@ func (g *gracefulManager) RunAtTerminate(ctx context.Context, terminate func())
 }
 
 // RunAtShutdown creates a go-routine to run the provided function at shutdown
-func (g *gracefulManager) RunAtShutdown(ctx context.Context, shutdown func()) {
+func (g *Manager) RunAtShutdown(ctx context.Context, shutdown func()) {
 	go func() {
 		select {
 		case <-g.IsShutdown():
@@ -132,7 +142,7 @@ func (g *gracefulManager) RunAtShutdown(ctx context.Context, shutdown func()) {
 }
 
 // RunAtHammer creates a go-routine to run the provided function at shutdown
-func (g *gracefulManager) RunAtHammer(ctx context.Context, hammer func()) {
+func (g *Manager) RunAtHammer(ctx context.Context, hammer func()) {
 	go func() {
 		select {
 		case <-g.IsHammer():
@@ -141,7 +151,7 @@ func (g *gracefulManager) RunAtHammer(ctx context.Context, hammer func()) {
 		}
 	}()
 }
-func (g *gracefulManager) doShutdown() {
+func (g *Manager) doShutdown() {
 	if !g.setStateTransition(stateRunning, stateShuttingDown) {
 		return
 	}
@@ -158,48 +168,47 @@ func (g *gracefulManager) doShutdown() {
 		g.doHammerTime(0)
 		<-time.After(1 * time.Second)
 		g.doTerminate()
+		g.WaitForTerminate()
+		g.lock.Lock()
+		close(g.done)
+		g.lock.Unlock()
 	}()
 }
 
-func (g *gracefulManager) doHammerTime(d time.Duration) {
+func (g *Manager) doHammerTime(d time.Duration) {
 	time.Sleep(d)
+	g.lock.Lock()
 	select {
 	case <-g.hammer:
 	default:
 		log.Warn("Setting Hammer condition")
 		close(g.hammer)
 	}
-
+	g.lock.Unlock()
 }
 
-func (g *gracefulManager) doTerminate() {
+func (g *Manager) doTerminate() {
 	if !g.setStateTransition(stateShuttingDown, stateTerminate) {
 		return
 	}
 	g.lock.Lock()
-	close(g.terminate)
+	select {
+	case <-g.terminate:
+	default:
+		log.Warn("Terminating")
+		close(g.terminate)
+	}
 	g.lock.Unlock()
 }
 
 // IsChild returns if the current process is a child of previous Gitea process
-func (g *gracefulManager) IsChild() bool {
+func (g *Manager) IsChild() bool {
 	return g.isChild
 }
 
 // IsShutdown returns a channel which will be closed at shutdown.
 // The order of closure is IsShutdown, IsHammer (potentially), IsTerminate
-func (g *gracefulManager) IsShutdown() <-chan struct{} {
-	g.lock.RLock()
-	if g.shutdown == nil {
-		g.lock.RUnlock()
-		g.lock.Lock()
-		if g.shutdown == nil {
-			g.shutdown = make(chan struct{})
-		}
-		defer g.lock.Unlock()
-		return g.shutdown
-	}
-	defer g.lock.RUnlock()
+func (g *Manager) IsShutdown() <-chan struct{} {
 	return g.shutdown
 }
 
@@ -207,65 +216,43 @@ func (g *gracefulManager) IsShutdown() <-chan struct{} {
 // The order of closure is IsShutdown, IsHammer (potentially), IsTerminate
 // Servers running within the running server wait group should respond to IsHammer
 // if not shutdown already
-func (g *gracefulManager) IsHammer() <-chan struct{} {
-	g.lock.RLock()
-	if g.hammer == nil {
-		g.lock.RUnlock()
-		g.lock.Lock()
-		if g.hammer == nil {
-			g.hammer = make(chan struct{})
-		}
-		defer g.lock.Unlock()
-		return g.hammer
-	}
-	defer g.lock.RUnlock()
+func (g *Manager) IsHammer() <-chan struct{} {
 	return g.hammer
 }
 
 // IsTerminate returns a channel which will be closed at terminate
 // The order of closure is IsShutdown, IsHammer (potentially), IsTerminate
 // IsTerminate will only close once all running servers have stopped
-func (g *gracefulManager) IsTerminate() <-chan struct{} {
-	g.lock.RLock()
-	if g.terminate == nil {
-		g.lock.RUnlock()
-		g.lock.Lock()
-		if g.terminate == nil {
-			g.terminate = make(chan struct{})
-		}
-		defer g.lock.Unlock()
-		return g.terminate
-	}
-	defer g.lock.RUnlock()
+func (g *Manager) IsTerminate() <-chan struct{} {
 	return g.terminate
 }
 
 // ServerDone declares a running server done and subtracts one from the
 // running server wait group. Users probably do not want to call this
 // and should use one of the RunWithShutdown* functions
-func (g *gracefulManager) ServerDone() {
+func (g *Manager) ServerDone() {
 	g.runningServerWaitGroup.Done()
 }
 
 // WaitForServers waits for all running servers to finish. Users should probably
 // instead use AtTerminate or IsTerminate
-func (g *gracefulManager) WaitForServers() {
+func (g *Manager) WaitForServers() {
 	g.runningServerWaitGroup.Wait()
 }
 
 // WaitForTerminate waits for all terminating actions to finish.
 // Only the main go-routine should use this
-func (g *gracefulManager) WaitForTerminate() {
+func (g *Manager) WaitForTerminate() {
 	g.terminateWaitGroup.Wait()
 }
 
-func (g *gracefulManager) getState() state {
+func (g *Manager) getState() state {
 	g.lock.RLock()
 	defer g.lock.RUnlock()
 	return g.state
 }
 
-func (g *gracefulManager) setStateTransition(old, new state) bool {
+func (g *Manager) setStateTransition(old, new state) bool {
 	if old != g.getState() {
 		return false
 	}
@@ -279,7 +266,7 @@ func (g *gracefulManager) setStateTransition(old, new state) bool {
 	return true
 }
 
-func (g *gracefulManager) setState(st state) {
+func (g *Manager) setState(st state) {
 	g.lock.Lock()
 	defer g.lock.Unlock()
 
@@ -288,6 +275,31 @@ func (g *gracefulManager) setState(st state) {
 
 // InformCleanup tells the cleanup wait group that we have either taken a listener
 // or will not be taking a listener
-func (g *gracefulManager) InformCleanup() {
+func (g *Manager) InformCleanup() {
 	g.createServerWaitGroup.Done()
 }
+
+// Done allows the manager to be viewed as a context.Context, it returns a channel that is closed when the server is finished terminating
+func (g *Manager) Done() <-chan struct{} {
+	return g.done
+}
+
+// Err allows the manager to be viewed as a context.Context done at Terminate, it returns ErrTerminate
+func (g *Manager) Err() error {
+	select {
+	case <-g.Done():
+		return ErrTerminate
+	default:
+		return nil
+	}
+}
+
+// Value allows the manager to be viewed as a context.Context done at Terminate, it has no values
+func (g *Manager) Value(key interface{}) interface{} {
+	return nil
+}
+
+// Deadline returns nil as there is no fixed Deadline for the manager, it allows the manager to be viewed as a context.Context
+func (g *Manager) Deadline() (deadline time.Time, ok bool) {
+	return
+}