├── .gitignore ├── .travis.yml ├── LICENSE ├── README.md ├── examples ├── server-channel.go ├── server-example.go └── simple-func.go ├── http.go ├── http_test.go ├── shutdown.go └── shutdown_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 2 | *.o 3 | *.a 4 | *.so 5 | 6 | # Folders 7 | _obj 8 | _test 9 | 10 | # Architecture specific extensions/prefixes 11 | *.[568vq] 12 | [568vq].out 13 | 14 | *.cgo1.go 15 | *.cgo2.c 16 | _cgo_defun.c 17 | _cgo_gotypes.go 18 | _cgo_export.* 19 | 20 | _testmain.go 21 | 22 | *.exe 23 | *.test 24 | *.prof 25 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 Klaus Post 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 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # old version 2 | 3 | This package has been updated to a new version. If you are starting a new project, use the new package: 4 | 5 | * Package home: https://github.com/klauspost/shutdown2 6 | * Godoc: https://godoc.org/github.com/klauspost/shutdown2 7 | 8 | Version 2 mainly contains minor adjustments to the API to make it a bit easier to use. 9 | 10 | This package remains here to maintain compatibility. 11 | 12 | # shutdown 13 | Shutdown management library for Go 14 | 15 | This package helps you manage shutdown code centrally, and provides functionality to execute code when a controlled shutdown occurs. 16 | 17 | This will enable you to save data, notify other services that your application is shutting down. 18 | 19 | * Package home: https://github.com/klauspost/shutdown 20 | * Godoc: https://godoc.org/github.com/klauspost/shutdown 21 | 22 | [![GoDoc][1]][2] [![Build Status][3]][4] 23 | 24 | [1]: https://godoc.org/github.com/klauspost/shutdown?status.svg 25 | [2]: https://godoc.org/github.com/klauspost/shutdown 26 | [3]: https://travis-ci.org/klauspost/shutdown.svg 27 | [4]: https://travis-ci.org/klauspost/shutdown 28 | 29 | # concept 30 | Managing shutdowns can be very tricky, often leading to races, crashes and strange behavior. 31 | This package will help you manage the shutdown process and will attempt to fix some of the common problems when dealing with shutting down. 32 | 33 | The shutdown package allow you to block shutdown while certain parts of your code is running. This is helpful to ensure that operations are not interrupted. 34 | 35 | The second part of the shutdown process is notifying goroutines in a select loop and calling functions in your code that handles various shutdown procedures, like closing databases, notifying other servers, deleting temporary files, etc. 36 | 37 | The second part of the process has three **stages**, which will enable you to do your shutdown in stages. This will enable you to rely on some parts, like logging, to work in the first two stages. There is no rules for what you should put in which stage, but things executing in stage one can safely rely on stage two not being executed yet. 38 | 39 | All operations have **timeouts**. This is to fix another big issue with shutdowns; applications that hang on shutdown. The timeout is for each stage of the shutdown process, and can be adjusted to your application needs. If a timeout is exceeded the next shutdown stage will be initiated regardless. 40 | 41 | Finally, you can always cancel a notifier, which will remove it from the shutdown queue. 42 | 43 | # usage 44 | 45 | First get the libary with `go get -u github.com/klauspost/shutdown`, and add it as an import to your code with `import github.com/klauspost/shutdown`. 46 | 47 | The next thing you probably want to do is to register Ctrl+c and system terminate. This will make all shutdown handlers run when any of these are sent to your program: 48 | ```Go 49 | shutdown.OnSignal(0, os.Interrupt, syscall.SIGTERM) 50 | ``` 51 | 52 | If you don't like the default timeout duration of 5 seconds, you can change it by calling the `SetTimeout` function: 53 | ```Go 54 | shutdown.SetTimeout(time.Second * 1) 55 | ``` 56 | Now the maximum delay for shutdown is **4 seconds**. The timeout is applied to each of the stages and that is also the maximum time to wait for the shutdown to begin. If you need to adjust a single stage, use `SetTimeoutN` function. 57 | 58 | Next you can register functions to run when shutdown runs: 59 | ```Go 60 | logFile := os.Create("log.txt") 61 | 62 | // Execute the function in the first stage of the shutdown process 63 | _ = shutdown.FirstFunc(func(interface{}){ 64 | logFile.Close() 65 | }, nil) 66 | 67 | // Execute this function in the second part of the shutdown process 68 | _ = shutdown.SecondFunc(func(interface{}){ 69 | _ = os.Delete("log.txt") 70 | }, nil) 71 | ``` 72 | As noted there are three stages. All functions in one stage are executed in parallel, but the package will wait for all functions in one stage to have finished before moving on to the next one. So your code cannot rely on any particular order of execution inside a single stage, but you are guaranteed that the First stage is finished before any functions from stage two are executed. 73 | 74 | You can send a parameter to your function, which is delivered as an `interface{}`. This way you can re-use the same function for similar tasks. See `simple-func.go` in the examples folder. 75 | 76 | This example above uses functions that are called, but you can also request channels that are notified on shutdown. This allows you do have shutdown handling in blocked select statements like this: 77 | 78 | ```Go 79 | go func() { 80 | // Get a stage 1 notification 81 | finish := shutdown.First() 82 | select { 83 | case n:= <-finish: 84 | log.Println("Closing") 85 | close(n) 86 | return 87 | } 88 | ``` 89 | It is important that you close the channel you receive. This is your way of signalling that you are done. If you do not close the channel you get shutdown will wait until the timeout has expired before proceeding to the next stage. 90 | 91 | If you for some reason don't need a notifier anymore you can cancel it. When a notifier has been cancelled it will no longer receive notifications, and the shutdown code will no longer wait for it on exit. 92 | ```Go 93 | go func() { 94 | // Get a stage 1 notification 95 | finish := shutdown.First() 96 | select { 97 | case n:= <-finish: 98 | close(n) 99 | return 100 | case <-otherchan: 101 | finish.Cancel() 102 | return 103 | } 104 | ``` 105 | Functions are cancelled the same way by cancelling the returned notifier. Be aware that if shutdown has been initiated you can no longer cancel notifiers, so you may need to aquire a shutdown lock (see below). 106 | 107 | The final thing you can do is to lock shutdown in parts of your code you do not want to be interrupted by a shutdown, or if the code relies on resources that are destroyed as part of the shutdown process. 108 | 109 | A simple example can be seen in this http handler: 110 | ```Go 111 | http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) { 112 | // Acquire a lock. 113 | // While this is held server will not shut down (except after timeout) 114 | if !shutdown.Lock() { 115 | // Shutdown has started, return that the service is unavailable 116 | w.WriteHeader(http.StatusServiceUnavailable) 117 | } 118 | defer shutdown.Unlock() 119 | io.WriteString(w, "Server running") 120 | }) 121 | ``` 122 | If shutdown is started, either by a signal or by another goroutine, it will wait until the lock is released. It is important always to release the lock, if shutdown.Lock() returns true. Otherwise the server will have to wait until the timeout has passed before it starts shutting down, which may not be what you want. 123 | 124 | Finally you can call `shutdown.Exit(exitcode)` to call all exit handlers and exit your application. This will wait for all locks to be released and notify all shutdown handlers and exit with the given exit code. If you want to do the exit yourself you can call the `shutdown.Shutdown()`, whihc does the same, but doesn't exit. Beware that you don't hold a lock when you call Exit/Shutdown. 125 | 126 | 127 | Also there are some things to be mindful of: 128 | * Notifiers **can** be created inside shutdown code, but only for stages **following** the current. So stage 1 notifiers can create stage 2 notifiers, but if they create a stage 1 notifier this will never be called. 129 | * Timeout can be changed once shutdown has been initiated, but it will only affect the **following** stages. 130 | * Notifiers returned from a function (eg. FirstFunc) can be used for selects. They will be notified, but the shutdown manager will not wait for them to finish, so using them for this is not recommended. 131 | * If a panic occurs inside a shutdown function call in your code, the panic will be recovered and **ignored** and the shutdown will proceed. A message is printed to `log`. If you want to handle panics, you must do it in your code. 132 | * When shutdown is initiated, it cannot be stopped. 133 | 134 | When you design with this do take care that this library is for **controlled** shutdown of your application. If you application crashes no shutdown handlers are run, so panics will still be fatal. You can of course still call the `Shutdown()` function if you recover a panic, but the library does nothing like this automatically. 135 | 136 | # why 3 stages? 137 | By limiting the design to "only" three stages enable you to clearly make design choices, and force you to run as many things as possible in parallel. With this you can write simple design docs. Lets look at a webserver example: 138 | 139 | * Preshutdown: Finish accepted requests, refuse new ones. 140 | * Stage 1: Notify clients, flush data to database, notify upstream servers we are offline. 141 | * Stage 2: Flush database bulk writers, messages, close databases. (no database writes) 142 | * Stage 3: Flush/close log/metrics writer. (no log writes) 143 | 144 | My intention is that this makes the shutdown process easier to manage, and encourage more concurrency, because you don't create a long daisy-chain of events, and doesn't force you to look through all your code to insert a single event correctly. 145 | 146 | Don't think of the 3-stages as something that must do all stages of your shutdown. A single function call can of course (and is intended to) contain several "substages". Shutting down the database can easily be several stages, but you only register a single stage in the shutdown manager. The important part is that nothing else in the same stage can use the database. 147 | 148 | # examples 149 | 150 | There are examples in the [examples folder](https://github.com/klauspost/shutdown/tree/master/examples). 151 | 152 | # license 153 | 154 | This code is published under an MIT license. See LICENSE file for more information. 155 | -------------------------------------------------------------------------------- /examples/server-channel.go: -------------------------------------------------------------------------------- 1 | // +build ignore 2 | 3 | package main 4 | 5 | // This example shows a server that has logging in a separate goroutine 6 | // 7 | // When the server is shut down (via ctrl+c for instance), it will flush 8 | // all queued log entries and close the file. 9 | // 10 | // To execute, use 'go run server-channel.go' 11 | 12 | import ( 13 | "github.com/klauspost/shutdown" 14 | "io" 15 | "log" 16 | "net/http" 17 | "os" 18 | "syscall" 19 | ) 20 | 21 | func main() { 22 | // Make shutdown catch Ctrl+c and system terminate 23 | shutdown.OnSignal(0, os.Interrupt, syscall.SIGTERM) 24 | 25 | logStream = make(chan string, 100) 26 | 27 | // start a logger receiver 28 | go logger() 29 | 30 | // Start a webserver 31 | http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) { 32 | if shutdown.Lock() { 33 | logStream <- req.URL.String() + "\n" 34 | shutdown.Unlock() 35 | } else { 36 | log.Println("Already shutting down") 37 | } 38 | }) 39 | log.Fatal(http.ListenAndServe(":8080", nil)) 40 | } 41 | 42 | var logStream chan string 43 | 44 | func logger() { 45 | logFile, _ := os.Create("log.txt") 46 | 47 | // Get a nofification when we are at third stage of shutting down 48 | exit := shutdown.Third() 49 | for { 50 | select { 51 | case v := <-exit: 52 | log.Println("Flushing log...") 53 | finished := false 54 | for !finished { 55 | select { 56 | case m := <-logStream: 57 | _, _ = io.WriteString(logFile, m) 58 | default: 59 | finished = true 60 | } 61 | } 62 | log.Println("Closing log...") 63 | logFile.Close() 64 | // Signal we are done 65 | close(v) 66 | return 67 | case v := <-logStream: 68 | _, _ = io.WriteString(logFile, v) 69 | } 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /examples/server-example.go: -------------------------------------------------------------------------------- 1 | // +build ignore 2 | 3 | package main 4 | 5 | import ( 6 | "github.com/klauspost/shutdown" 7 | "io" 8 | "log" 9 | "net/http" 10 | "os" 11 | "syscall" 12 | "time" 13 | ) 14 | 15 | // This example shows a server that has message processing in a separate goroutine 16 | // 17 | // When the server is shut down (via ctrl+c for instance), it will notify an upstream 18 | // server, and all new incoming requests will get a 'StatusServiceUnavailable' (503) 19 | // response code. 20 | // 21 | // The server will finish all pending requests before shutdown is initiated, 22 | // so all requests are handled gracefully. 23 | // 24 | // To execute, use 'go run server-example.go' 25 | // 26 | // Open the server at http://localhost:8080 27 | // To shut down the server, go to http://localhost:8080/?shutdown=true or press ctrl+c 28 | 29 | // A Sample Webserver 30 | func HelloServer(w http.ResponseWriter, req *http.Request) { 31 | // Tracks all running requests 32 | if shutdown.Lock() { 33 | defer shutdown.Unlock() 34 | } else { 35 | // Shutdown has started, return that the service is unavailable 36 | w.WriteHeader(http.StatusServiceUnavailable) 37 | w.Write([]byte("Server is now shutting down")) 38 | return 39 | } 40 | 41 | if req.FormValue("shutdown") != "" { 42 | request <- "shutdown" 43 | log.Println("Requesting server shutdown") 44 | // We start the exit in a separate go-routine, otherwise this request will have 45 | // to wait for shutdown to be completed. 46 | go shutdown.Exit(0) 47 | } else { 48 | // Add artificial delay 49 | time.Sleep(time.Second * 5) 50 | request <- "greet" 51 | } 52 | io.WriteString(w, <-reply) 53 | } 54 | 55 | func main() { 56 | // Make shutdown catch Ctrl+c and system terminate 57 | shutdown.OnSignal(0, os.Interrupt, syscall.SIGTERM) 58 | 59 | // In the first stage we will make sure all request have finished 60 | shutdown.FirstFunc(func(interface{}) { 61 | log.Println("Notify upstream we are going offline") 62 | // TODO: Send a request upstream 63 | }, nil) 64 | 65 | // Start a service 66 | go dataLoop() 67 | 68 | // Start a webserver 69 | http.HandleFunc("/", HelloServer) 70 | log.Fatal(http.ListenAndServe(":8080", nil)) 71 | } 72 | 73 | var request = make(chan string) 74 | var reply = make(chan string) 75 | 76 | func dataLoop() { 77 | // We register for Second stage shutdown notification, 78 | // since we don't want to stop this service while requests are still being handled. 79 | end := shutdown.Second() 80 | for { 81 | select { 82 | case v := <-request: 83 | if v == "greet" { 84 | reply <- "hello world\n" 85 | } else if v == "shutdown" { 86 | reply <- "initiating server shutdown\n" 87 | } else { 88 | reply <- "unknown command\n" 89 | } 90 | case n := <-end: 91 | log.Println("Exiting data loop") 92 | close(request) 93 | close(reply) 94 | close(n) 95 | return 96 | } 97 | } 98 | } 99 | -------------------------------------------------------------------------------- /examples/simple-func.go: -------------------------------------------------------------------------------- 1 | // +build ignore 2 | 3 | package main 4 | 5 | import ( 6 | "github.com/klauspost/shutdown" 7 | "log" 8 | "net/http" 9 | "os" 10 | "syscall" 11 | ) 12 | 13 | // This example shows a server that has logging to a file 14 | // 15 | // When the webserver is closed, it will close the file when all requests have 16 | // been finished. 17 | // 18 | // In a real world, you would not want multiple goroutines writing to the same file 19 | // 20 | // To execute, use 'go run simple-func.go' 21 | 22 | // This is the function we would like to execute at shutdown. 23 | func closeFile(i interface{}) { 24 | f := i.(*os.File) 25 | log.Println("Closing", f.Name()+"...") 26 | f.Close() 27 | } 28 | 29 | func main() { 30 | // Make shutdown catch Ctrl+c and system terminate 31 | shutdown.OnSignal(0, os.Interrupt, syscall.SIGTERM) 32 | 33 | // Create a log file 34 | var logFile *os.File 35 | logFile, _ = os.Create("log.txt") 36 | 37 | // When shutdown is initiated, close the file 38 | shutdown.FirstFunc(closeFile, logFile) 39 | 40 | // Start a webserver 41 | http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) { 42 | // Get a lock, and write to the file if we get it. 43 | // While we have the lock the file will not be closed. 44 | if shutdown.Lock() { 45 | _, _ = logFile.WriteString(req.URL.String() + "\n") 46 | shutdown.Unlock() 47 | } 48 | }) 49 | log.Fatal(http.ListenAndServe(":8080", nil)) 50 | } 51 | -------------------------------------------------------------------------------- /http.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2015 Klaus Post, released under MIT License. See LICENSE file. 2 | 3 | package shutdown 4 | 5 | import ( 6 | "net/http" 7 | ) 8 | 9 | // WrapHandler will return an http Handler 10 | // That will lock shutdown until all have completed 11 | // and will return http.StatusServiceUnavailable if 12 | // shutdown has been initiated. 13 | func WrapHandler(h http.Handler) http.Handler { 14 | fn := func(w http.ResponseWriter, r *http.Request) { 15 | if !Lock() { 16 | w.WriteHeader(http.StatusServiceUnavailable) 17 | return 18 | } 19 | // We defer, so panics will not keep a lock 20 | defer Unlock() 21 | h.ServeHTTP(w, r) 22 | } 23 | return http.HandlerFunc(fn) 24 | } 25 | 26 | // WrapHandlerFunc will return an http.HandlerFunc 27 | // that will lock shutdown until all have completed. 28 | // The handler will return http.StatusServiceUnavailable if 29 | // shutdown has been initiated. 30 | func WrapHandlerFunc(h http.HandlerFunc) http.HandlerFunc { 31 | fn := func(w http.ResponseWriter, r *http.Request) { 32 | if !Lock() { 33 | w.WriteHeader(http.StatusServiceUnavailable) 34 | return 35 | } 36 | // We defer, so panics will not keep a lock 37 | defer Unlock() 38 | h(w, r) 39 | } 40 | return http.HandlerFunc(fn) 41 | } 42 | -------------------------------------------------------------------------------- /http_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2015 Klaus Post, released under MIT License. See LICENSE file. 2 | 3 | package shutdown 4 | 5 | import ( 6 | "bytes" 7 | "fmt" 8 | "html" 9 | "net/http" 10 | "net/http/httptest" 11 | "os" 12 | "syscall" 13 | "testing" 14 | "time" 15 | ) 16 | 17 | // This example creates a custom function handler 18 | // and wraps the handler, so all request will 19 | // finish before shutdown is started. 20 | // 21 | // If requests take too long to finish (see the shutdown will proceed 22 | // and clients will be disconnected when the server shuts down. 23 | // To modify the timeout use SetTimeoutN(Preshutdown, duration) 24 | func ExampleWrapHandlerFunc() { 25 | // Set a custom timeout, if the 5 second default doesn't fit your needs. 26 | SetTimeoutN(Preshutdown, time.Second*30) 27 | // Catch OS signals 28 | OnSignal(0, os.Interrupt, syscall.SIGTERM) 29 | 30 | // Example handler function 31 | fn := func(w http.ResponseWriter, r *http.Request) { 32 | fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path)) 33 | } 34 | 35 | // Wrap the handler function 36 | http.HandleFunc("/", WrapHandlerFunc(fn)) 37 | 38 | // Start the server 39 | http.ListenAndServe(":8080", nil) 40 | } 41 | 42 | // This example creates a fileserver 43 | // and wraps the handler, so all request will 44 | // finish before shutdown is started. 45 | // 46 | // If requests take too long to finish the shutdown will proceed 47 | // and clients will be disconnected when the server shuts down. 48 | // To modify the timeout use SetTimeoutN(Preshutdown, duration) 49 | func ExampleWrapHandler() { 50 | // Set a custom timeout, if the 5 second default doesn't fit your needs. 51 | SetTimeoutN(Preshutdown, time.Second*30) 52 | // Catch OS signals 53 | OnSignal(0, os.Interrupt, syscall.SIGTERM) 54 | 55 | // Create a fileserver handler 56 | fh := http.FileServer(http.Dir("/examples")) 57 | 58 | // Wrap the handler function 59 | http.Handle("/", WrapHandler(fh)) 60 | 61 | // Start the server 62 | http.ListenAndServe(":8080", nil) 63 | } 64 | 65 | func TestWrapHandlerBasic(t *testing.T) { 66 | reset() 67 | defer close(startTimer(t)) 68 | var finished = false 69 | fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 70 | finished = true 71 | }) 72 | 73 | wrapped := WrapHandler(fn) 74 | res := httptest.NewRecorder() 75 | req, _ := http.NewRequest("", "", bytes.NewBufferString("")) 76 | wrapped.ServeHTTP(res, req) 77 | if res.Code == http.StatusServiceUnavailable { 78 | t.Fatal("Expected result code NOT to be", http.StatusServiceUnavailable, "got", res.Code) 79 | } 80 | if !finished { 81 | t.Fatal("Handler was not executed") 82 | } 83 | 84 | Shutdown() 85 | finished = false 86 | res = httptest.NewRecorder() 87 | wrapped.ServeHTTP(res, req) 88 | if res.Code != http.StatusServiceUnavailable { 89 | t.Fatal("Expected result code to be", http.StatusServiceUnavailable, " got", res.Code) 90 | } 91 | if finished { 92 | t.Fatal("Unexpected execution of funtion") 93 | } 94 | } 95 | 96 | func TestWrapHandlerFuncBasic(t *testing.T) { 97 | reset() 98 | defer close(startTimer(t)) 99 | var finished = false 100 | fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 101 | finished = true 102 | }) 103 | 104 | wrapped := WrapHandlerFunc(fn) 105 | res := httptest.NewRecorder() 106 | req, _ := http.NewRequest("", "", bytes.NewBufferString("")) 107 | wrapped(res, req) 108 | if res.Code == http.StatusServiceUnavailable { 109 | t.Fatal("Expected result code NOT to be", http.StatusServiceUnavailable, "got", res.Code) 110 | } 111 | if !finished { 112 | t.Fatal("Handler was not executed") 113 | } 114 | 115 | Shutdown() 116 | finished = false 117 | res = httptest.NewRecorder() 118 | wrapped(res, req) 119 | if res.Code != http.StatusServiceUnavailable { 120 | t.Fatal("Expected result code to be", http.StatusServiceUnavailable, " got", res.Code) 121 | } 122 | if finished { 123 | t.Fatal("Unexpected execution of funtion") 124 | } 125 | } 126 | 127 | // Test if panics locks shutdown. 128 | func TestWrapHandlerPanic(t *testing.T) { 129 | reset() 130 | SetTimeout(time.Second) 131 | defer close(startTimer(t)) 132 | fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 133 | panic("test panic") 134 | }) 135 | 136 | wrapped := WrapHandler(fn) 137 | res := httptest.NewRecorder() 138 | req, _ := http.NewRequest("", "", bytes.NewBufferString("")) 139 | func() { 140 | defer func() { 141 | recover() 142 | }() 143 | wrapped.ServeHTTP(res, req) 144 | }() 145 | 146 | // There should be no locks held, so it should finish immediately 147 | tn := time.Now() 148 | Shutdown() 149 | dur := time.Now().Sub(tn) 150 | if dur > time.Millisecond*500 { 151 | t.Fatalf("timeout time was unexpected:%v", time.Now().Sub(tn)) 152 | } 153 | } 154 | 155 | // Test if panics locks shutdown. 156 | func TestWrapHandlerFuncPanic(t *testing.T) { 157 | reset() 158 | SetTimeout(time.Millisecond * 200) 159 | defer close(startTimer(t)) 160 | fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 161 | panic("test panic") 162 | }) 163 | 164 | wrapped := WrapHandlerFunc(fn) 165 | res := httptest.NewRecorder() 166 | req, _ := http.NewRequest("", "", bytes.NewBufferString("")) 167 | func() { 168 | defer func() { 169 | recover() 170 | }() 171 | wrapped(res, req) 172 | }() 173 | 174 | // There should be no locks held, so it should finish immediately 175 | tn := time.Now() 176 | Shutdown() 177 | dur := time.Now().Sub(tn) 178 | if dur > time.Millisecond*100 { 179 | t.Fatalf("timeout time was unexpected:%v", time.Now().Sub(tn)) 180 | } 181 | } 182 | 183 | // Tests that shutdown doesn't complete until handler function has returned 184 | func TestWrapHandlerOrder(t *testing.T) { 185 | reset() 186 | defer close(startTimer(t)) 187 | var finished = make(chan bool) 188 | var wait = make(chan bool) 189 | var waiting = make(chan bool) 190 | fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 191 | close(waiting) 192 | <-wait 193 | }) 194 | 195 | wrapped := WrapHandler(fn) 196 | 197 | go func() { 198 | res := httptest.NewRecorder() 199 | req, _ := http.NewRequest("", "", bytes.NewBufferString("")) 200 | wrapped.ServeHTTP(res, req) 201 | close(finished) 202 | }() 203 | 204 | release := time.After(time.Millisecond * 100) 205 | completed := make(chan bool) 206 | testOK := make(chan bool) 207 | go func() { 208 | select { 209 | case <-release: 210 | select { 211 | case <-finished: 212 | panic("Shutdown was already finished") 213 | case <-completed: 214 | panic("Shutdown had already completed") 215 | default: 216 | } 217 | close(wait) 218 | close(testOK) 219 | } 220 | }() 221 | <-waiting 222 | tn := time.Now() 223 | Shutdown() 224 | dur := time.Now().Sub(tn) 225 | if dur > time.Millisecond*400 { 226 | t.Fatalf("timeout time was unexpected:%v", time.Now().Sub(tn)) 227 | } 228 | close(completed) 229 | // We should make sure the release has run before exiting 230 | <-testOK 231 | 232 | select { 233 | case <-finished: 234 | default: 235 | t.Fatal("Function had not finished") 236 | } 237 | } 238 | 239 | // Tests that shutdown doesn't complete until handler function has returned 240 | func TestWrapHandlerFuncOrder(t *testing.T) { 241 | reset() 242 | defer close(startTimer(t)) 243 | var finished = make(chan bool) 244 | var wait = make(chan bool) 245 | var waiting = make(chan bool) 246 | fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 247 | close(waiting) 248 | <-wait 249 | }) 250 | 251 | wrapped := WrapHandlerFunc(fn) 252 | 253 | go func() { 254 | res := httptest.NewRecorder() 255 | req, _ := http.NewRequest("", "", bytes.NewBufferString("")) 256 | wrapped(res, req) 257 | close(finished) 258 | }() 259 | 260 | release := time.After(time.Millisecond * 100) 261 | completed := make(chan bool) 262 | testOK := make(chan bool) 263 | go func() { 264 | select { 265 | case <-release: 266 | select { 267 | case <-finished: 268 | panic("Shutdown was already finished") 269 | case <-completed: 270 | panic("Shutdown had already completed") 271 | default: 272 | } 273 | close(wait) 274 | close(testOK) 275 | } 276 | }() 277 | <-waiting 278 | tn := time.Now() 279 | Shutdown() 280 | dur := time.Now().Sub(tn) 281 | if dur > time.Millisecond*400 { 282 | t.Fatalf("timeout time was unexpected:%v", time.Now().Sub(tn)) 283 | } 284 | close(completed) 285 | // We should make sure the release has run before exiting 286 | <-testOK 287 | 288 | select { 289 | case <-finished: 290 | default: 291 | t.Fatal("Function had not finished") 292 | } 293 | } 294 | -------------------------------------------------------------------------------- /shutdown.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2015 Klaus Post, released under MIT License. See LICENSE file. 2 | 3 | // Package shutdown provides management of your shutdown process. 4 | // 5 | // The package will enable you to get notifications for your application and handle the shutdown process. 6 | // 7 | // See more information about the how to use it in the README.md file 8 | // 9 | // Package home: https://github.com/klauspost/shutdown 10 | // 11 | // This package has been updated to a new version. If you are starting a new project, use the new package: 12 | // 13 | // * Package home: https://github.com/klauspost/shutdown2 14 | // * Godoc: https://godoc.org/github.com/klauspost/shutdown2 15 | // 16 | // Version 2 mainly contains minor adjustments to the API to make it a bit easier to use. 17 | // 18 | // This package remains here to maintain compatibility. 19 | package shutdown 20 | 21 | import ( 22 | "log" 23 | "os" 24 | "os/signal" 25 | "sync" 26 | "time" 27 | ) 28 | 29 | // Valid values for this is exported as variables. 30 | type Stage struct { 31 | n int 32 | } 33 | 34 | // Logger used for output. 35 | // This can be exchanged with your own. 36 | var Logger = log.New(os.Stderr, "[shutdown]: ", log.LstdFlags) 37 | 38 | var Preshutdown = Stage{0} // Indicates stage when waiting for locks to be released. 39 | var Stage1 = Stage{1} // Indicates first stage of timeouts. 40 | var Stage2 = Stage{2} // Indicates second stage of timeouts. 41 | var Stage3 = Stage{3} // Indicates third stage of timeouts. 42 | 43 | // Notifier is a channel, that will be sent a channel 44 | // once the application shuts down. 45 | // When you have performed your shutdown actions close the channel you are given. 46 | type Notifier chan chan struct{} 47 | 48 | type fnNotify struct { 49 | client Notifier 50 | internal Notifier 51 | cancel chan struct{} 52 | } 53 | 54 | var sqM sync.Mutex // Mutex for below 55 | var shutdownQueue [4][]Notifier 56 | var shutdownFnQueue [4][]fnNotify 57 | var shutdownFinished = make(chan struct{}, 0) // Closed when shutdown has finished 58 | 59 | var srM sync.RWMutex // Mutex for below 60 | var shutdownRequested = false 61 | var timeouts = [4]time.Duration{5 * time.Second, 5 * time.Second, 5 * time.Second, 5 * time.Second} 62 | 63 | // SetTimeout sets maximum delay to wait for each stage to finish. 64 | // When the timeout has expired for a stage the next stage will be initiated. 65 | func SetTimeout(d time.Duration) { 66 | srM.Lock() 67 | for i := range timeouts { 68 | timeouts[i] = d 69 | } 70 | srM.Unlock() 71 | } 72 | 73 | // SetTimeoutN set maximum delay to wait for a specific stage to finish. 74 | // When the timeout expired for a stage the next stage will be initiated. 75 | // The stage can be obtained by using the exported variables called 'Stage1, etc. 76 | func SetTimeoutN(s Stage, d time.Duration) { 77 | srM.Lock() 78 | timeouts[s.n] = d 79 | srM.Unlock() 80 | } 81 | 82 | // Cancel a Notifier. 83 | // This will remove a notifier from the shutdown queue, 84 | // and it will not be signalled when shutdown starts. 85 | // If the shutdown has already started this will not have any effect. 86 | func (s *Notifier) Cancel() { 87 | srM.RLock() 88 | if shutdownRequested { 89 | srM.RUnlock() 90 | return 91 | } 92 | srM.RUnlock() 93 | sqM.Lock() 94 | var a chan chan struct{} 95 | var b chan chan struct{} 96 | a = *s 97 | for n, sdq := range shutdownQueue { 98 | for i, qi := range sdq { 99 | b = qi 100 | if a == b { 101 | shutdownQueue[n] = append(shutdownQueue[n][:i], shutdownQueue[n][i+1:]...) 102 | } 103 | } 104 | for i, fn := range shutdownFnQueue[n] { 105 | b = fn.client 106 | if a == b { 107 | // Find the matching internal and remove that. 108 | for i := range shutdownQueue[n] { 109 | b = shutdownQueue[n][i] 110 | if fn.internal == b { 111 | shutdownQueue[n] = append(shutdownQueue[n][:i], shutdownQueue[n][i+1:]...) 112 | } 113 | } 114 | // Cancel, so the goroutine exits. 115 | close(fn.cancel) 116 | // Remove this 117 | shutdownFnQueue[n] = append(shutdownFnQueue[n][:i], shutdownFnQueue[n][i+1:]...) 118 | } 119 | } 120 | } 121 | sqM.Unlock() 122 | } 123 | 124 | // PreShutdown will return a Notifier that will be fired as soon as the shutdown 125 | // is signalled, before locks are released. 126 | // This allows to for instance send signals to upstream servers not to send more requests. 127 | func PreShutdown() Notifier { 128 | return onShutdown(0) 129 | } 130 | 131 | type ShutdownFn func(interface{}) 132 | 133 | // PreShutdownFunc registers a function that will be called as soon as the shutdown 134 | // is signalled, before locks are released. 135 | // This allows to for instance send signals to upstream servers not to send more requests. 136 | func PreShutdownFunc(fn ShutdownFn, v interface{}) Notifier { 137 | return onFunc(0, fn, v) 138 | } 139 | 140 | // First returns a notifier that will be called in the first stage of shutdowns 141 | func First() Notifier { 142 | return onShutdown(1) 143 | } 144 | 145 | // FirstFunc executes a function in the first stage of the shutdown 146 | func FirstFunc(fn ShutdownFn, v interface{}) Notifier { 147 | return onFunc(1, fn, v) 148 | } 149 | 150 | // Second returns a notifier that will be called in the second stage of shutdowns 151 | func Second() Notifier { 152 | return onShutdown(2) 153 | } 154 | 155 | // SecondFunc executes a function in the second stage of the shutdown 156 | func SecondFunc(fn ShutdownFn, v interface{}) Notifier { 157 | return onFunc(2, fn, v) 158 | } 159 | 160 | // Third returns a notifier that will be called in the third stage of shutdowns 161 | func Third() Notifier { 162 | return onShutdown(3) 163 | } 164 | 165 | // ThirdFunc executes a function in the third stage of the shutdown 166 | // The returned Notifier is only really useful for cancelling the shutdown function 167 | func ThirdFunc(fn ShutdownFn, v interface{}) Notifier { 168 | return onFunc(3, fn, v) 169 | } 170 | 171 | // Create a function notifier. 172 | func onFunc(prio int, fn ShutdownFn, i interface{}) Notifier { 173 | f := fnNotify{ 174 | internal: onShutdown(prio), 175 | cancel: make(chan struct{}), 176 | client: make(Notifier, 1), 177 | } 178 | go func() { 179 | select { 180 | case <-f.cancel: 181 | return 182 | case c := <-f.internal: 183 | { 184 | defer func() { 185 | if r := recover(); r != nil { 186 | Logger.Println("Panic in shutdown function:", r) 187 | } 188 | if c != nil { 189 | close(c) 190 | } 191 | }() 192 | fn(i) 193 | } 194 | } 195 | }() 196 | sqM.Lock() 197 | shutdownFnQueue[prio] = append(shutdownFnQueue[prio], f) 198 | sqM.Unlock() 199 | return f.client 200 | } 201 | 202 | // onShutdown will request a shutdown notifier. 203 | func onShutdown(prio int) Notifier { 204 | sqM.Lock() 205 | n := make(Notifier, 1) 206 | shutdownQueue[prio] = append(shutdownQueue[prio], n) 207 | sqM.Unlock() 208 | return n 209 | } 210 | 211 | // OnSignal will start the shutdown when any of the given signals arrive 212 | // 213 | // A good shutdown default is 214 | // shutdown.OnSignal(0, os.Interrupt, syscall.SIGTERM) 215 | // which will do shutdown on Ctrl+C and when the program is terminated. 216 | func OnSignal(exitCode int, sig ...os.Signal) { 217 | // capture signal and shut down. 218 | c := make(chan os.Signal, 1) 219 | signal.Notify(c, sig...) 220 | go func() { 221 | for _ = range c { 222 | Shutdown() 223 | os.Exit(exitCode) 224 | } 225 | }() 226 | } 227 | 228 | // Exit performs shutdown operations and exits with the given exit code. 229 | func Exit(code int) { 230 | Shutdown() 231 | os.Exit(code) 232 | } 233 | 234 | // Shutdown will signal all notifiers in three stages. 235 | // It will first check that all locks have been released - see Lock() 236 | func Shutdown() { 237 | srM.Lock() 238 | if shutdownRequested { 239 | srM.Unlock() 240 | // Wait till shutdown finished 241 | <-shutdownFinished 242 | return 243 | } 244 | shutdownRequested = true 245 | srM.Unlock() 246 | 247 | // Add a pre-shutdown function that waits for all locks to be released. 248 | PreShutdownFunc(func(interface{}) { 249 | srM.Lock() 250 | wait := wg 251 | srM.Unlock() 252 | wait.Wait() 253 | }, nil) 254 | 255 | sqM.Lock() 256 | for stage := 0; stage < 4; stage++ { 257 | srM.Lock() 258 | to := timeouts[stage] 259 | srM.Unlock() 260 | 261 | queue := shutdownQueue[stage] 262 | if len(queue) == 0 { 263 | continue 264 | } 265 | if stage == 0 { 266 | Logger.Println("Initiating shutdown") 267 | } else { 268 | Logger.Println("Shutdown stage", stage) 269 | } 270 | wait := make([]chan struct{}, len(queue)) 271 | 272 | // Send notification to all waiting 273 | for i := range queue { 274 | wait[i] = make(chan struct{}) 275 | queue[i] <- wait[i] 276 | } 277 | 278 | // Send notification to all function notifiers, but don't wait 279 | for _, notifier := range shutdownFnQueue[stage] { 280 | notifier.client <- make(chan struct{}) 281 | close(notifier.client) 282 | } 283 | 284 | // We don't lock while we are waiting for notifiers to return 285 | sqM.Unlock() 286 | 287 | // Wait for all to return, no more than the shutdown delay 288 | timeout := time.After(to) 289 | 290 | brwait: 291 | for i := range wait { 292 | select { 293 | case <-wait[i]: 294 | case <-timeout: 295 | Logger.Println("timeout waiting to shutdown, forcing shutdown") 296 | break brwait 297 | } 298 | } 299 | sqM.Lock() 300 | } 301 | // Reset - mainly for tests. 302 | shutdownQueue = [4][]Notifier{} 303 | shutdownFnQueue = [4][]fnNotify{} 304 | close(shutdownFinished) 305 | sqM.Unlock() 306 | } 307 | 308 | // Started returns true if shutdown has been started. 309 | // Note that shutdown can have been started before you check the value. 310 | func Started() bool { 311 | srM.RLock() 312 | started := shutdownRequested 313 | srM.RUnlock() 314 | return started 315 | } 316 | 317 | var wg *sync.WaitGroup 318 | 319 | func init() { 320 | wg = &sync.WaitGroup{} 321 | } 322 | 323 | // Wait will wait until shutdown has finished. 324 | // This can be used to keep a main function from exiting 325 | // until shutdown has been called, either by a goroutine 326 | // or a signal. 327 | func Wait() { 328 | <-shutdownFinished 329 | } 330 | 331 | // Lock will signal that you have a function running, 332 | // that you do not want to be interrupted by a shutdown. 333 | // 334 | // If the function returns false shutdown has already been initiated, 335 | // and you did not get a lock. You should therefore not call Unlock. 336 | // 337 | // If the function returned true, you must call Unlock() once to release the lock. 338 | // 339 | // You should not hold a lock when you start a shutdown. 340 | func Lock() bool { 341 | srM.RLock() 342 | s := shutdownRequested 343 | if !s { 344 | wg.Add(1) 345 | } 346 | srM.RUnlock() 347 | return !s 348 | } 349 | 350 | // Unlock will release a shutdown lock. 351 | // This may only be called if you have previously called Lock and it has 352 | // returned true 353 | func Unlock() { 354 | wg.Done() 355 | } 356 | -------------------------------------------------------------------------------- /shutdown_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2015 Klaus Post, released under MIT License. See LICENSE file. 2 | 3 | package shutdown 4 | 5 | import ( 6 | "fmt" 7 | "net/http" 8 | "sync" 9 | "testing" 10 | "time" 11 | ) 12 | 13 | func reset() { 14 | SetTimeout(1 * time.Second) 15 | sqM.Lock() 16 | defer sqM.Unlock() 17 | srM.Lock() 18 | defer srM.Unlock() 19 | shutdownRequested = false 20 | wg = &sync.WaitGroup{} 21 | shutdownQueue = [4][]Notifier{} 22 | shutdownFnQueue = [4][]fnNotify{} 23 | shutdownFinished = make(chan struct{}) 24 | } 25 | 26 | func startTimer(t *testing.T) chan struct{} { 27 | finished := make(chan struct{}, 0) 28 | srM.RLock() 29 | var to time.Duration 30 | for i := range timeouts { 31 | to += timeouts[i] 32 | } 33 | srM.RUnlock() 34 | // Add some extra time. 35 | toc := time.After((to * 10) / 9) 36 | go func() { 37 | select { 38 | case <-toc: 39 | panic("unexpected timeout while running test") 40 | return 41 | case <-finished: 42 | return 43 | 44 | } 45 | }() 46 | return finished 47 | } 48 | 49 | func TestBasic(t *testing.T) { 50 | reset() 51 | defer close(startTimer(t)) 52 | f := First() 53 | ok := false 54 | go func() { 55 | select { 56 | case n := <-f: 57 | ok = true 58 | close(n) 59 | } 60 | }() 61 | Shutdown() 62 | if !ok { 63 | t.Fatal("did not get expected shutdown signal") 64 | } 65 | if !Started() { 66 | t.Fatal("shutdown not marked started") 67 | } 68 | // Should just return at once. 69 | Shutdown() 70 | // Should also return at once. 71 | Wait() 72 | } 73 | 74 | func TestPreShutdown(t *testing.T) { 75 | reset() 76 | defer close(startTimer(t)) 77 | f := PreShutdown() 78 | ok := false 79 | Lock() 80 | go func() { 81 | select { 82 | case n := <-f: 83 | ok = true 84 | Unlock() 85 | close(n) 86 | } 87 | }() 88 | tn := time.Now() 89 | Shutdown() 90 | dur := time.Now().Sub(tn) 91 | if dur > time.Second { 92 | t.Fatalf("timeout time was hit unexpected:%v", time.Now().Sub(tn)) 93 | } 94 | 95 | if !ok { 96 | t.Fatal("did not get expected shutdown signal") 97 | } 98 | if !Started() { 99 | t.Fatal("shutdown not marked started") 100 | } 101 | } 102 | 103 | func TestCancel(t *testing.T) { 104 | reset() 105 | defer close(startTimer(t)) 106 | f := First() 107 | ok := false 108 | go func() { 109 | select { 110 | case n := <-f: 111 | ok = true 112 | close(n) 113 | } 114 | }() 115 | f.Cancel() 116 | Shutdown() 117 | if ok { 118 | t.Fatal("got unexpected shutdown signal") 119 | } 120 | } 121 | 122 | func TestCancel2(t *testing.T) { 123 | reset() 124 | defer close(startTimer(t)) 125 | f2 := First() 126 | f := First() 127 | var ok, ok2 bool 128 | 129 | go func() { 130 | select { 131 | case n := <-f: 132 | ok = true 133 | close(n) 134 | } 135 | }() 136 | go func() { 137 | select { 138 | case n := <-f2: 139 | ok2 = true 140 | close(n) 141 | } 142 | }() 143 | f.Cancel() 144 | Shutdown() 145 | if ok { 146 | t.Fatal("got unexpected shutdown signal") 147 | } 148 | if !ok2 { 149 | t.Fatal("missing shutdown signal") 150 | } 151 | } 152 | 153 | func TestWait(t *testing.T) { 154 | reset() 155 | defer close(startTimer(t)) 156 | ok := make(chan bool) 157 | go func() { 158 | Wait() 159 | close(ok) 160 | }() 161 | // Wait a little - enough to fail very often. 162 | time.Sleep(time.Millisecond * 10) 163 | 164 | select { 165 | case <-ok: 166 | t.Fatal("Wait returned before shutdown finished") 167 | default: 168 | } 169 | 170 | Shutdown() 171 | 172 | // ok should return, otherwise we wait for timeout, which will fail the test 173 | <-ok 174 | } 175 | 176 | func TestTimeout(t *testing.T) { 177 | reset() 178 | SetTimeout(time.Millisecond * 100) 179 | defer close(startTimer(t)) 180 | f := First() 181 | go func() { 182 | select { 183 | case <-f: 184 | } 185 | }() 186 | tn := time.Now() 187 | Shutdown() 188 | dur := time.Now().Sub(tn) 189 | if dur > time.Second || dur < time.Millisecond*50 { 190 | t.Fatalf("timeout time was unexpected:%v", time.Now().Sub(tn)) 191 | } 192 | if !Started() { 193 | t.Fatal("got unexpected shutdown signal") 194 | } 195 | } 196 | 197 | func TestTimeoutN(t *testing.T) { 198 | reset() 199 | SetTimeout(time.Second * 2) 200 | SetTimeoutN(Stage1, time.Millisecond*100) 201 | defer close(startTimer(t)) 202 | f := First() 203 | go func() { 204 | select { 205 | case <-f: 206 | } 207 | }() 208 | tn := time.Now() 209 | Shutdown() 210 | dur := time.Now().Sub(tn) 211 | if dur > time.Second || dur < time.Millisecond*50 { 212 | t.Fatalf("timeout time was unexpected:%v", time.Now().Sub(tn)) 213 | } 214 | if !Started() { 215 | t.Fatal("got unexpected shutdown signal") 216 | } 217 | } 218 | 219 | func TestTimeoutN2(t *testing.T) { 220 | reset() 221 | SetTimeout(time.Millisecond * 100) 222 | SetTimeoutN(Stage2, time.Second*2) 223 | defer close(startTimer(t)) 224 | f := First() 225 | go func() { 226 | select { 227 | case <-f: 228 | } 229 | }() 230 | tn := time.Now() 231 | Shutdown() 232 | dur := time.Now().Sub(tn) 233 | if dur > time.Second || dur < time.Millisecond*50 { 234 | t.Fatalf("timeout time was unexpected:%v", time.Now().Sub(tn)) 235 | } 236 | if !Started() { 237 | t.Fatal("got unexpected shutdown signal") 238 | } 239 | } 240 | 241 | func TestLock(t *testing.T) { 242 | reset() 243 | defer close(startTimer(t)) 244 | f := First() 245 | ok := false 246 | go func() { 247 | select { 248 | case n := <-f: 249 | ok = true 250 | close(n) 251 | } 252 | }() 253 | got := Lock() 254 | if !got { 255 | t.Fatal("Unable to aquire lock") 256 | } 257 | Unlock() 258 | for i := 0; i < 10; i++ { 259 | go func() { 260 | if Lock() { 261 | time.Sleep(time.Second) 262 | Unlock() 263 | } 264 | }() 265 | } 266 | Shutdown() 267 | if !ok { 268 | t.Fatal("shutdown signal not received") 269 | } 270 | if !Started() { 271 | t.Fatal("expected that shutdown had started") 272 | } 273 | } 274 | 275 | func TestLockUnrelease(t *testing.T) { 276 | reset() 277 | defer close(startTimer(t)) 278 | SetTimeout(time.Millisecond * 100) 279 | got := Lock() 280 | if !got { 281 | t.Fatal("Unable to aquire lock") 282 | } 283 | tn := time.Now() 284 | Shutdown() 285 | dur := time.Now().Sub(tn) 286 | if dur > time.Second || dur < time.Millisecond*50 { 287 | t.Fatalf("timeout time was unexpected:%v", time.Now().Sub(tn)) 288 | } 289 | if !Started() { 290 | t.Fatal("expected that shutdown had started") 291 | } 292 | // Unlock to be nice 293 | Unlock() 294 | } 295 | 296 | func TestOrder(t *testing.T) { 297 | reset() 298 | defer close(startTimer(t)) 299 | 300 | t3 := Third() 301 | if Started() { 302 | t.Fatal("shutdown started unexpectedly") 303 | } 304 | 305 | t2 := Second() 306 | if Started() { 307 | t.Fatal("shutdown started unexpectedly") 308 | } 309 | 310 | t1 := First() 311 | if Started() { 312 | t.Fatal("shutdown started unexpectedly") 313 | } 314 | 315 | t0 := PreShutdown() 316 | if Started() { 317 | t.Fatal("shutdown started unexpectedly") 318 | } 319 | 320 | var ok0, ok1, ok2, ok3 bool 321 | go func() { 322 | for { 323 | select { 324 | //t0 must be first 325 | case n := <-t0: 326 | if ok0 || ok1 || ok2 || ok3 { 327 | t.Fatal("unexpected order", ok0, ok1, ok2, ok3) 328 | } 329 | ok0 = true 330 | close(n) 331 | case n := <-t1: 332 | if !ok0 || ok1 || ok2 || ok3 { 333 | t.Fatal("unexpected order", ok0, ok1, ok2, ok3) 334 | } 335 | ok1 = true 336 | close(n) 337 | case n := <-t2: 338 | if !ok0 || !ok1 || ok2 || ok3 { 339 | t.Fatal("unexpected order", ok0, ok1, ok2, ok3) 340 | } 341 | ok2 = true 342 | close(n) 343 | case n := <-t3: 344 | if !ok0 || !ok1 || !ok2 || ok3 { 345 | t.Fatal("unexpected order", ok0, ok1, ok2, ok3) 346 | } 347 | ok3 = true 348 | close(n) 349 | return 350 | } 351 | } 352 | }() 353 | if ok0 || ok1 || ok2 || ok3 { 354 | t.Fatal("shutdown has already happened", ok0, ok1, ok2, ok3) 355 | } 356 | 357 | Shutdown() 358 | if !ok0 || !ok1 || !ok2 || !ok3 { 359 | t.Fatal("did not get expected shutdown signal", ok0, ok1, ok2, ok3) 360 | } 361 | } 362 | 363 | func TestRecursive(t *testing.T) { 364 | reset() 365 | defer close(startTimer(t)) 366 | 367 | if Started() { 368 | t.Fatal("shutdown started unexpectedly") 369 | } 370 | 371 | t1 := First() 372 | if Started() { 373 | t.Fatal("shutdown started unexpectedly") 374 | } 375 | 376 | var ok1, ok2, ok3 bool 377 | go func() { 378 | for { 379 | select { 380 | case n := <-t1: 381 | ok1 = true 382 | t2 := Second() 383 | close(n) 384 | select { 385 | case n := <-t2: 386 | ok2 = true 387 | t3 := Third() 388 | close(n) 389 | select { 390 | case n := <-t3: 391 | ok3 = true 392 | close(n) 393 | return 394 | } 395 | } 396 | } 397 | } 398 | }() 399 | if ok1 || ok2 || ok3 { 400 | t.Fatal("shutdown has already happened", ok1, ok2, ok3) 401 | } 402 | 403 | Shutdown() 404 | if !ok1 || !ok2 || !ok3 { 405 | t.Fatal("did not get expected shutdown signal", ok1, ok2, ok3) 406 | } 407 | } 408 | 409 | func TestBasicFn(t *testing.T) { 410 | reset() 411 | defer close(startTimer(t)) 412 | gotcall := false 413 | 414 | // Register a function 415 | _ = FirstFunc(func(i interface{}) { 416 | gotcall = i.(bool) 417 | }, true) 418 | 419 | // Start shutdown 420 | Shutdown() 421 | if !gotcall { 422 | t.Fatal("did not get expected shutdown signal") 423 | } 424 | } 425 | 426 | func setBool(i interface{}) { 427 | set := i.(*bool) 428 | *set = true 429 | } 430 | 431 | func TestFnOrder(t *testing.T) { 432 | reset() 433 | defer close(startTimer(t)) 434 | 435 | var ok1, ok2, ok3 bool 436 | _ = ThirdFunc(setBool, &ok3) 437 | if Started() { 438 | t.Fatal("shutdown started unexpectedly") 439 | } 440 | 441 | _ = SecondFunc(setBool, &ok2) 442 | if Started() { 443 | t.Fatal("shutdown started unexpectedly") 444 | } 445 | 446 | _ = FirstFunc(setBool, &ok1) 447 | if Started() { 448 | t.Fatal("shutdown started unexpectedly") 449 | } 450 | 451 | if ok1 || ok2 || ok3 { 452 | t.Fatal("shutdown has already happened", ok1, ok2, ok3) 453 | } 454 | 455 | Shutdown() 456 | 457 | if !ok1 || !ok2 || !ok3 { 458 | t.Fatal("did not get expected shutdown signal", ok1, ok2, ok3) 459 | } 460 | } 461 | 462 | func TestFnRecursive(t *testing.T) { 463 | reset() 464 | defer close(startTimer(t)) 465 | 466 | var ok1, ok2, ok3 bool 467 | 468 | _ = FirstFunc(func(i interface{}) { 469 | set := i.(*bool) 470 | *set = true 471 | _ = SecondFunc(func(i interface{}) { 472 | set := i.(*bool) 473 | *set = true 474 | _ = ThirdFunc(func(i interface{}) { 475 | set := i.(*bool) 476 | *set = true 477 | }, &ok3) 478 | }, &ok2) 479 | }, &ok1) 480 | 481 | if Started() { 482 | t.Fatal("shutdown started unexpectedly") 483 | } 484 | 485 | if ok1 || ok2 || ok3 { 486 | t.Fatal("shutdown has already happened", ok1, ok2, ok3) 487 | } 488 | 489 | Shutdown() 490 | 491 | if !ok1 || !ok2 || !ok3 { 492 | t.Fatal("did not get expected shutdown signal", ok1, ok2, ok3) 493 | } 494 | } 495 | 496 | // When setting First or Second inside stage three they should be ignored. 497 | func TestFnRecursiveRev(t *testing.T) { 498 | reset() 499 | defer close(startTimer(t)) 500 | 501 | var ok1, ok2, ok3 bool 502 | 503 | _ = ThirdFunc(func(i interface{}) { 504 | set := i.(*bool) 505 | *set = true 506 | _ = SecondFunc(func(i interface{}) { 507 | set := i.(*bool) 508 | *set = true 509 | }, &ok2) 510 | _ = FirstFunc(func(i interface{}) { 511 | set := i.(*bool) 512 | *set = true 513 | }, &ok1) 514 | }, &ok3) 515 | 516 | if Started() { 517 | t.Fatal("shutdown started unexpectedly") 518 | } 519 | 520 | if ok1 || ok2 || ok3 { 521 | t.Fatal("shutdown has already happened", ok1, ok2, ok3) 522 | } 523 | 524 | Shutdown() 525 | 526 | if ok1 || ok2 || !ok3 { 527 | t.Fatal("did not get expected shutdown signal", ok1, ok2, ok3) 528 | } 529 | } 530 | 531 | func TestFnCancel(t *testing.T) { 532 | reset() 533 | defer close(startTimer(t)) 534 | var g0, g1, g2, g3 bool 535 | 536 | // Register a function 537 | notp := PreShutdownFunc(func(i interface{}) { 538 | g0 = i.(bool) 539 | }, true) 540 | not1 := FirstFunc(func(i interface{}) { 541 | g1 = i.(bool) 542 | }, true) 543 | not2 := SecondFunc(func(i interface{}) { 544 | g2 = i.(bool) 545 | }, true) 546 | not3 := ThirdFunc(func(i interface{}) { 547 | g3 = i.(bool) 548 | }, true) 549 | 550 | notp.Cancel() 551 | not1.Cancel() 552 | not2.Cancel() 553 | not3.Cancel() 554 | 555 | // Start shutdown 556 | Shutdown() 557 | if g1 || g2 || g3 || g0 { 558 | t.Fatal("got unexpected shutdown signal", g0, g1, g2, g3) 559 | } 560 | } 561 | 562 | func TestFnPanic(t *testing.T) { 563 | reset() 564 | defer close(startTimer(t)) 565 | gotcall := false 566 | 567 | // Register a function 568 | _ = FirstFunc(func(i interface{}) { 569 | gotcall = i.(bool) 570 | panic("This is expected") 571 | }, true) 572 | 573 | // Start shutdown 574 | Shutdown() 575 | if !gotcall { 576 | t.Fatal("did not get expected shutdown signal") 577 | } 578 | } 579 | 580 | func TestFnNotify(t *testing.T) { 581 | reset() 582 | defer close(startTimer(t)) 583 | gotcall := false 584 | 585 | // Register a function 586 | fn := FirstFunc(func(i interface{}) { 587 | gotcall = i.(bool) 588 | }, true) 589 | 590 | // Start shutdown 591 | Shutdown() 592 | 593 | // This must have a notification 594 | _, ok := <-fn 595 | if !ok { 596 | t.Fatal("Notifier was closed before a notification") 597 | } 598 | // After this the channel must be closed 599 | _, ok = <-fn 600 | if ok { 601 | t.Fatal("Notifier was not closed after initial notification") 602 | } 603 | if !gotcall { 604 | t.Fatal("did not get expected shutdown signal") 605 | } 606 | } 607 | 608 | func TestFnSingleCancel(t *testing.T) { 609 | reset() 610 | defer close(startTimer(t)) 611 | 612 | var ok1, ok2, ok3, okcancel bool 613 | _ = ThirdFunc(func(i interface{}) { 614 | set := i.(*bool) 615 | *set = true 616 | }, &ok3) 617 | if Started() { 618 | t.Fatal("shutdown started unexpectedly") 619 | } 620 | 621 | _ = SecondFunc(func(i interface{}) { 622 | set := i.(*bool) 623 | *set = true 624 | }, &ok2) 625 | if Started() { 626 | t.Fatal("shutdown started unexpectedly") 627 | } 628 | 629 | cancel := SecondFunc(func(i interface{}) { 630 | set := i.(*bool) 631 | *set = true 632 | }, &okcancel) 633 | if Started() { 634 | t.Fatal("shutdown started unexpectedly") 635 | } 636 | 637 | _ = FirstFunc(func(i interface{}) { 638 | set := i.(*bool) 639 | *set = true 640 | }, &ok1) 641 | if Started() { 642 | t.Fatal("shutdown started unexpectedly") 643 | } 644 | 645 | if ok1 || ok2 || ok3 || okcancel { 646 | t.Fatal("shutdown has already happened", ok1, ok2, ok3, okcancel) 647 | } 648 | 649 | cancel.Cancel() 650 | 651 | Shutdown() 652 | 653 | if !ok1 || !ok2 || !ok3 || okcancel { 654 | t.Fatal("did not get expected shutdown signal", ok1, ok2, ok3, okcancel) 655 | } 656 | } 657 | 658 | // Get a notifier and perform our own code when we shutdown 659 | func ExampleNotifier() { 660 | shutdown := First() 661 | select { 662 | case n := <-shutdown: 663 | // Do shutdown code ... 664 | 665 | // Signal we are done 666 | close(n) 667 | } 668 | } 669 | 670 | // Get a notifier and perform our own function when we shutdown 671 | func ExampleShutdownFn() { 672 | _ = FirstFunc(func(i interface{}) { 673 | // This function is called on shutdown 674 | fmt.Println(i.(string)) 675 | }, "Example parameter") 676 | 677 | // Will print the parameter when Shutdown() is called 678 | } 679 | 680 | // Note that the same effect of this example can also be achieved using the 681 | // WrapHandlerFunc helper. 682 | func ExampleLock() { 683 | http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) { 684 | // Get a lock while we have the lock, the server will not shut down. 685 | if Lock() { 686 | defer Unlock() 687 | } else { 688 | // We are currently shutting down, return http.StatusServiceUnavailable 689 | w.WriteHeader(http.StatusServiceUnavailable) 690 | return 691 | } 692 | // ... 693 | }) 694 | http.ListenAndServe(":8080", nil) 695 | } 696 | 697 | // Change timeout for a single stage 698 | func ExampleSetTimeoutN() { 699 | // Set timout for all stages 700 | SetTimeout(time.Second) 701 | 702 | // But give second stage more time 703 | SetTimeoutN(Stage2, time.Second*10) 704 | } 705 | 706 | // This is an example, that could be your main function. 707 | // 708 | // We wait for jobs to finish in another goroutine, from 709 | // where we initialize the shutdown. 710 | // 711 | // This is of course not a real-world problem, but there are many 712 | // cases where you would want to initialize shutdown from other places than 713 | // your main function, and where you would still like it to be able to 714 | // do some final cleanup. 715 | func ExampleWait() { 716 | x := make([]struct{}, 10) 717 | var wg sync.WaitGroup 718 | 719 | wg.Add(len(x)) 720 | for i := range x { 721 | go func(i int) { 722 | time.Sleep(time.Millisecond * time.Duration(i)) 723 | wg.Done() 724 | }(i) 725 | } 726 | 727 | // ignore this reset, for test purposes only 728 | reset() 729 | 730 | // Wait for the jobs above to finish 731 | go func() { 732 | wg.Wait() 733 | fmt.Println("jobs done") 734 | Shutdown() 735 | }() 736 | 737 | // Since this is main, we wait for a shutdown to occur before 738 | // exiting. 739 | Wait() 740 | fmt.Println("exiting main") 741 | 742 | // Note than the output will always be in this order. 743 | 744 | // Output: jobs done 745 | // exiting main 746 | } 747 | --------------------------------------------------------------------------------