├── .github ├── sus.png └── workflows │ └── test.yml ├── .gitignore ├── LICENSE ├── README.md ├── examples ├── README.md ├── blog │ ├── 01_done_channel_wait_group │ │ └── main.go │ ├── 02_done_channel_wait_group_force │ │ └── main.go │ ├── 03_done_channel_wait_group_force_timeout │ │ └── main.go │ ├── 04_context_wait_group_force_timeout │ │ └── main.go │ └── 05_context_errgroup_force_timeout │ │ └── main.go └── lib │ └── http_example │ └── main.go ├── go.mod ├── go.sum ├── graceful.go ├── graceful_test.go ├── signal_handler.go ├── signal_handler_test.go ├── timeout_handler.go └── timeout_handler_test.go /.github/sus.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/itzloop/gograce/5a6c6bdaf1217b428ee3ec66c9b01884683639d0/.github/sus.png -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | # This workflow will build a golang project 2 | # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-go 3 | 4 | name: Go Test 5 | 6 | on: 7 | push: 8 | branches: [ "master" ] 9 | pull_request: 10 | branches: [ "master" ] 11 | 12 | jobs: 13 | 14 | build: 15 | runs-on: ubuntu-latest 16 | strategy: 17 | matrix: 18 | go-version: [ '1.21.x' ] 19 | steps: 20 | - uses: actions/checkout@v3 21 | 22 | - name: Set up Go ${{ matrix.go-version }} 23 | uses: actions/setup-go@v4 24 | with: 25 | go-version: ${{ matrix.go-version }} 26 | 27 | - name: Build 28 | run: go build -v ./... 29 | 30 | - name: Test 31 | run: go test -v ./... 32 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | *.out 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Sina Shabani 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Gograce 2 | > gograce let's you run your programs gracefully managing signal handling, cleanup timeouts and force quit for you. 3 | 4 | ![Red Sus](./.github/sus.png) 5 | 6 | ![GitHub Workflow Status (with event)](https://img.shields.io/github/actions/workflow/status/itzloop/gograce/test.yml) 7 | ![GitHub License](https://img.shields.io/github/license/itzloop/gograce) 8 | [![Go Report Card](https://goreportcard.com/badge/github.com/itzloop/gograce)](https://goreportcard.com/report/github.com/itzloop/gograce) 9 | [![PkgGoDev](https://pkg.go.dev/badge/mod/github.com/itzloop/gograce)](https://pkg.go.dev/mod/github.com/itzloop/gograce) 10 | 11 | ## Usage 12 | ```go 13 | package main 14 | 15 | import ( 16 | "context" 17 | "github.com/itzloop/gograce" 18 | "time" 19 | ) 20 | 21 | func main() { 22 | grace := gograce.NewGraceful(gograce.Options{ 23 | // Timeout sets a hard deadline for cleanup phase. If time out is specified, 24 | // gograce will wait for that amount and then terminates forcefully 25 | Timeout: 10 * time.Second, 26 | 27 | // This controls whether or not sending terminate signal twice will forcefully 28 | // terminate the application 29 | NoForceQuit: false, 30 | 31 | // Setting this will limit the number of go-routines running at the same time. 32 | MaxGoRoutines: 0, 33 | 34 | // Setting this will overwrite the default signals 35 | Signals: nil, 36 | }) 37 | 38 | app := App{} 39 | grace.GoWithContext(app.Start) 40 | grace.GoWithContext(app.Close) 41 | grace.FatalWait() 42 | } 43 | 44 | type App struct{} 45 | 46 | func (app *App) Start(ctx context.Context) error { 47 | // run stuff 48 | } 49 | 50 | 51 | func (app *App) Close(ctx context.Context) error { 52 | <-ctx.Done() 53 | // do cleanup 54 | } 55 | ``` 56 | 57 | For more information on how to use it refer to [examples](/examples/README.md) readme. 58 | 59 | ## Testing 60 | 61 | ```bash 62 | $ go test -v ./... 63 | ``` 64 | 65 | ## Contributing 66 | 67 | TODO 68 | -------------------------------------------------------------------------------- /examples/README.md: -------------------------------------------------------------------------------- 1 | # Examples 2 | 3 | ## Blog examples 4 | 5 | These are the examples used in the blog post 6 | 7 | - [Basic Example](./blog/01_done_channel_wait_group/main.go) that i used in the blog post as a base and improved it step 8 | by step. 9 | - [Example with Done Channel and WaitGroup with Force Quit](./blog/02_done_channel_wait_group_force/main.go) 10 | - [Example with Timeouts](./blog/03_done_channel_wait_group_force_timeout/main.go) 11 | - [Example with context.Context instead of Done Channel](./blog/04_context_wait_group_force_timeout/main.go) 12 | - [Example with errgroup.Group instead of WaitGroups](./blog/05_context_errgroup_force_timeout/main.go) 13 | 14 | ## Library Example 15 | 16 | - [HTTP Server](lib/http_example/main.go) 17 | 18 | Creates a simples http.Server with: 19 | - `/long-running-job` that sleeps for 10 seconds to demonstrate all active connections will finish before termination. 20 | - `/hello` that will write `hi` back :). 21 | 22 | ### Running the example 23 | 24 | ```bash 25 | $ go run examples/lib/http_example/main.go 26 | ``` 27 | 28 | You now have a HTTP server that has to endpoints: 29 | 1. `/hello` which simply writes `hi` back. 30 | 2. `log-runing-job` which is simple `time.Sleep` and then writes `done` back. 31 | 32 | Now when you send a HTTP request to `/log-running-job`, and press `ctrl+c` in server, the code will wait for all requests to finish and then terminates. 33 | -------------------------------------------------------------------------------- /examples/blog/01_done_channel_wait_group/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "os" 6 | "os/signal" 7 | "sync" 8 | "syscall" 9 | "time" 10 | ) 11 | 12 | func init() { 13 | log.SetFlags(0) 14 | log.SetOutput(os.Stdout) 15 | } 16 | 17 | func main() { 18 | // create the channel and pass it to application 19 | done := make(chan struct{}) 20 | 21 | // create a sync.WaitGroup and add 1 to it for each go-routine 22 | wg := sync.WaitGroup{} 23 | wg.Add(1) 24 | go func() { 25 | defer wg.Done() 26 | runApplication(done) 27 | }() 28 | 29 | sig := make(chan os.Signal, 1) 30 | signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) 31 | s := <-sig 32 | log.Println("received signal:", s) 33 | close(done) // we close the channel when a signal is received 34 | 35 | // we wait for all the go-routines 36 | wg.Wait() 37 | } 38 | 39 | func runApplication(done <-chan struct{}) { 40 | defer cleanupApplication() 41 | 42 | for { 43 | select { 44 | case _, ok := <-done: 45 | if !ok { 46 | // exit from application when the channel is closed 47 | return 48 | } 49 | default: 50 | log.Println("doing stuff...") 51 | time.Sleep(time.Second) 52 | } 53 | } 54 | } 55 | 56 | func cleanupApplication() { 57 | time.Sleep(5 * time.Second) 58 | log.Println("application done.") 59 | } 60 | -------------------------------------------------------------------------------- /examples/blog/02_done_channel_wait_group_force/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "os" 6 | "os/signal" 7 | "sync" 8 | "syscall" 9 | "time" 10 | ) 11 | 12 | func init() { 13 | log.SetFlags(0) 14 | log.SetOutput(os.Stdout) 15 | } 16 | 17 | func main() { 18 | // start signal handler 19 | done := signalHandler() 20 | 21 | // create a sync.WaitGroup and add 1 to it for each go-routine 22 | wg := sync.WaitGroup{} 23 | wg.Add(1) 24 | go func() { 25 | defer wg.Done() 26 | runApplication(done) 27 | }() 28 | 29 | // we wait for all the go-routines 30 | wg.Wait() 31 | } 32 | 33 | func signalHandler() <-chan struct{} { 34 | // create the channel and pass it to application 35 | done := make(chan struct{}) 36 | go func() { 37 | sig := make(chan os.Signal, 1) 38 | signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) 39 | s := <-sig 40 | log.Printf("received signal '%s', gracefully quitting...", s) 41 | close(done) // we close the channel when a signal is received 42 | s = <-sig // listen again to force quit 43 | log.Printf("received signal '%s', forcefully quitting...", s) 44 | os.Exit(1) // forcefully terminate the program 45 | }() 46 | 47 | return done 48 | } 49 | 50 | func runApplication(done <-chan struct{}) { 51 | defer cleanupApplication() 52 | 53 | for { 54 | select { 55 | case _, ok := <-done: 56 | if !ok { 57 | // exit from application when the channel is closed 58 | return 59 | } 60 | default: 61 | log.Println("doing stuff...") 62 | time.Sleep(time.Second) 63 | } 64 | } 65 | } 66 | 67 | func cleanupApplication() { 68 | time.Sleep(5 * time.Second) 69 | log.Println("application done.") 70 | } 71 | -------------------------------------------------------------------------------- /examples/blog/03_done_channel_wait_group_force_timeout/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "os" 6 | "os/signal" 7 | "sync" 8 | "syscall" 9 | "time" 10 | ) 11 | 12 | func init() { 13 | log.SetFlags(0) 14 | log.SetOutput(os.Stdout) 15 | } 16 | 17 | func main() { 18 | // start signal handler 19 | done := signalHandler() 20 | 21 | // create a sync.WaitGroup and add 1 to it for each go-routine 22 | wg := sync.WaitGroup{} 23 | wg.Add(1) 24 | go func() { 25 | defer wg.Done() 26 | runApplication(done) 27 | }() 28 | 29 | // run the timeoutHandler handler to have a hard time limit on cleanup phase 30 | go timeoutHandler(done, 2*time.Second) 31 | 32 | // we wait for all the go-routines 33 | wg.Wait() 34 | } 35 | 36 | func signalHandler() <-chan struct{} { 37 | // create the channel and pass it to application 38 | done := make(chan struct{}) 39 | go func() { 40 | sig := make(chan os.Signal, 1) 41 | signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) 42 | s := <-sig 43 | log.Printf("received signal '%s', gracefully quitting...", s) 44 | close(done) // we close the channel when a signal is received 45 | s = <-sig // listen again to force quit 46 | log.Printf("received signal '%s', forcefully quitting...", s) 47 | os.Exit(1) // forcefully terminate the program 48 | }() 49 | 50 | return done 51 | } 52 | 53 | func timeoutHandler(done <-chan struct{}, timeout time.Duration) { 54 | <-done // make sure we are in termination phase 55 | 56 | // create a timer to be able to handle timeouts 57 | time.AfterFunc(timeout, func() { 58 | log.Println("timeoutHandler: cleanup phase timeout reached, forcefully quitting...") 59 | os.Exit(1) 60 | }) 61 | } 62 | 63 | func runApplication(done <-chan struct{}) { 64 | defer cleanupApplication() 65 | 66 | for { 67 | select { 68 | case _, ok := <-done: 69 | if !ok { 70 | // exit from application when the channel is closed 71 | return 72 | } 73 | default: 74 | log.Println("doing stuff...") 75 | time.Sleep(time.Second) 76 | } 77 | } 78 | } 79 | 80 | func cleanupApplication() { 81 | time.Sleep(5 * time.Second) 82 | log.Println("application done.") 83 | } 84 | -------------------------------------------------------------------------------- /examples/blog/04_context_wait_group_force_timeout/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "log" 7 | "os" 8 | "os/signal" 9 | "sync" 10 | "syscall" 11 | "time" 12 | ) 13 | 14 | func init() { 15 | log.SetFlags(0) 16 | log.SetOutput(os.Stdout) 17 | } 18 | 19 | func main() { 20 | ctx := context.Background() 21 | // start signal handler 22 | ctx = signalHandler(ctx) 23 | 24 | // create a sync.WaitGroup and add 1 to it for each go-routine 25 | wg := sync.WaitGroup{} 26 | wg.Add(1) 27 | go func() { 28 | defer wg.Done() 29 | err := runApplication(ctx) 30 | if err != nil { 31 | log.Printf("main: error in application: %v\n", err) 32 | } 33 | }() 34 | 35 | // run the timeoutHandler handler to have a hard time limit on cleanup phase 36 | go timeoutHandler(ctx, 30*time.Second) 37 | 38 | // we wait for all the go-routines 39 | wg.Wait() 40 | } 41 | 42 | func signalHandler(ctx context.Context) context.Context { 43 | ctx, cancel := context.WithCancel(ctx) 44 | go func() { 45 | sig := make(chan os.Signal, 1) 46 | signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) 47 | s := <-sig 48 | log.Printf("received signal '%s', gracefully quitting...", s) 49 | cancel() // instead of calling close directly we call context.CancelFunc when a signal is received 50 | s = <-sig // listen again to force quit 51 | log.Printf("received signal '%s', forcefully quitting...", s) 52 | os.Exit(1) // forcefully terminate the program 53 | }() 54 | 55 | return ctx 56 | } 57 | 58 | func timeoutHandler(ctx context.Context, timeout time.Duration) { 59 | <-ctx.Done() // make sure we are in termination phase 60 | 61 | // create a timer to be able to handle timeouts 62 | time.AfterFunc(timeout, func() { 63 | log.Println("timeoutHandler: cleanup phase timeout reached, forcefully quitting...") 64 | os.Exit(1) 65 | }) 66 | } 67 | 68 | func runApplication(ctx context.Context) (err error) { 69 | defer func() { 70 | err = cleanupApplication() 71 | }() 72 | 73 | for { 74 | select { 75 | case <-ctx.Done(): 76 | err = ctx.Err() 77 | if err == context.Canceled { 78 | err = nil 79 | } 80 | 81 | return 82 | default: 83 | log.Println("doing stuff...") 84 | time.Sleep(time.Second) 85 | } 86 | } 87 | } 88 | 89 | func cleanupApplication() error { 90 | time.Sleep(5 * time.Second) 91 | log.Println("application done.") 92 | return errors.New("intentional error") 93 | } 94 | -------------------------------------------------------------------------------- /examples/blog/05_context_errgroup_force_timeout/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "golang.org/x/sync/errgroup" 7 | "log" 8 | "os" 9 | "os/signal" 10 | "syscall" 11 | "time" 12 | ) 13 | 14 | func init() { 15 | log.SetFlags(0) 16 | log.SetOutput(os.Stdout) 17 | } 18 | 19 | func main() { 20 | ctx := context.Background() 21 | 22 | // start signal handler 23 | ctx = signalHandler(ctx) 24 | 25 | // create a group from errgroup package 26 | group, ctx := errgroup.WithContext(ctx) 27 | 28 | group.Go(func() error { 29 | return runApplication(ctx) 30 | }) 31 | 32 | // run the timeoutHandler handler to have a hard time limit on cleanup phase 33 | go timeoutHandler(ctx, 30*time.Second) 34 | 35 | if err := group.Wait(); err != nil { 36 | log.Printf("one of the go-routines failed: %v", err) 37 | os.Exit(1) 38 | } 39 | } 40 | 41 | func signalHandler(ctx context.Context) context.Context { 42 | ctx, cancel := context.WithCancel(ctx) 43 | go func() { 44 | sig := make(chan os.Signal, 1) 45 | signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) 46 | s := <-sig 47 | log.Printf("received signal '%s', gracefully quitting...", s) 48 | cancel() // instead of calling close directly we call context.CancelFunc when a signal is received 49 | s = <-sig // listen again to force quit 50 | log.Printf("received signal '%s', forcefully quitting...", s) 51 | os.Exit(1) // forcefully terminate the program 52 | }() 53 | 54 | return ctx 55 | } 56 | 57 | func timeoutHandler(ctx context.Context, timeout time.Duration) { 58 | <-ctx.Done() // make sure we are in termination phase 59 | 60 | // create a timer to be able to handle timeouts 61 | time.AfterFunc(timeout, func() { 62 | log.Println("timeoutHandler: cleanup phase timeout reached, forcefully quitting...") 63 | os.Exit(1) 64 | }) 65 | } 66 | 67 | func runApplication(ctx context.Context) (err error) { 68 | defer func() { 69 | err = cleanupApplication() 70 | }() 71 | 72 | for { 73 | select { 74 | case <-ctx.Done(): 75 | err = ctx.Err() 76 | if err == context.Canceled { 77 | err = nil 78 | } 79 | 80 | return 81 | default: 82 | log.Println("doing stuff...") 83 | time.Sleep(time.Second) 84 | } 85 | } 86 | } 87 | 88 | func cleanupApplication() error { 89 | time.Sleep(5 * time.Second) 90 | //log.Println("application done.") 91 | return errors.New("intentional error") 92 | } 93 | -------------------------------------------------------------------------------- /examples/lib/http_example/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "github.com/itzloop/gograce" 7 | "log" 8 | "net" 9 | "net/http" 10 | "time" 11 | ) 12 | 13 | func main() { 14 | // create a grace instance 15 | grace := gograce.NewGraceful(gograce.Options{ 16 | Timeout: 15 * time.Second, // wait 15 seconds and forcefully terminate the application 17 | NoForceQuit: false, // by pressing Ctrl+C twice, app will terminate immediately 18 | MaxGoRoutines: 0, // set no limit for the number of go-routines running at the same time 19 | Signals: nil, // use the defaultSignals in signal.Notify. 20 | }) 21 | 22 | // create a simple http server 23 | exampleHTTPServer := NewExampleHTTPServer(":8000") 24 | 25 | // add start and close operations to grace instance 26 | grace.GoWithContext(exampleHTTPServer.start) 27 | grace.GoWithContext(exampleHTTPServer.close) 28 | 29 | // wait for all go-routines or the cancel signal and 30 | // if any error is encountered, call log.Fatal 31 | grace.FatalWait() 32 | } 33 | 34 | type ExampleHTTPServer struct { 35 | // addr is the listen address for the http.Server 36 | addr string 37 | 38 | // an instance of http.Server that is used to server request 39 | httpServer *http.Server 40 | } 41 | 42 | // NewExampleHTTPServer creates an instance of ExampleHTTPServer 43 | func NewExampleHTTPServer(addr string) *ExampleHTTPServer { 44 | mux := http.NewServeMux() 45 | mux.HandleFunc("/log-running-job", func(writer http.ResponseWriter, request *http.Request) { 46 | time.Sleep(10 * time.Second) 47 | if _, err := writer.Write([]byte("done")); err != nil { 48 | log.Printf("log-running-job: failed to write response: %v", err) 49 | } 50 | }) 51 | mux.HandleFunc("/hello", func(writer http.ResponseWriter, request *http.Request) { 52 | if _, err := writer.Write([]byte("hi")); err != nil { 53 | log.Printf("log-running-job: failed to write response: %v", err) 54 | } 55 | }) 56 | 57 | server := http.Server{ 58 | Addr: addr, 59 | Handler: mux, 60 | } 61 | 62 | return &ExampleHTTPServer{ 63 | addr: addr, 64 | httpServer: &server, 65 | } 66 | } 67 | 68 | // start starts the httpServer and sets the http.Server.BaseContext. 69 | func (s *ExampleHTTPServer) start(ctx context.Context) (err error) { 70 | s.httpServer.BaseContext = func(_ net.Listener) context.Context { 71 | return ctx 72 | } 73 | 74 | err = s.httpServer.ListenAndServe() 75 | if err != nil && err != http.ErrServerClosed { 76 | return err 77 | } 78 | 79 | return nil 80 | } 81 | 82 | // backup runs an imaginary backup routine 83 | func (s *ExampleHTTPServer) backup() (err error) { 84 | log.Printf("ExampleHTTPServer.backup: backing up some imaginary stuff") 85 | time.Sleep(time.Second * 2) 86 | log.Printf("ExampleHTTPServer.backup: backuped everything") 87 | return nil 88 | } 89 | 90 | // shutdown shuts down the http server so no new connectionos 91 | // are accepted and running connections will have time to finish 92 | func (s *ExampleHTTPServer) shutdown() (err error) { 93 | log.Printf("ExampleHTTPServer.shutdown: shutting down http server") 94 | return s.httpServer.Shutdown(context.Background()) 95 | } 96 | 97 | // close runs shutdown and backup and errors.Join their errors and returns 98 | func (s *ExampleHTTPServer) close(ctx context.Context) (err error) { 99 | <-ctx.Done() 100 | err = errors.Join(s.shutdown()) 101 | err = errors.Join(s.backup()) 102 | return 103 | } 104 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/itzloop/gograce 2 | 3 | go 1.21.2 4 | 5 | require ( 6 | github.com/stretchr/testify v1.8.4 7 | golang.org/x/sync v0.4.0 8 | ) 9 | 10 | require ( 11 | github.com/davecgh/go-spew v1.1.1 // indirect 12 | github.com/pmezard/go-difflib v1.0.0 // indirect 13 | gopkg.in/yaml.v3 v3.0.1 // indirect 14 | ) 15 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 2 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 3 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 4 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 5 | github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= 6 | github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= 7 | golang.org/x/sync v0.4.0 h1:zxkM55ReGkDlKSM+Fu41A+zmbZuaPVbGMzvvdUPznYQ= 8 | golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= 9 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 10 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 11 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 12 | -------------------------------------------------------------------------------- /graceful.go: -------------------------------------------------------------------------------- 1 | package gograce 2 | 3 | import ( 4 | "context" 5 | "log" 6 | "os" 7 | "time" 8 | 9 | "golang.org/x/sync/errgroup" 10 | ) 11 | 12 | type Options struct { 13 | // Timeout defines how long should the program wait before forcefully exiting. 14 | // a zero-value indicates no timeout. 15 | Timeout time.Duration 16 | 17 | // NoForceQuit disables the force quit feature. After the first termination signal, any further signals 18 | // will be ignored. 19 | NoForceQuit bool 20 | 21 | // MaxGoRoutines defines how many go-routines can be started. This value is passed to SetLimit on errgroup.Group. 22 | // a zero-value or negative indicates no limit. 23 | MaxGoRoutines int 24 | 25 | // TODO custom signals? 26 | // Signals let's you overwrite graceful.defaultSignals. 27 | // a zero-value or an empty slice indicate no overwrite 28 | Signals []os.Signal 29 | } 30 | 31 | type Graceful struct { 32 | // TODO is this a good idea? [issue#22602](https://github.com/golang/go/issues/22602) 33 | // TODO I don't like doing a bunch of g.Go(func() error { return f(ctx) }) 34 | // TODO instead i'd like to do g.Go(f) 35 | ctx context.Context 36 | g *errgroup.Group 37 | sh *SignalHandler 38 | th *TimeoutHandler 39 | } 40 | 41 | // NewGraceful calls NewGracefulWithContext with context.Background() 42 | func NewGraceful(opts Options) *Graceful { 43 | return NewGracefulWithContext(context.Background(), opts) 44 | } 45 | 46 | // NewGracefulWithContext will create a SignalHandler and a TimeoutHandler 47 | // which are started automatically. 48 | func NewGracefulWithContext(ctx context.Context, opts Options) *Graceful { 49 | var ( 50 | g *errgroup.Group 51 | graceful = &Graceful{} 52 | signals = defaultSignals[:] 53 | ) 54 | 55 | // run signal handler 56 | if len(opts.Signals) != 0 { 57 | signals = opts.Signals 58 | } 59 | 60 | // Create signal handler 61 | graceful.sh, ctx = NewSignalHandler(ctx, SignalHandlerOptions{ 62 | Force: !opts.NoForceQuit, 63 | Signals: signals, 64 | }) 65 | 66 | if opts.Timeout != 0 { 67 | graceful.th = NewTimeoutHandler(ctx, TimeoutHandlerOptions{ 68 | Timeout: opts.Timeout, 69 | }) 70 | } 71 | 72 | g, ctx = errgroup.WithContext(ctx) 73 | 74 | if opts.MaxGoRoutines != 0 { 75 | g.SetLimit(opts.MaxGoRoutines) 76 | } 77 | 78 | graceful.g = g 79 | graceful.ctx = ctx 80 | 81 | return graceful 82 | } 83 | 84 | // GoWithContext is convenient wrapper for (*errgroup.Group).Go that 85 | // accepts a functions that takes a context as input instead of not 86 | // having any input. 87 | func (grace *Graceful) GoWithContext(f func(ctx context.Context) error) { 88 | grace.g.Go(func() error { 89 | return f(grace.ctx) 90 | }) 91 | } 92 | 93 | // Go calls (*errgroup.Group).Go() internally 94 | func (grace *Graceful) Go(f func() error) { 95 | grace.g.Go(f) 96 | } 97 | 98 | // Wait calls (*errgroup.Group).Wait() and returns the error 99 | func (grace *Graceful) Wait() error { 100 | return grace.g.Wait() 101 | } 102 | 103 | // FatalWait calls Wait but log.Fatals when an error is received 104 | func (grace *Graceful) FatalWait() { 105 | if err := grace.Wait(); err != nil { 106 | log.Fatalln(err.Error()) 107 | } 108 | } 109 | -------------------------------------------------------------------------------- /graceful_test.go: -------------------------------------------------------------------------------- 1 | package gograce 2 | 3 | import ( 4 | "context" 5 | "sync" 6 | "syscall" 7 | "testing" 8 | "time" 9 | 10 | "github.com/stretchr/testify/assert" 11 | ) 12 | 13 | func TestGracefulForce(t *testing.T) { 14 | grace := NewGracefulWithContext(context.Background(), Options{ 15 | NoForceQuit: false, 16 | }) 17 | 18 | var ( 19 | started bool 20 | ended bool 21 | forceHanlerCalled bool 22 | wg = sync.WaitGroup{} 23 | ) 24 | 25 | grace.GoWithContext(func(ctx context.Context) error { 26 | started = true 27 | return nil 28 | }) 29 | 30 | grace.GoWithContext(func(ctx context.Context) error { 31 | time.Sleep(100 * time.Millisecond) 32 | ended = true 33 | return nil 34 | }) 35 | 36 | wg.Add(1) 37 | grace.sh.forceFunc = func() { 38 | defer wg.Done() 39 | forceHanlerCalled = true 40 | } 41 | 42 | go func() { 43 | grace.sh.sigChan <- syscall.SIGINT 44 | grace.sh.sigChan <- syscall.SIGINT 45 | }() 46 | 47 | wg.Wait() 48 | assert.True(t, forceHanlerCalled) 49 | assert.True(t, started) 50 | assert.False(t, ended) 51 | } 52 | 53 | func TestGracefulTimeout(t *testing.T) { 54 | grace := NewGracefulWithContext(context.Background(), Options{ 55 | Timeout: time.Millisecond, 56 | NoForceQuit: false, 57 | }) 58 | 59 | var ( 60 | started bool 61 | ended bool 62 | timeoutHandlerFunc bool 63 | wg = sync.WaitGroup{} 64 | ) 65 | 66 | grace.GoWithContext(func(ctx context.Context) error { 67 | <-ctx.Done() 68 | started = true 69 | return nil 70 | }) 71 | 72 | grace.GoWithContext(func(ctx context.Context) error { 73 | <-ctx.Done() 74 | time.Sleep(10 * time.Second) 75 | ended = true 76 | return nil 77 | }) 78 | 79 | go func() { 80 | grace.sh.sigChan <- syscall.SIGINT 81 | }() 82 | 83 | wg.Add(1) 84 | grace.th.timeoutFunc = func() { 85 | defer wg.Done() 86 | timeoutHandlerFunc = true 87 | } 88 | 89 | wg.Wait() 90 | 91 | assert.True(t, timeoutHandlerFunc) 92 | assert.True(t, started) 93 | assert.False(t, ended) 94 | 95 | } 96 | 97 | func TestGraceful(t *testing.T) { 98 | grace := NewGracefulWithContext(context.Background(), Options{ 99 | Timeout: 1 * time.Second, 100 | NoForceQuit: false, 101 | }) 102 | 103 | var ( 104 | started bool 105 | ended bool 106 | ) 107 | 108 | grace.GoWithContext(func(ctx context.Context) error { 109 | <-ctx.Done() 110 | started = true 111 | return nil 112 | }) 113 | 114 | grace.GoWithContext(func(ctx context.Context) error { 115 | <-ctx.Done() 116 | ended = true 117 | return nil 118 | }) 119 | 120 | go func() { 121 | grace.sh.sigChan <- syscall.SIGINT 122 | }() 123 | 124 | grace.Wait() 125 | 126 | assert.True(t, started) 127 | assert.True(t, ended) 128 | 129 | } 130 | 131 | func TestGracefulAppExits(t *testing.T) { 132 | ctx, cancel := context.WithCancel(context.Background()) 133 | grace := NewGracefulWithContext(ctx, Options{ 134 | Timeout: 1 * time.Second, 135 | NoForceQuit: false, 136 | }) 137 | 138 | var ( 139 | started bool 140 | ended bool 141 | ) 142 | 143 | grace.GoWithContext(func(ctx context.Context) error { 144 | defer cancel() 145 | time.Sleep(100 * time.Millisecond) 146 | started = true 147 | return nil 148 | }) 149 | 150 | grace.GoWithContext(func(ctx context.Context) error { 151 | <-ctx.Done() 152 | ended = true 153 | return nil 154 | }) 155 | 156 | grace.Wait() 157 | 158 | assert.True(t, started) 159 | assert.True(t, ended) 160 | } 161 | -------------------------------------------------------------------------------- /signal_handler.go: -------------------------------------------------------------------------------- 1 | package gograce 2 | 3 | import ( 4 | "context" 5 | "log" 6 | "os" 7 | "sync/atomic" 8 | "syscall" 9 | 10 | "os/signal" 11 | ) 12 | 13 | type ForceFunc func() 14 | 15 | var ( 16 | defaultSignals = [...]os.Signal{syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP} 17 | ) 18 | 19 | // SignalHandlerOptions 20 | type SignalHandlerOptions struct { 21 | // Force enables quiting forcefully (by sending one of the Signals twice) 22 | // when graceful shutdown is in progress 23 | Force bool 24 | 25 | // Signals overwrites the defaultSignals. 26 | Signals []os.Signal 27 | 28 | // ForceFunc is called when Force = true and one of the Signals is sent twice. 29 | // If ForceFunc is nil, defaultForceFunc will be used which is os.Exit(1). 30 | ForceFunc ForceFunc 31 | } 32 | 33 | // A SignalHandler listens for signals and handles graceful and forceful shutdown 34 | // When a signal has been sent twice SignalHandler will call forceFunc. Default 35 | // forceFunc is os.Exit(1) so application will terminate. 36 | type SignalHandler struct { 37 | signals []os.Signal 38 | force bool 39 | sigChan chan os.Signal 40 | 41 | forceFunc ForceFunc 42 | 43 | started atomic.Bool 44 | } 45 | 46 | // NewSignalHandler will create a signal handler based on the desired opts given. 47 | // It will then start the signal handler as well. 48 | func NewSignalHandler(ctx context.Context, opts SignalHandlerOptions) (*SignalHandler, context.Context) { 49 | if len(opts.Signals) == 0 { 50 | opts.Signals = defaultSignals[:] 51 | } 52 | 53 | if opts.ForceFunc == nil { 54 | opts.ForceFunc = defaultForceFunc 55 | } 56 | 57 | sh := &SignalHandler{ 58 | signals: opts.Signals, 59 | force: opts.Force, 60 | started: atomic.Bool{}, 61 | forceFunc: opts.ForceFunc, 62 | } 63 | 64 | ctx = sh.Start(ctx) 65 | 66 | return sh, ctx 67 | } 68 | 69 | // Start will start the SignalHandler by listening to signals and parent context. 70 | // If at anypoint parent context gets canceled, Start will return. It is safe 71 | // but useless to call Start from multiple go-routines because it will start it 72 | // the first and you have to Close it first to be able to Start it again. 73 | func (s *SignalHandler) Start(ctx context.Context) context.Context { 74 | if s.started.Swap(true) { 75 | return ctx // TODO should this be nil or not? 76 | } 77 | 78 | // at any point we need to stop execution when the 79 | // parent context gets canceled so we make a copy 80 | // of it. 81 | parentCtx := ctx 82 | ctx, cancel := context.WithCancel(ctx) 83 | s.sigChan = make(chan os.Signal, 1) 84 | 85 | go func() { 86 | signal.Notify(s.sigChan, s.signals...) 87 | 88 | defer s.Close() 89 | var ( 90 | sig os.Signal 91 | ok bool 92 | ) 93 | 94 | select { 95 | case sig, ok = <-s.sigChan: 96 | if !ok { 97 | log.Println("signal channel closed quiting...") 98 | return 99 | } 100 | log.Printf("received signal '%s', gracefully quitting...\n", sig) 101 | cancel() 102 | case <-parentCtx.Done(): 103 | log.Printf("parent context canceled\n") 104 | return 105 | } 106 | 107 | if s.force { 108 | select { 109 | case sig, ok = <-s.sigChan: 110 | if !ok { 111 | log.Println("signal channel closed quiting...") 112 | return 113 | } 114 | 115 | log.Printf("received signal '%s', forcefully quitting...\n", sig) 116 | cancel() 117 | case <-parentCtx.Done(): 118 | log.Printf("parent context canceled, while waiting for second signal\n") 119 | return 120 | } 121 | 122 | s.Close() 123 | 124 | s.forceFunc() 125 | return 126 | } 127 | 128 | s.Close() 129 | }() 130 | 131 | return ctx 132 | } 133 | 134 | // Close closes sigChan. Calls to close only work when SignalHandler has been started 135 | // and other wise it has no effect. It is also safe to call it from multiple go-routines. 136 | func (sh *SignalHandler) Close() { 137 | if sh.started.Swap(false) { 138 | return 139 | } 140 | 141 | close(sh.sigChan) 142 | } 143 | 144 | func defaultForceFunc() { 145 | os.Exit(1) 146 | } 147 | -------------------------------------------------------------------------------- /signal_handler_test.go: -------------------------------------------------------------------------------- 1 | package gograce 2 | 3 | import ( 4 | "context" 5 | "sync" 6 | "syscall" 7 | "testing" 8 | "time" 9 | 10 | "github.com/stretchr/testify/require" 11 | ) 12 | 13 | func TestSignalHandler(t *testing.T) { 14 | t.Run("without force", func(t *testing.T) { 15 | sh, ctx := NewSignalHandler(context.Background(), SignalHandlerOptions{ 16 | Force: false, 17 | }) 18 | 19 | require.True(t, sh.started.Load()) 20 | 21 | go func() { 22 | sh.sigChan <- syscall.SIGINT 23 | }() 24 | 25 | <-ctx.Done() 26 | require.ErrorIs(t, ctx.Err(), context.Canceled) 27 | }) 28 | 29 | t.Run("with force", func(t *testing.T) { 30 | var forceCalled bool 31 | sh, ctx := NewSignalHandler(context.Background(), SignalHandlerOptions{ 32 | Force: true, 33 | }) 34 | 35 | sh.forceFunc = func() { 36 | forceCalled = true 37 | } 38 | 39 | go func() { 40 | sh.sigChan <- syscall.SIGINT 41 | sh.sigChan <- syscall.SIGINT 42 | }() 43 | 44 | <-ctx.Done() 45 | require.ErrorIs(t, ctx.Err(), context.Canceled) 46 | require.True(t, forceCalled) 47 | require.False(t, sh.started.Load()) 48 | }) 49 | 50 | t.Run("cancel parent context", func(t *testing.T) { 51 | var forceCalled bool 52 | ctx, cancel := context.WithCancel(context.Background()) 53 | sh, ctx := NewSignalHandler(ctx, SignalHandlerOptions{ 54 | Force: true, 55 | ForceFunc: func() { 56 | forceCalled = true 57 | }, 58 | }) 59 | 60 | cancel() 61 | 62 | time.Sleep(time.Millisecond * 100) 63 | 64 | require.ErrorIs(t, ctx.Err(), context.Canceled) 65 | require.False(t, forceCalled) 66 | require.False(t, sh.started.Load()) 67 | }) 68 | 69 | t.Run("multiple start and close", func(t *testing.T) { 70 | sh, ctx := NewSignalHandler(context.Background(), SignalHandlerOptions{ 71 | Force: false, 72 | }) 73 | 74 | require.True(t, sh.started.Load()) 75 | 76 | go func() { 77 | sh.sigChan <- syscall.SIGINT 78 | }() 79 | 80 | <-ctx.Done() 81 | require.ErrorIs(t, ctx.Err(), context.Canceled) 82 | 83 | sh.Start(context.Background()) 84 | 85 | require.True(t, sh.started.Load()) 86 | 87 | go func() { 88 | sh.sigChan <- syscall.SIGINT 89 | }() 90 | 91 | <-ctx.Done() 92 | require.ErrorIs(t, ctx.Err(), context.Canceled) 93 | 94 | }) 95 | 96 | t.Run("calling close", func(t *testing.T) { 97 | sh, _ := NewSignalHandler(context.Background(), SignalHandlerOptions{ 98 | Force: false, 99 | }) 100 | 101 | require.True(t, sh.started.Load()) 102 | 103 | sh.Close() 104 | 105 | require.False(t, sh.started.Load()) 106 | }) 107 | 108 | t.Run("calling close while force quiting", func(t *testing.T) { 109 | var forceCalled bool 110 | sh, _ := NewSignalHandler(context.Background(), SignalHandlerOptions{ 111 | Force: true, 112 | ForceFunc: func() { 113 | forceCalled = true 114 | }, 115 | }) 116 | 117 | require.True(t, sh.started.Load()) 118 | 119 | sh.sigChan <- syscall.SIGINT 120 | 121 | time.Sleep(100 * time.Millisecond) 122 | require.True(t, sh.started.Load()) 123 | 124 | sh.Close() 125 | 126 | require.False(t, sh.started.Load()) 127 | require.False(t, forceCalled) 128 | }) 129 | 130 | t.Run("calling start multiple times", func(t *testing.T) { 131 | var ( 132 | forceCalled bool 133 | wg = sync.WaitGroup{} 134 | ) 135 | 136 | wg.Add(1) 137 | sh, _ := NewSignalHandler(context.Background(), SignalHandlerOptions{ 138 | Force: true, 139 | ForceFunc: func() { 140 | defer wg.Done() 141 | forceCalled = true 142 | }, 143 | }) 144 | 145 | sh.Start(context.Background()) 146 | sh.Start(context.Background()) 147 | sh.Start(context.Background()) 148 | sh.Start(context.Background()) 149 | 150 | require.True(t, sh.started.Load()) 151 | 152 | sh.sigChan <- syscall.SIGINT 153 | sh.sigChan <- syscall.SIGINT 154 | wg.Wait() 155 | 156 | require.False(t, sh.started.Load()) 157 | require.True(t, forceCalled) 158 | }) 159 | } 160 | -------------------------------------------------------------------------------- /timeout_handler.go: -------------------------------------------------------------------------------- 1 | package gograce 2 | 3 | import ( 4 | "context" 5 | "log" 6 | "os" 7 | "time" 8 | ) 9 | 10 | // TimeoutFunc will be called when the the 11 | type TimeoutFunc func() 12 | 13 | // TimeoutHandlerOptions 14 | type TimeoutHandlerOptions struct { 15 | // Timeout is the value that is passed to time.AfterFunc. 16 | Timeout time.Duration 17 | 18 | // TimeoutFunc is the function that is passed to time.AfterFunc. 19 | TimeoutFunc TimeoutFunc 20 | } 21 | 22 | // TimeoutHandler will set a hard limit for graceful shutdown. If that limit 23 | // is reached the program will call timeoutFunc. defaultTimeoutFunc is os.Exit(1) 24 | // so this will terminate the application. 25 | type TimeoutHandler struct { 26 | timeout time.Duration 27 | 28 | timeoutFunc TimeoutFunc 29 | } 30 | 31 | // NewTimeoutHandler 32 | func NewTimeoutHandler(ctx context.Context, opts TimeoutHandlerOptions) *TimeoutHandler { 33 | if opts.TimeoutFunc == nil { 34 | opts.TimeoutFunc = defaultTimeoutFunc 35 | } 36 | 37 | th := &TimeoutHandler{ 38 | timeout: opts.Timeout, 39 | timeoutFunc: opts.TimeoutFunc, 40 | } 41 | 42 | go th.Start(ctx) 43 | 44 | return th 45 | } 46 | 47 | func (th *TimeoutHandler) Start(ctx context.Context) { 48 | <-ctx.Done() // make sure we are in termination phase 49 | 50 | // create a timer to be able to handle timeouts 51 | time.AfterFunc(th.timeout, func() { 52 | log.Println("timeoutHandler: cleanup phase timeout reached, forcefully quitting...") 53 | th.timeoutFunc() 54 | }) 55 | } 56 | 57 | func defaultTimeoutFunc() { 58 | os.Exit(1) 59 | } 60 | -------------------------------------------------------------------------------- /timeout_handler_test.go: -------------------------------------------------------------------------------- 1 | package gograce 2 | 3 | import ( 4 | "context" 5 | "sync" 6 | "testing" 7 | "time" 8 | 9 | "github.com/stretchr/testify/require" 10 | ) 11 | 12 | func TestTimeoutHandler(t *testing.T) { 13 | var ( 14 | timeoutFuncCalled bool 15 | wg = sync.WaitGroup{} 16 | timeoutFunc = func() { 17 | defer wg.Done() 18 | timeoutFuncCalled = true 19 | } 20 | ) 21 | 22 | ctx, cancel := context.WithCancel(context.Background()) 23 | 24 | wg.Add(1) 25 | NewTimeoutHandler(ctx, TimeoutHandlerOptions{ 26 | Timeout: 100 * time.Millisecond, 27 | TimeoutFunc: timeoutFunc, 28 | }) 29 | 30 | cancel() 31 | 32 | wg.Wait() 33 | 34 | require.True(t, timeoutFuncCalled) 35 | } 36 | --------------------------------------------------------------------------------