diff options
Diffstat (limited to 'modules/graceful/manager.go')
| -rw-r--r-- | modules/graceful/manager.go | 50 |
1 files changed, 16 insertions, 34 deletions
diff --git a/modules/graceful/manager.go b/modules/graceful/manager.go index f783573607..8766cfca0e 100644 --- a/modules/graceful/manager.go +++ b/modules/graceful/manager.go @@ -6,6 +6,7 @@ package graceful import ( "context" + "runtime/pprof" "sync" "time" @@ -62,7 +63,6 @@ type WithCallback func(callback func()) // Similarly the callback function provided to atTerminate must return once termination is complete. // 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. type RunnableWithShutdownFns func(atShutdown, atTerminate func(func())) // RunWithShutdownFns takes a function that has both atShutdown and atTerminate callbacks @@ -70,7 +70,6 @@ type RunnableWithShutdownFns func(atShutdown, atTerminate func(func())) // Similarly the callback function provided to atTerminate must return once termination is complete. // 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 *Manager) RunWithShutdownFns(run RunnableWithShutdownFns) { g.runningServerWaitGroup.Add(1) defer g.runningServerWaitGroup.Done() @@ -98,32 +97,6 @@ func (g *Manager) RunWithShutdownFns(run RunnableWithShutdownFns) { }) } -// RunnableWithShutdownChan is a runnable with functions to run at shutdown and terminate. -// After the atShutdown channel is closed, the main function must return once shutdown is complete. -// (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. -type RunnableWithShutdownChan func(atShutdown <-chan struct{}, atTerminate WithCallback) - -// RunWithShutdownChan takes a function that has channel to watch for shutdown and atTerminate callbacks -// After the atShutdown channel is closed, the main function must return once shutdown is complete. -// (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 *Manager) RunWithShutdownChan(run RunnableWithShutdownChan) { - g.runningServerWaitGroup.Add(1) - defer g.runningServerWaitGroup.Done() - defer func() { - if err := recover(); err != nil { - log.Critical("PANIC during RunWithShutdownChan: %v\nStacktrace: %s", err, log.Stack(2)) - g.doShutdown() - } - }() - run(g.IsShutdown(), func(atTerminate func()) { - g.RunAtTerminate(atTerminate) - }) -} - // 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.) @@ -136,7 +109,9 @@ func (g *Manager) RunWithShutdownContext(run func(context.Context)) { g.doShutdown() } }() - run(g.ShutdownContext()) + ctx := g.ShutdownContext() + pprof.SetGoroutineLabels(ctx) // We don't have a label to restore back to but I think this is fine + run(ctx) } // RunAtTerminate adds to the terminate wait group and creates a go-routine to run the provided function at termination @@ -198,6 +173,8 @@ func (g *Manager) doShutdown() { } g.lock.Lock() g.shutdownCtxCancel() + atShutdownCtx := pprof.WithLabels(g.hammerCtx, pprof.Labels("graceful-lifecycle", "post-shutdown")) + pprof.SetGoroutineLabels(atShutdownCtx) for _, fn := range g.toRunAtShutdown { go fn() } @@ -214,7 +191,7 @@ func (g *Manager) doShutdown() { g.doTerminate() g.WaitForTerminate() g.lock.Lock() - g.doneCtxCancel() + g.managerCtxCancel() g.lock.Unlock() }() } @@ -227,6 +204,8 @@ func (g *Manager) doHammerTime(d time.Duration) { default: log.Warn("Setting Hammer condition") g.hammerCtxCancel() + atHammerCtx := pprof.WithLabels(g.terminateCtx, pprof.Labels("graceful-lifecycle", "post-hammer")) + pprof.SetGoroutineLabels(atHammerCtx) for _, fn := range g.toRunAtHammer { go fn() } @@ -244,6 +223,9 @@ func (g *Manager) doTerminate() { default: log.Warn("Terminating") g.terminateCtxCancel() + atTerminateCtx := pprof.WithLabels(g.managerCtx, pprof.Labels("graceful-lifecycle", "post-terminate")) + pprof.SetGoroutineLabels(atTerminateCtx) + for _, fn := range g.toRunAtTerminate { go fn() } @@ -331,20 +313,20 @@ func (g *Manager) InformCleanup() { // 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.doneCtx.Done() + return g.managerCtx.Done() } // Err allows the manager to be viewed as a context.Context done at Terminate func (g *Manager) Err() error { - return g.doneCtx.Err() + return g.managerCtx.Err() } // Value allows the manager to be viewed as a context.Context done at Terminate func (g *Manager) Value(key interface{}) interface{} { - return g.doneCtx.Value(key) + return g.managerCtx.Value(key) } // 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 g.doneCtx.Deadline() + return g.managerCtx.Deadline() } |