2019-11-21 21:32:02 +03:00
// Copyright 2019 The Gitea Authors. All rights reserved.
2022-11-27 21:20:29 +03:00
// SPDX-License-Identifier: MIT
2019-11-21 21:32:02 +03:00
package graceful
import (
2019-11-30 17:40:22 +03:00
"context"
2022-03-25 15:47:12 +03:00
"runtime/pprof"
2019-12-15 12:51:28 +03:00
"sync"
2019-11-21 21:32:02 +03:00
"time"
"code.gitea.io/gitea/modules/log"
2019-11-30 17:40:22 +03:00
"code.gitea.io/gitea/modules/process"
2019-11-21 21:32:02 +03:00
"code.gitea.io/gitea/modules/setting"
)
type state uint8
const (
stateInit state = iota
stateRunning
stateShuttingDown
stateTerminate
)
2022-08-21 09:50:27 +03:00
// There are some places that could inherit sockets:
2019-11-21 21:32:02 +03:00
//
// * HTTP or HTTPS main listener
2022-08-21 09:50:27 +03:00
// * HTTP or HTTPS install listener
2019-11-21 21:32:02 +03:00
// * HTTP redirection fallback
2022-08-21 09:50:27 +03:00
// * Builtin SSH listener
2019-11-21 21:32:02 +03:00
//
// If you add an additional place you must increment this number
// and add a function to call manager.InformCleanup if it's not going to be used
2020-10-20 00:03:08 +03:00
const numberOfServersToCreate = 4
2019-11-21 21:32:02 +03:00
// Manager represents the graceful server manager interface
2019-12-15 12:51:28 +03:00
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 ( )
} )
2019-11-21 21:32:02 +03:00
}
2021-05-15 17:22:26 +03:00
// WithCallback is a runnable to call when the caller has finished
type WithCallback func ( callback func ( ) )
2019-11-30 17:40:22 +03:00
// RunnableWithShutdownFns is a runnable with functions to run at shutdown and terminate
// After the callback to atShutdown is called and is complete, the main function must return.
// 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.
2021-05-15 17:22:26 +03:00
type RunnableWithShutdownFns func ( atShutdown , atTerminate func ( func ( ) ) )
2019-11-30 17:40:22 +03:00
// RunWithShutdownFns takes a function that has both atShutdown and atTerminate callbacks
// After the callback to atShutdown is called and is complete, the main function must return.
// 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.
2019-12-15 12:51:28 +03:00
func ( g * Manager ) RunWithShutdownFns ( run RunnableWithShutdownFns ) {
2019-11-30 17:40:22 +03:00
g . runningServerWaitGroup . Add ( 1 )
defer g . runningServerWaitGroup . Done ( )
2020-05-15 03:06:00 +03:00
defer func ( ) {
if err := recover ( ) ; err != nil {
log . Critical ( "PANIC during RunWithShutdownFns: %v\nStacktrace: %s" , err , log . Stack ( 2 ) )
g . doShutdown ( )
}
} ( )
2021-05-15 17:22:26 +03:00
run ( func ( atShutdown func ( ) ) {
g . lock . Lock ( )
defer g . lock . Unlock ( )
g . toRunAtShutdown = append ( g . toRunAtShutdown ,
func ( ) {
defer func ( ) {
if err := recover ( ) ; err != nil {
log . Critical ( "PANIC during RunWithShutdownFns: %v\nStacktrace: %s" , err , log . Stack ( 2 ) )
g . doShutdown ( )
}
} ( )
2019-11-30 17:40:22 +03:00
atShutdown ( )
2021-05-15 17:22:26 +03:00
} )
} , func ( atTerminate func ( ) ) {
g . RunAtTerminate ( atTerminate )
2019-11-30 17:40:22 +03:00
} )
}
// 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.)
2019-12-15 12:51:28 +03:00
func ( g * Manager ) RunWithShutdownContext ( run func ( context . Context ) ) {
2019-11-30 17:40:22 +03:00
g . runningServerWaitGroup . Add ( 1 )
defer g . runningServerWaitGroup . Done ( )
2020-05-15 03:06:00 +03:00
defer func ( ) {
if err := recover ( ) ; err != nil {
log . Critical ( "PANIC during RunWithShutdownContext: %v\nStacktrace: %s" , err , log . Stack ( 2 ) )
g . doShutdown ( )
}
} ( )
2022-03-25 15:47:12 +03:00
ctx := g . ShutdownContext ( )
pprof . SetGoroutineLabels ( ctx ) // We don't have a label to restore back to but I think this is fine
run ( ctx )
2019-11-30 17:40:22 +03:00
}
// RunAtTerminate adds to the terminate wait group and creates a go-routine to run the provided function at termination
2021-05-15 17:22:26 +03:00
func ( g * Manager ) RunAtTerminate ( terminate func ( ) ) {
2019-11-30 17:40:22 +03:00
g . terminateWaitGroup . Add ( 1 )
2021-05-15 17:22:26 +03:00
g . lock . Lock ( )
defer g . lock . Unlock ( )
g . toRunAtTerminate = append ( g . toRunAtTerminate ,
func ( ) {
defer g . terminateWaitGroup . Done ( )
defer func ( ) {
if err := recover ( ) ; err != nil {
log . Critical ( "PANIC during RunAtTerminate: %v\nStacktrace: %s" , err , log . Stack ( 2 ) )
}
} ( )
2019-11-30 17:40:22 +03:00
terminate ( )
2021-05-15 17:22:26 +03:00
} )
2019-11-30 17:40:22 +03:00
}
// RunAtShutdown creates a go-routine to run the provided function at shutdown
2019-12-15 12:51:28 +03:00
func ( g * Manager ) RunAtShutdown ( ctx context . Context , shutdown func ( ) ) {
2021-05-15 17:22:26 +03:00
g . lock . Lock ( )
defer g . lock . Unlock ( )
g . toRunAtShutdown = append ( g . toRunAtShutdown ,
func ( ) {
defer func ( ) {
if err := recover ( ) ; err != nil {
log . Critical ( "PANIC during RunAtShutdown: %v\nStacktrace: %s" , err , log . Stack ( 2 ) )
}
} ( )
select {
case <- ctx . Done ( ) :
return
default :
shutdown ( )
2020-05-15 03:06:00 +03:00
}
2021-05-15 17:22:26 +03:00
} )
2019-11-30 17:40:22 +03:00
}
// RunAtHammer creates a go-routine to run the provided function at shutdown
2021-05-15 17:22:26 +03:00
func ( g * Manager ) RunAtHammer ( hammer func ( ) ) {
g . lock . Lock ( )
defer g . lock . Unlock ( )
g . toRunAtHammer = append ( g . toRunAtHammer ,
func ( ) {
defer func ( ) {
if err := recover ( ) ; err != nil {
log . Critical ( "PANIC during RunAtHammer: %v\nStacktrace: %s" , err , log . Stack ( 2 ) )
}
} ( )
2019-11-30 17:40:22 +03:00
hammer ( )
2021-05-15 17:22:26 +03:00
} )
2019-11-30 17:40:22 +03:00
}
2022-01-20 20:46:10 +03:00
2019-12-15 12:51:28 +03:00
func ( g * Manager ) doShutdown ( ) {
2019-11-21 21:32:02 +03:00
if ! g . setStateTransition ( stateRunning , stateShuttingDown ) {
2022-02-19 19:36:25 +03:00
g . DoImmediateHammer ( )
2019-11-21 21:32:02 +03:00
return
}
g . lock . Lock ( )
2021-05-15 17:22:26 +03:00
g . shutdownCtxCancel ( )
2022-03-25 15:47:12 +03:00
atShutdownCtx := pprof . WithLabels ( g . hammerCtx , pprof . Labels ( "graceful-lifecycle" , "post-shutdown" ) )
pprof . SetGoroutineLabels ( atShutdownCtx )
2021-05-15 17:22:26 +03:00
for _ , fn := range g . toRunAtShutdown {
go fn ( )
}
2019-11-21 21:32:02 +03:00
g . lock . Unlock ( )
if setting . GracefulHammerTime >= 0 {
go g . doHammerTime ( setting . GracefulHammerTime )
}
go func ( ) {
g . WaitForServers ( )
2019-11-30 17:40:22 +03:00
// Mop up any remaining unclosed events.
g . doHammerTime ( 0 )
2019-11-21 21:32:02 +03:00
<- time . After ( 1 * time . Second )
g . doTerminate ( )
2019-12-15 12:51:28 +03:00
g . WaitForTerminate ( )
g . lock . Lock ( )
2022-03-25 15:47:12 +03:00
g . managerCtxCancel ( )
2019-12-15 12:51:28 +03:00
g . lock . Unlock ( )
2019-11-21 21:32:02 +03:00
} ( )
}
2019-12-15 12:51:28 +03:00
func ( g * Manager ) doHammerTime ( d time . Duration ) {
2019-11-21 21:32:02 +03:00
time . Sleep ( d )
2019-12-15 12:51:28 +03:00
g . lock . Lock ( )
2019-11-21 21:32:02 +03:00
select {
2021-05-15 17:22:26 +03:00
case <- g . hammerCtx . Done ( ) :
2019-11-21 21:32:02 +03:00
default :
log . Warn ( "Setting Hammer condition" )
2021-05-15 17:22:26 +03:00
g . hammerCtxCancel ( )
2022-03-25 15:47:12 +03:00
atHammerCtx := pprof . WithLabels ( g . terminateCtx , pprof . Labels ( "graceful-lifecycle" , "post-hammer" ) )
pprof . SetGoroutineLabels ( atHammerCtx )
2021-05-15 17:22:26 +03:00
for _ , fn := range g . toRunAtHammer {
go fn ( )
}
2019-11-21 21:32:02 +03:00
}
2019-12-15 12:51:28 +03:00
g . lock . Unlock ( )
2019-11-21 21:32:02 +03:00
}
2019-12-15 12:51:28 +03:00
func ( g * Manager ) doTerminate ( ) {
2019-11-21 21:32:02 +03:00
if ! g . setStateTransition ( stateShuttingDown , stateTerminate ) {
return
}
g . lock . Lock ( )
2019-12-15 12:51:28 +03:00
select {
2021-05-15 17:22:26 +03:00
case <- g . terminateCtx . Done ( ) :
2019-12-15 12:51:28 +03:00
default :
log . Warn ( "Terminating" )
2021-05-15 17:22:26 +03:00
g . terminateCtxCancel ( )
2022-03-25 15:47:12 +03:00
atTerminateCtx := pprof . WithLabels ( g . managerCtx , pprof . Labels ( "graceful-lifecycle" , "post-terminate" ) )
pprof . SetGoroutineLabels ( atTerminateCtx )
2021-05-15 17:22:26 +03:00
for _ , fn := range g . toRunAtTerminate {
go fn ( )
}
2019-12-15 12:51:28 +03:00
}
2019-11-21 21:32:02 +03:00
g . lock . Unlock ( )
}
// IsChild returns if the current process is a child of previous Gitea process
2019-12-15 12:51:28 +03:00
func ( g * Manager ) IsChild ( ) bool {
2019-11-21 21:32:02 +03:00
return g . isChild
}
// IsShutdown returns a channel which will be closed at shutdown.
// The order of closure is IsShutdown, IsHammer (potentially), IsTerminate
2019-12-15 12:51:28 +03:00
func ( g * Manager ) IsShutdown ( ) <- chan struct { } {
2021-05-15 17:22:26 +03:00
return g . shutdownCtx . Done ( )
2019-11-21 21:32:02 +03:00
}
// IsHammer returns a channel which will be closed at hammer
// 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
2019-12-15 12:51:28 +03:00
func ( g * Manager ) IsHammer ( ) <- chan struct { } {
2021-05-15 17:22:26 +03:00
return g . hammerCtx . Done ( )
2019-11-21 21:32:02 +03:00
}
// 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
2019-12-15 12:51:28 +03:00
func ( g * Manager ) IsTerminate ( ) <- chan struct { } {
2021-05-15 17:22:26 +03:00
return g . terminateCtx . Done ( )
2019-11-21 21:32:02 +03:00
}
// 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
2019-12-15 12:51:28 +03:00
func ( g * Manager ) ServerDone ( ) {
2019-11-21 21:32:02 +03:00
g . runningServerWaitGroup . Done ( )
}
// WaitForServers waits for all running servers to finish. Users should probably
// instead use AtTerminate or IsTerminate
2019-12-15 12:51:28 +03:00
func ( g * Manager ) WaitForServers ( ) {
2019-11-21 21:32:02 +03:00
g . runningServerWaitGroup . Wait ( )
}
// WaitForTerminate waits for all terminating actions to finish.
// Only the main go-routine should use this
2019-12-15 12:51:28 +03:00
func ( g * Manager ) WaitForTerminate ( ) {
2019-11-21 21:32:02 +03:00
g . terminateWaitGroup . Wait ( )
}
2019-12-15 12:51:28 +03:00
func ( g * Manager ) getState ( ) state {
2019-11-21 21:32:02 +03:00
g . lock . RLock ( )
defer g . lock . RUnlock ( )
return g . state
}
2019-12-15 12:51:28 +03:00
func ( g * Manager ) setStateTransition ( old , new state ) bool {
2019-11-21 21:32:02 +03:00
if old != g . getState ( ) {
return false
}
g . lock . Lock ( )
if g . state != old {
g . lock . Unlock ( )
return false
}
g . state = new
g . lock . Unlock ( )
return true
}
2019-12-15 12:51:28 +03:00
func ( g * Manager ) setState ( st state ) {
2019-11-21 21:32:02 +03:00
g . lock . Lock ( )
defer g . lock . Unlock ( )
g . state = st
}
2022-08-21 09:50:27 +03:00
// InformCleanup tells the cleanup wait group that we have either taken a listener or will not be taking a listener.
// At the moment the total number of servers (numberOfServersToCreate) are pre-defined as a const before global init,
// so this function MUST be called if a server is not used.
2019-12-15 12:51:28 +03:00
func ( g * Manager ) InformCleanup ( ) {
2019-11-21 21:32:02 +03:00
g . createServerWaitGroup . Done ( )
}
2019-12-15 12:51:28 +03:00
// 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 { } {
2022-03-25 15:47:12 +03:00
return g . managerCtx . Done ( )
2019-12-15 12:51:28 +03:00
}
2021-05-15 17:22:26 +03:00
// Err allows the manager to be viewed as a context.Context done at Terminate
2019-12-15 12:51:28 +03:00
func ( g * Manager ) Err ( ) error {
2022-03-25 15:47:12 +03:00
return g . managerCtx . Err ( )
2019-12-15 12:51:28 +03:00
}
2021-05-15 17:22:26 +03:00
// Value allows the manager to be viewed as a context.Context done at Terminate
2019-12-15 12:51:28 +03:00
func ( g * Manager ) Value ( key interface { } ) interface { } {
2022-03-25 15:47:12 +03:00
return g . managerCtx . Value ( key )
2019-12-15 12:51:28 +03:00
}
// 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 ) {
2022-03-25 15:47:12 +03:00
return g . managerCtx . Deadline ( )
2019-12-15 12:51:28 +03:00
}