├── go.mod
├── logger.go
├── .gitignore
├── .github
├── workflows
│ └── pr.yaml
└── ISSUE_TEMPLATE
│ ├── feature_request.yaml
│ └── bug_report.yaml
├── LICENSE
├── types.go
├── default.go
├── doc.go
├── pipeline.go
├── README.md
└── pipeline_test.go
/go.mod:
--------------------------------------------------------------------------------
1 | module github.com/meshapi/go-shutdown
2 |
3 | go 1.19
4 |
--------------------------------------------------------------------------------
/logger.go:
--------------------------------------------------------------------------------
1 | package shutdown
2 |
3 | import (
4 | "log"
5 | )
6 |
7 | // StandardLogger is a logger that uses the standard log package.
8 | type StandardLogger struct {
9 | logger *log.Logger
10 | }
11 |
12 | // NewStandardLogger creates a logger wrapping a log.Logger instance.
13 | func NewStandardLogger(logger *log.Logger) StandardLogger {
14 | return StandardLogger{logger: logger}
15 | }
16 |
17 | func (d StandardLogger) Info(text string) {
18 | d.logger.Printf(text)
19 | }
20 |
21 | func (d StandardLogger) Error(text string) {
22 | d.logger.Printf(text)
23 | }
24 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # If you prefer the allow list template instead of the deny list, see community template:
2 | # https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore
3 | #
4 | # Binaries for programs and plugins
5 | *.exe
6 | *.exe~
7 | *.dll
8 | *.so
9 | *.dylib
10 |
11 | # Test binary, built with `go test -c`
12 | *.test
13 |
14 | # Output of the go coverage tool, specifically when used with LiteIDE
15 | *.out
16 |
17 | # Dependency directories (remove the comment below to include it)
18 | # vendor/
19 |
20 | # Go workspace file
21 | go.work
22 |
--------------------------------------------------------------------------------
/.github/workflows/pr.yaml:
--------------------------------------------------------------------------------
1 | on:
2 | pull_request:
3 | types:
4 | - opened
5 | - reopened
6 | - synchronize
7 | - edited
8 | branches:
9 | - 'main'
10 |
11 | jobs:
12 | pr-checks:
13 | runs-on: ubuntu-latest
14 | steps:
15 | - uses: actions/checkout@v3
16 |
17 | - name: Set up Go
18 | uses: actions/setup-go@v4
19 | with:
20 | go-version: '1.19'
21 |
22 | - name: Test
23 | run: go test -v ./...
24 |
25 | - name: Lint
26 | uses: golangci/golangci-lint-action@v3
27 | with:
28 | version: v1.54
29 |
--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/feature_request.yaml:
--------------------------------------------------------------------------------
1 | name: Feature request
2 | description: "Suggest an idea for this project."
3 | labels: [enhancement]
4 | body:
5 |
6 | - type: textarea
7 | id: problem
8 | attributes:
9 | label: Your feature request related to a problem? Please describe.
10 | placeholder: "A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]"
11 | validations:
12 | required: true
13 |
14 | - type: textarea
15 | id: solution
16 | attributes:
17 | label: Describe the solution you'd like.
18 | placeholder: "A clear and concise description of what you want to happen."
19 | validations:
20 | required: true
21 |
22 | - type: textarea
23 | id: alternatives
24 | attributes:
25 | label: Describe alternatives you've considered.
26 | placeholder: "A clear and concise description of any alternative solutions or features you've considered."
27 | validations:
28 | required: true
29 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2024 Mesh API
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 |
--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/bug_report.yaml:
--------------------------------------------------------------------------------
1 | name: Bug Report
2 | description: "Create a report to help us improve."
3 | labels: [bug]
4 | body:
5 | - type: textarea
6 | id: problem
7 | attributes:
8 | label: Description of the problem
9 | placeholder: Your problem description
10 | validations:
11 | required: true
12 |
13 | - type: textarea
14 | id: version
15 | attributes:
16 | label: Version of go-shutdown
17 | value: "version of go-shutdown"
18 | validations:
19 | required: true
20 |
21 | - type: textarea
22 | id: go-env
23 | attributes:
24 | label: Go environment
25 | value: |-
26 |
27 |
28 | ```console
29 | $ go version && go env
30 | # paste output here
31 | ```
32 |
33 |
34 | validations:
35 | required: true
36 |
37 | - type: textarea
38 | id: code-example
39 | attributes:
40 | label: A minimal reproducible example or link to a public repository
41 | description: if your problem is related to a private repository, a minimal reproducible example is required.
42 | value: |-
43 |
44 |
45 | ```go
46 | // add your code here
47 | ```
48 |
49 |
50 | validations:
51 | required: true
52 |
53 | - type: checkboxes
54 | id: validation
55 | attributes:
56 | label: Validation
57 | options:
58 | - label: Yes, I've included all information above (version, config, etc.).
59 | required: true
60 |
--------------------------------------------------------------------------------
/types.go:
--------------------------------------------------------------------------------
1 | package shutdown
2 |
3 | import (
4 | "context"
5 | "fmt"
6 | )
7 |
8 | // Logger describes ability to write logs.
9 | type Logger interface {
10 | // Info writes an information log.
11 | Info(text string)
12 |
13 | // Error writes an error log.
14 | Error(text string)
15 | }
16 |
17 | // Handler describes ability to handle a graceful shutdown.
18 | type Handler interface {
19 | HandleShutdown(ctx context.Context) error
20 | }
21 |
22 | // NamedHandler is a handler that has a specific name.
23 | type NamedHandler interface {
24 | Handler
25 |
26 | Name() string
27 | }
28 |
29 | type namedHandler struct {
30 | Handler
31 | name string
32 | }
33 |
34 | func (n namedHandler) Name() string {
35 | return n.name
36 | }
37 |
38 | // HandleFunc describes various different function signatures that can be used as a shutdown handler function.
39 | type HandleFunc interface {
40 | func() | func() error | func(context.Context) | func(context.Context) error
41 | }
42 |
43 | // HandlerWithName returns a named handler with a specific name.
44 | func HandlerWithName(name string, handler Handler) NamedHandler {
45 | return namedHandler{name: name, Handler: handler}
46 | }
47 |
48 | // HandlerFuncWithName returns a named handler for a function with a specific name.
49 | //
50 | // Accepted function signatures:
51 | // func ()
52 | // func (context.Context)
53 | // func () error
54 | // func (context.Context) error
55 | func HandlerFuncWithName[H HandleFunc](name string, handleFunc H) NamedHandler {
56 | if handleFunc, ok := any(handleFunc).(func()); ok {
57 | return namedHandler{name: name, Handler: handlerFuncNoError(handleFunc)}
58 | }
59 |
60 | if handleFunc, ok := any(handleFunc).(func() error); ok {
61 | return namedHandler{name: name, Handler: handlerFunc(handleFunc)}
62 | }
63 |
64 | if handleFunc, ok := any(handleFunc).(func(context.Context)); ok {
65 | return namedHandler{name: name, Handler: handlerFuncContextNoError(handleFunc)}
66 | }
67 |
68 | if handleFunc, ok := any(handleFunc).(func(context.Context) error); ok {
69 | return namedHandler{name: name, Handler: handlerFuncContext(handleFunc)}
70 | }
71 |
72 | panic(fmt.Sprintf("unexpected function signature for handler: %T", handleFunc))
73 | }
74 |
75 | type (
76 | handlerFuncNoError func()
77 | handlerFunc func() error
78 | handlerFuncContext func(context.Context) error
79 | handlerFuncContextNoError func(context.Context)
80 | )
81 |
82 | func (n handlerFuncNoError) HandleShutdown(context.Context) error {
83 | n()
84 | return nil
85 | }
86 |
87 | func (h handlerFunc) HandleShutdown(ctx context.Context) error {
88 | return h()
89 | }
90 |
91 | func (h handlerFuncContext) HandleShutdown(ctx context.Context) error {
92 | return h(ctx)
93 | }
94 |
95 | func (n handlerFuncContextNoError) HandleShutdown(ctx context.Context) error {
96 | n(ctx)
97 | return nil
98 | }
99 |
--------------------------------------------------------------------------------
/default.go:
--------------------------------------------------------------------------------
1 | package shutdown
2 |
3 | import (
4 | "context"
5 | "sync"
6 | "time"
7 | )
8 |
9 | var (
10 | defaultPipeline *Manager
11 | defaultLock sync.Mutex
12 | )
13 |
14 | // Default returns the default manager. This method is thread-safe.
15 | func Default() *Manager {
16 | if defaultPipeline == nil {
17 | defaultLock.Lock()
18 | defer defaultLock.Unlock()
19 |
20 | // if after acquiring the lock, pipeline is still nil, then set it.
21 | if defaultPipeline == nil {
22 | defaultPipeline = New()
23 | }
24 | }
25 |
26 | return defaultPipeline
27 | }
28 |
29 | // WaitForInterrupt blocks until an interrupt signal is received and all shutdown steps have been executed.
30 | func WaitForInterrupt() {
31 | Default().WaitForInterrupt()
32 | }
33 |
34 | // Trigger starts the shutdown pipeline immediately. It will acquire a lock on the pipeline so all changes to the
35 | // pipeline get blocked until the pipeline has completed. Panics and errors are all handled.
36 | func Trigger(ctx context.Context) {
37 | Default().Trigger(ctx)
38 | }
39 |
40 | // AddParallelSequence is similar to AddSequence but it will execute the handlers all at the same time.
41 | // AddParallelSequence(a) and AddParallelSequence(b) is not the same as AddParallelSequence(a, b). In the former, a
42 | // runs and upon completion, b starts whereas in the latter case a and b both get started at the same time.
43 | func AddParallelSequence(handlers ...NamedHandler) {
44 | Default().AddParallelSequence(handlers...)
45 | }
46 |
47 | // AddSequence adds sequencial steps meaning that these handlers will be executed one at a time and in the same order
48 | // given.
49 | // Calling AddSequence(a) and AddSequence(b) is same as AddSequence(a, b)
50 | func AddSequence(handlers ...NamedHandler) {
51 | Default().AddSequence(handlers...)
52 | }
53 |
54 | // AddSteps adds parallel shutdown steps. These steps will be executed at the same time together or along
55 | // with previously added steps if they are also able to run in parallel. In another word, calling AddSteps(a) and
56 | // AddSteps(b) is same as AddSteps(a, b)
57 | func AddSteps(handlers ...NamedHandler) {
58 | Default().AddSteps(handlers...)
59 | }
60 |
61 | // SetLogger sets the shutdown logger. If set to nil, no logs will be written.
62 | func SetLogger(logger Logger) {
63 | Default().SetLogger(logger)
64 | }
65 |
66 | // SetCompletionFunc sets a function to get called after all of the shutdown steps have been executed. Regardless of
67 | // panics or errors, this function will always get executed as the very last step. Even when a the pipeline times out,
68 | // this function gets called before returning.
69 | func SetCompletionFunc(f func()) {
70 | Default().SetCompletionFunc(f)
71 | }
72 |
73 | // SetTimeout sets the shutdown pipeline timeout. This indicates that when shutdown is triggered, the entire pipeline
74 | // iteration must finish within the duration specified.
75 | //
76 | // NOTE: If the pipeline times out, the shutdown method is still called and some of the steps in the pipeline will
77 | // still get scheduled but the blocking method (Trigger or WaitForInterrupt) will return immediately without waiting
78 | // for the rest of the shutdown steps to complete.
79 | func SetTimeout(duration time.Duration) {
80 | Default().SetTimeout(duration)
81 | }
82 |
--------------------------------------------------------------------------------
/doc.go:
--------------------------------------------------------------------------------
1 | // Package shutdown contains utilities to gracefully handle kernel interrupts and shutdown procedures.
2 | //
3 | // This package provides means to define a shutdown pipeline, which can contain sequential and parallel steps to
4 | // complete a shutdown and also ways to trigger the pipeline either via signals or manually.
5 | //
6 | // To create a pipeline manager, use the following code:
7 | //
8 | // shutdownManager := New()
9 | //
10 | // There is a default shutdown manager that can be accessed via Default() method and many package-level functions are
11 | // available as a shortcut to accessing the default manager's methods. The default manager uses the standard log
12 | // package for logging.
13 | //
14 | // The shutdown manager allows the addition of steps to the shutdown procedure, with some steps capable of running in
15 | // parallel, while others run sequentially. The shutdown procedure can be triggered manually or by subscribing to
16 | // kernel SIGTERM interrupt signals.
17 | //
18 | // There are three main methods for adding steps into the shutdown manager. The shutdown manager monitors the
19 | // added steps and organizes them into sequential groups. Unless explicitly specified to form a distinct group,
20 | // consecutive steps with the same parallel status are grouped together.
21 | //
22 | // * AddSteps: Adds steps that can run concurrently in separate goroutines.
23 | //
24 | // Example:
25 | // In the example below, handlers 'a' and 'b' are called concurrently.
26 | //
27 | // AddSteps(
28 | // HandlerFuncWithName("a", func(){}),
29 | // HandlerFuncWithName("b", func(){}))
30 | //
31 | // NOTE: Neighboring steps with the same parallelism status are grouped together. Thus the following lines have the
32 | // same effect as the code above.
33 | //
34 | // AddSteps(HandlerFuncWithName("a", func(){}))
35 | // AddSteps(HandlerFuncWithName("b", func(){}))
36 | //
37 | // * AddSequence: Adds steps that run one after another in the given order, without separate goroutines.
38 | //
39 | // Example:
40 | // In the example below, handler 'a' is called first, followed by handler 'b'.
41 | //
42 | // AddSequence(
43 | // HandlerFuncWithName("a", func(){}),
44 | // HandlerFuncWithName("b", func(){}))
45 | //
46 | // * AddParallelSequence: Adds steps to be executed in parallel, with an explicit instruction for the manager to wait
47 | // for all steps prior to finish before starting the execution of this group.
48 | //
49 | // Example:
50 | // In the example below handlers:
51 | // 1. Handlers 'a' and 'b' are called concurrently.
52 | // 2. After 'a' and 'b' finish, handlers 'c', 'd', and 'e' are called in parallel.
53 | // 3. After 'c', 'd', and 'e' finish, handler 'e' is called, and upon completion, handler 'f' is called.
54 | //
55 | // AddSteps(
56 | // HandlerFuncWithName("a", func(){}),
57 | // HandlerFuncWithName("b", func(){})).
58 | // AddParallelSequence(
59 | // HandlerFuncWithName("c", func(){}),
60 | // HandlerFuncWithName("d", func(){})).
61 | // AddSteps(HandlerFuncWithName("e", func(){})).
62 | // AddSequence(HandlerFuncWithName("f", func(){}))
63 | //
64 | // Additional features include:
65 | // * SetLogger: Set a custom logger for the shutdown pipeline (default is plog.Default()).
66 | // * SetTimeout: Define a timeout for the execution of the entire shutdown pipeline.
67 | // * SetCompletionFunc: Add a callback function to handle the end of the shutdown pipeline, ensuring it gets called in
68 | // any case, even in the presence of panics, errors, or timeouts.
69 | package shutdown
70 |
--------------------------------------------------------------------------------
/pipeline.go:
--------------------------------------------------------------------------------
1 | package shutdown
2 |
3 | import (
4 | "context"
5 | "fmt"
6 | "log"
7 | "os"
8 | "os/signal"
9 | "sync"
10 | "syscall"
11 | "time"
12 | )
13 |
14 | // Manager is a shutdown pipeline manager that can be configured to run through a number of parallel and sequencial
15 | // steps when a shutdown is triggered. In order to start the shutdown procedure upon receiving a kernel Interrupt
16 | // signal, use WaitForInterrupt() blocking method.
17 | type Manager struct {
18 | steps []shutdownStep
19 | timeout time.Duration
20 | completionFuncnc func()
21 | logger Logger
22 | lock sync.Mutex
23 | }
24 |
25 | // New creates a new shutdown pipeline.
26 | func New() *Manager {
27 | return &Manager{
28 | logger: NewStandardLogger(log.Default()),
29 | }
30 | }
31 |
32 | // SetTimeout sets the shutdown pipeline timeout. This indicates that when shutdown is triggered, the entire pipeline
33 | // iteration must finish within the duration specified.
34 | //
35 | // NOTE: If the pipeline times out, the shutdown method is still called and some of the steps in the pipeline will
36 | // still get scheduled but the blocking method (Trigger or WaitForInterrupt) will return immediately without waiting
37 | // for the rest of the shutdown steps to complete.
38 | func (m *Manager) SetTimeout(duration time.Duration) {
39 | m.lock.Lock()
40 | defer m.lock.Unlock()
41 |
42 | m.timeout = duration
43 | }
44 |
45 | // SetCompletionFunc sets a function to get called after all of the shutdown steps have been executed. Regardless of
46 | // panics or errors, this function will always get executed as the very last step. Even when a the pipeline times out,
47 | // this function gets called before returning.
48 | func (m *Manager) SetCompletionFunc(f func()) {
49 | m.lock.Lock()
50 | defer m.lock.Unlock()
51 |
52 | m.completionFuncnc = f
53 | }
54 |
55 | // SetLogger sets the shutdown logger. If set to nil, no logs will be written.
56 | func (m *Manager) SetLogger(logger Logger) {
57 | m.lock.Lock()
58 | defer m.lock.Unlock()
59 |
60 | m.logger = logger
61 | }
62 |
63 | // AddSteps adds parallel shutdown steps. These steps will be executed at the same time together or along
64 | // with previously added steps if they are also able to run in parallel. In another word, calling AddSteps(a) and
65 | // AddSteps(b) is same as AddSteps(a, b)
66 | func (m *Manager) AddSteps(handlers ...NamedHandler) {
67 | m.lock.Lock()
68 | defer m.lock.Unlock()
69 |
70 | if len(m.steps) == 0 {
71 | m.steps = append(m.steps, shutdownStep{handlers: handlers, parallel: true})
72 | return
73 | }
74 |
75 | lastStep := m.steps[len(m.steps)-1]
76 | if lastStep.parallel {
77 | lastStep.handlers = append(lastStep.handlers, handlers...)
78 | m.steps[len(m.steps)-1] = lastStep
79 | return
80 | }
81 |
82 | m.steps = append(m.steps, shutdownStep{handlers: handlers, parallel: true})
83 | }
84 |
85 | // AddSequence adds sequencial steps meaning that these handlers will be executed one at a time and in the same order
86 | // given.
87 | // Calling AddSequence(a) and AddSequence(b) is same as AddSequence(a, b)
88 | func (m *Manager) AddSequence(handlers ...NamedHandler) {
89 | m.lock.Lock()
90 | defer m.lock.Unlock()
91 |
92 | if len(m.steps) == 0 {
93 | m.steps = append(m.steps, shutdownStep{handlers: handlers, parallel: false})
94 | return
95 | }
96 |
97 | lastStep := m.steps[len(m.steps)-1]
98 | if !lastStep.parallel {
99 | lastStep.handlers = append(lastStep.handlers, handlers...)
100 | m.steps[len(m.steps)-1] = lastStep
101 | return
102 | }
103 |
104 | m.steps = append(m.steps, shutdownStep{handlers: handlers, parallel: false})
105 | }
106 |
107 | // AddParallelSequence is similar to AddSequence but it will execute the handlers all at the same time.
108 | // AddParallelSequence(a) and AddParallelSequence(b) is not the same as AddParallelSequence(a, b). In the former, a
109 | // runs and upon completion, b starts whereas in the latter case a and b both get started at the same time.
110 | func (m *Manager) AddParallelSequence(handlers ...NamedHandler) {
111 | m.lock.Lock()
112 | defer m.lock.Unlock()
113 |
114 | m.steps = append(m.steps, shutdownStep{handlers: handlers, parallel: true})
115 | }
116 |
117 | // Trigger starts the shutdown pipeline immediately. It will acquire a lock on the pipeline so all changes to the
118 | // pipeline get blocked until the pipeline has completed. Panics and errors are all handled.
119 | func (m *Manager) Trigger(ctx context.Context) {
120 | m.lock.Lock()
121 | defer m.lock.Unlock()
122 |
123 | if len(m.steps) == 0 {
124 | return
125 | }
126 |
127 | if m.timeout != 0 {
128 | newCtx, cancel := context.WithTimeout(ctx, m.timeout)
129 | ctx = newCtx
130 | defer cancel()
131 | }
132 |
133 | errorCount := 0
134 | resultChannel := make(chan handlerResult)
135 |
136 | mainLoop:
137 | for _, step := range m.steps {
138 | remainingHandlers := len(step.handlers)
139 |
140 | go func() {
141 | for _, handler := range step.handlers {
142 | if step.parallel {
143 | go func(h NamedHandler) {
144 | m.executeHandler(ctx, h, resultChannel)
145 | }(handler)
146 | } else {
147 | m.executeHandler(ctx, handler, resultChannel)
148 | }
149 | }
150 | }()
151 |
152 | for remainingHandlers > 0 {
153 | select {
154 | case result := <-resultChannel:
155 | if result.Err != nil {
156 | errorCount++
157 | m.err(result.HandlerName + " shutdown failed: " + result.Err.Error())
158 | } else {
159 | m.info(result.HandlerName + " shutdown completed")
160 | }
161 | remainingHandlers--
162 | case <-ctx.Done():
163 | m.err("context canceled")
164 | errorCount++
165 | break mainLoop
166 | }
167 | }
168 | }
169 |
170 | if m.completionFuncnc != nil {
171 | m.completionFuncnc()
172 | }
173 |
174 | if errorCount > 0 {
175 | m.err(fmt.Sprintf("shutdown pipeline completed with %d errors", errorCount))
176 | } else {
177 | m.info("shutdown pipeline completed with no errors")
178 | }
179 | }
180 |
181 | func (m *Manager) info(text string) {
182 | if m.logger != nil {
183 | m.logger.Info(text)
184 | }
185 | }
186 |
187 | func (m *Manager) err(text string) {
188 | if m.logger != nil {
189 | m.logger.Error(text)
190 | }
191 | }
192 |
193 | func (m *Manager) executeHandler(ctx context.Context, handler NamedHandler, resultChannel chan<- handlerResult) {
194 | var err error
195 |
196 | defer func() {
197 | if panicErr := recover(); panicErr != nil {
198 | resultChannel <- handlerResult{HandlerName: handler.Name(), Err: fmt.Errorf("panic: %s", panicErr)}
199 | } else {
200 | resultChannel <- handlerResult{HandlerName: handler.Name(), Err: err}
201 | }
202 | }()
203 |
204 | err = handler.HandleShutdown(ctx)
205 | }
206 |
207 | // WaitForInterrupt blocks until an interrupt signal is received and all shutdown steps have been executed.
208 | func (m *Manager) WaitForInterrupt() {
209 | exit := make(chan os.Signal, 1)
210 | signal.Notify(exit, os.Interrupt, syscall.SIGTERM)
211 | <-exit
212 |
213 | m.info("received interrupt signal, starting shutdown procedures...")
214 |
215 | m.Trigger(context.Background())
216 | }
217 |
218 | type shutdownStep struct {
219 | handlers []NamedHandler
220 | parallel bool
221 | }
222 |
223 | type handlerResult struct {
224 | Err error
225 | HandlerName string
226 | }
227 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Graceful shutdown utility for Go
2 |
3 | 
4 | 
5 | 
6 | 
7 | 
8 |
9 | The go-shutdown package offers a collection of utilities designed for managing shutdown procedures,
10 | primarily focused on gracefully handling termination signals from the kernel, which is a common and
11 | expected use case.
12 |
13 | Ensuring graceful shutdowns is crucial for services to support zero-downtime deployments among other reasons.
14 | During a deployment, the application that is currently running receives a SIGTERM signal and the process
15 | can gracefully finish the work in progress and stop taking new requests in order to ensure no active task is
16 | interrupted.
17 |
18 | While this may seem straightforward, it is often overlooked when starting a new service or project.
19 | The purpose of this package is to make it as easy as writing 2-3 lines of code to handle shutdowns
20 | gracefully and also to remain flexible for more advanced usecases.
21 |
22 |
23 | ## Features
24 |
25 | * [Use different concurrency options for steps in the shutdown pipeline](#configure-your-shutdown-pipeline)
26 | * [Manual trigger or wait for SIGTERM signal](#shutdown-trigger)
27 | * [Set timeout](#timeout)
28 |
29 | ## Installation
30 |
31 | ```bash
32 | go get github.com/meshapi/go-shutdown
33 | ```
34 |
35 | ## Getting started
36 |
37 | The way to handle shutdowns is to create and configure a shutdown pipeline and decide how it should get triggered.
38 | There is a default shutdown pipeline for convenience but you can create separate pipelines as well.
39 | Configuring the pipeline involves defining the steps and their concurrency. Finally deciding on triggering the pipeline
40 | manually or subscribing to process signals.
41 |
42 | ### Configure your shutdown pipeline
43 |
44 | To create a new shutdown pipeline use `shutdown.New` or simply use the package level methods to configure the default
45 | shutdown pipeline.
46 |
47 | There are three different methods to add new steps to the shutdown procedure, each with distinct concurrency behaviors.
48 | Each has the same method signature which takes `shutdown.NamedHandler` instances. For logging purposes, each
49 | handler/step must have a name.
50 |
51 | Shortcut method `shutdown.HandlerWithName(string,Handler)` creates a `shutdown.NamedHandler` from a `shutdown.Handler`.
52 |
53 | Shortcut method `shutdown.HandlerFuncWithName(string,HandleFunc)` creates a `shutdown.NamedHandler` from a callback and
54 | accepts any of the following function types:
55 | - `func()`
56 | - `func() error`
57 | - `func(context.Context)`
58 | - `func(context.Context) error`
59 |
60 | As new steps are added to the shutdown pipeline, the pipeline organizes the steps into sequential groups.
61 | Unless explicitly specified to form a distinct group, consecutive steps with the same parallel status
62 | are grouped together.
63 |
64 | * `AddSteps`: Adds steps that can run concurrently in separate goroutines.
65 |
66 | Example: In the example below, handlers 'a' and 'b' are called concurrently.
67 |
68 | ```go
69 | shutdown.AddSteps(
70 | shutdown.HandlerFuncWithName("a", func(){}),
71 | shutdown.HandlerFuncWithName("b", func(){}))
72 | ```
73 |
74 | > NOTE: Neighboring steps with the same parallelism status are grouped together. Thus the following lines have the
75 | same effect as the code above.
76 |
77 | ```go
78 | shutdown.AddSteps(shutdown.HandlerFuncWithName("a", func(){}))
79 | shutdown.AddSteps(shutdown.HandlerFuncWithName("b", func(){}))
80 | ```
81 |
82 | * `AddSequence`: Adds steps that run one after another in the given order, without separate goroutines.
83 |
84 | Example: In the example below, handler 'a' is called first, followed by handler 'b'.
85 |
86 | ```go
87 | shutdown.AddSequence(
88 | shutdown.HandlerFuncWithName("a", func(){}),
89 | shutdown.HandlerFuncWithName("b", func(){}))
90 | ```
91 |
92 | * `AddParallelSequence`: Adds steps to be executed in parallel, with an explicit instruction for the manager to wait
93 | for all steps prior to finish before starting the execution of this group.
94 |
95 | Example: With the configuration code below:
96 | 1. Handlers 'a' and 'b' are called concurrently.
97 | 2. After 'a' and 'b' finish, handlers 'c', 'd', and 'e' are called in parallel.
98 | 3. After 'c', 'd', and 'e' finish, handler 'e' is called, and upon completion, handler 'f' is called.
99 |
100 | ```go
101 | manager := shutdown.New()
102 |
103 | manager.AddSteps(
104 | shutdown.HandlerFuncWithName("a", func(){}),
105 | shutdown.HandlerFuncWithName("b", func(){})).
106 | manager.AddParallelSequence(
107 | shutdown.HandlerFuncWithName("c", func(){}),
108 | shutdown.HandlerFuncWithName("d", func(){})).
109 | manager.AddSteps(shutdown.HandlerFuncWithName("e", func(){})).
110 | manager.AddSequence(shutdown.HandlerFuncWithName("f", func(){}))
111 | ```
112 |
113 | > NOTE: Many servers have a graceful shutdown method and they can be used with
114 | > `shutdown.HandlerFuncWithName(string,HandleFunc)` method. The code below is an example of an HTTP server shutdown:
115 |
116 | ```go
117 | httpServer := &http.Server{}
118 | go func() {
119 | if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
120 | log.Fatalf("HTTP server failed: %s", err)
121 | }
122 | }()
123 |
124 | shutdown.AddSteps(shutdown.HandlerFuncWithName("http", httpServer.Shutdown))
125 | shutdown.WaitForInterrupt()
126 | ```
127 |
128 | ### Shutdown trigger
129 |
130 | To manually trigger a shutdown, use `Trigger(context.Context)` method on a shutdown manager instance.
131 | `shutdown.Trigger` is a shortcut method to trigger the default pipeline's shutdown procedure.
132 |
133 | To trigger the shutdown when a SIGTERM is received, simply call `WaitForInterrupt()` method. This method blocks until
134 | SIGTERM signal is received and the shutdown pipeline has concluded or deadline reached.
135 |
136 | ```go
137 | manager := shutdown.New()
138 |
139 | manager.Trigger(context.Background()) // manual trigger.
140 | manager.WaitForInterrupt() // block until SIGTERM is received and shutdown procedures have finished.
141 | ```
142 |
143 | ### Logging
144 |
145 | In order to remain flexible with your logging tool, no choice over the logger is made here, instead any type that has
146 | `Info` and `Error` methods, can be used as a logger.
147 |
148 | The default shutdown pipeline uses the `log` package from the standard library but when you create new instances, no
149 | logger is set and when no logger is available, no logging will be made.
150 |
151 | Use the `SetLogger(Logger)` method to set a logger.
152 |
153 | ```go
154 | var myLogger shutdown.Logger
155 |
156 | manager := shutdown.New()
157 | manager.SetLogger(myLogger) // set the logger on the newly created shutdown pipeline.
158 |
159 | shutdown.SetLogger(myLogger) // update the logger on the default pipeline.
160 | ```
161 |
162 | ### Timeout
163 |
164 | Use method `SetTimeout(time.Duration)` to set a deadline on the shutdown pipeline's completion time. By default this is
165 | not specified.
166 |
167 | ```go
168 | shutdown.SetTimeout(15*time.Second)
169 | ```
170 |
171 | ### Completion callback
172 |
173 | A completion function callback can be set via `SetCompletionFunc`. This function will get called in any case, even if
174 | there are panics, errors or timeouts.
175 |
176 | ----------------
177 |
178 | ## Contributions
179 |
180 | Contributions are absolutely welcome and please write tests for the new functionality added/modified.
181 | Additionally, we ask that you include a test function reproducing the issue raised to save everyone's time.
182 |
--------------------------------------------------------------------------------
/pipeline_test.go:
--------------------------------------------------------------------------------
1 | package shutdown_test
2 |
3 | import (
4 | "bytes"
5 | "context"
6 | "errors"
7 | "fmt"
8 | "io"
9 | "sort"
10 | "strings"
11 | "sync"
12 | "testing"
13 | "time"
14 |
15 | "github.com/meshapi/go-shutdown"
16 | )
17 |
18 | type LogWrapper struct {
19 | Writer io.Writer
20 | }
21 |
22 | func (l LogWrapper) Info(text string) {
23 | _, _ = l.Writer.Write([]byte(text))
24 | }
25 |
26 | func (l LogWrapper) Error(text string) {
27 | _, _ = l.Writer.Write([]byte(text))
28 | }
29 |
30 | // EventTime is to store at what relative time a process completes or should complete.
31 | type EventTime map[string]time.Duration
32 |
33 | func (e EventTime) String() string {
34 | keys := []string{}
35 | for key := range e {
36 | keys = append(keys, key)
37 | }
38 | sort.Strings(keys)
39 | writer := &strings.Builder{}
40 | for _, key := range keys {
41 | _, _ = fmt.Fprintf(writer, "%s\t%d\n", key, e[key].Milliseconds())
42 | }
43 | return writer.String()
44 | }
45 |
46 | // MustParseEventTime parsed :,... format.
47 | func MustParseEventTime(value string) EventTime {
48 | parts := strings.Split(value, ",")
49 | result := EventTime{}
50 | for _, part := range parts {
51 | sections := strings.Split(part, ":")
52 | duration, err := time.ParseDuration(sections[1])
53 | if err != nil {
54 | panic("failed to parse time duration " + sections[1])
55 | }
56 | result[sections[0]] = duration
57 | }
58 | return result
59 | }
60 |
61 | // SequenceMonitor captures events and marks their time of execution and asserts if they occurred in a certain time
62 | // frame.
63 | type SequenceMonitor struct {
64 | events EventTime
65 | startTime time.Time
66 | lock sync.Mutex
67 | }
68 |
69 | // Mark stores the relative time of completion of the given procedure. Needs to be called after StartRecording()
70 | func (s *SequenceMonitor) Mark(name string) {
71 | s.lock.Lock()
72 | defer s.lock.Unlock()
73 |
74 | s.events[name] = time.Since(s.startTime)
75 | }
76 |
77 | // EventTime returns the current recorded event time.
78 | func (s *SequenceMonitor) EventTime() EventTime {
79 | return s.events
80 | }
81 |
82 | // StartRecording sets the start time and initializes the event time.
83 | func (s *SequenceMonitor) StartRecording() {
84 | s.lock.Lock()
85 | defer s.lock.Unlock()
86 |
87 | s.events = EventTime{}
88 | s.startTime = time.Now()
89 | }
90 |
91 | // Matches returns whether or not the captured events and their timelines matches the input.
92 | func (s *SequenceMonitor) Matches(input EventTime) bool {
93 | s.lock.Lock()
94 | defer s.lock.Unlock()
95 |
96 | if len(s.events) != len(input) {
97 | return false
98 | }
99 |
100 | for inputKey, inputDuration := range input {
101 | expectedDuration, ok := s.events[inputKey]
102 | if !ok || !s.approximatelySameDuration(inputDuration, expectedDuration) {
103 | return false
104 | }
105 | }
106 |
107 | return true
108 | }
109 |
110 | // if the duration is accurate to the 10ms precision, then this function returns true.
111 | func (s *SequenceMonitor) approximatelySameDuration(d1, d2 time.Duration) bool {
112 | return (d1.Milliseconds() - d1.Milliseconds()%100) == (d2.Milliseconds() - d2.Milliseconds()%100)
113 | }
114 |
115 | func contextWithMonitor(monitor *SequenceMonitor) context.Context {
116 | //nolint:staticcheck
117 | return context.WithValue(context.Background(), "monitor", monitor)
118 | }
119 |
120 | // ShutdownMarker is a test shutdown handler that marks the completion of events using the SequenceMonitor from the context.
121 | type ShutdownMarker struct {
122 | name string
123 | err error
124 | panic bool
125 | }
126 |
127 | func (s ShutdownMarker) Name() string {
128 | return s.name
129 | }
130 |
131 | func (s ShutdownMarker) HandleShutdown(ctx context.Context) error {
132 | monitor := ctx.Value("monitor").(*SequenceMonitor)
133 | if monitor == nil {
134 | panic("no monitor available")
135 | }
136 |
137 | time.Sleep(100 * time.Millisecond)
138 | monitor.Mark(s.name)
139 | if s.panic {
140 | panic("shutdown panic")
141 | }
142 | return s.err
143 | }
144 |
145 | func TestTriggerRuntime(t *testing.T) {
146 | testCases := []struct {
147 | Name string
148 | Setup func(*shutdown.Manager)
149 | Timeout time.Duration
150 | ExpectedEventTime EventTime
151 | }{
152 | {
153 | // in this test, since all steps are parallel, they should all complete at the 10ms mark.
154 | Name: "AllParallel",
155 | Setup: func(m *shutdown.Manager) {
156 | m.AddSteps(ShutdownMarker{name: "a"}, ShutdownMarker{name: "b"})
157 | m.AddSteps(ShutdownMarker{name: "c"})
158 | },
159 | ExpectedEventTime: MustParseEventTime("a:100ms,b:100ms,c:100ms"),
160 | },
161 | {
162 | // in this test, the first two are in one parallel group, the second group is also in a parallel group so the
163 | // first group should complete around the same time and after that the second group should start and finish
164 | // around the same time.
165 | Name: "TwoParallelSequences",
166 | Setup: func(m *shutdown.Manager) {
167 | m.AddSteps(ShutdownMarker{name: "a"}, ShutdownMarker{name: "b"})
168 | m.AddParallelSequence(ShutdownMarker{name: "c"}, ShutdownMarker{name: "d"})
169 | },
170 | ExpectedEventTime: MustParseEventTime("a:100ms,b:100ms,c:200ms,d:200ms"),
171 | },
172 | {
173 | // in this test, the first group runs in parallel and when they're all completed the second group runs in
174 | // sequence and one after each other. So a and b complete at the same time, after that c runs and completes then
175 | // d runs and completes.
176 | Name: "ParallelAndSequence",
177 | Setup: func(m *shutdown.Manager) {
178 | m.AddSteps(
179 | ShutdownMarker{name: "a"},
180 | ShutdownMarker{name: "b"})
181 | m.AddSequence(
182 | ShutdownMarker{name: "c"},
183 | ShutdownMarker{name: "d"})
184 | },
185 | ExpectedEventTime: MustParseEventTime("a:100ms,b:100ms,c:200ms,d:300ms"),
186 | },
187 | {
188 | // in this test, we have an ordered set of steps but one of them errors out.
189 | Name: "SequenceWithError",
190 | Setup: func(m *shutdown.Manager) {
191 | m.AddSequence(
192 | ShutdownMarker{name: "a", err: errors.New("failed")},
193 | ShutdownMarker{name: "b"})
194 | },
195 | ExpectedEventTime: MustParseEventTime("a:100ms,b:200ms"),
196 | },
197 | {
198 | // in this test, we have an ordered set of steps but one of them panics.
199 | Name: "SequenceWithPanic",
200 | Setup: func(m *shutdown.Manager) {
201 | m.AddSequence(
202 | ShutdownMarker{name: "a", panic: true},
203 | ShutdownMarker{name: "b"})
204 | },
205 | ExpectedEventTime: MustParseEventTime("a:100ms,b:200ms"),
206 | },
207 | {
208 | // in this test, timeout is tested. Shutdown should still get called and the first step should complete but the
209 | // second step should no longer block the execution.
210 | Name: "SequenceWithTimeout",
211 | Setup: func(m *shutdown.Manager) {
212 | m.AddSequence(
213 | ShutdownMarker{name: "a", err: errors.New("failed")},
214 | ShutdownMarker{name: "b"})
215 | m.SetTimeout(120 * time.Millisecond) // time to complete a but not b
216 | },
217 | ExpectedEventTime: MustParseEventTime("a:100ms"),
218 | },
219 | }
220 |
221 | for _, tt := range testCases {
222 | t.Run(tt.Name, func(t *testing.T) {
223 | useLogger := true
224 | shutdownCalled := false
225 | for i := 0; i < 2; i++ {
226 | monitor := &SequenceMonitor{}
227 | pipeline := shutdown.New()
228 | var logger shutdown.Logger
229 | logData := &bytes.Buffer{}
230 |
231 | if useLogger {
232 | logger = LogWrapper{Writer: logData}
233 | pipeline.SetLogger(logger)
234 | useLogger = false
235 | } else {
236 | pipeline.SetLogger(nil)
237 | }
238 |
239 | if !shutdownCalled {
240 | pipeline.SetCompletionFunc(func() {
241 | shutdownCalled = true
242 | })
243 | }
244 | pipeline.SetTimeout(tt.Timeout)
245 |
246 | tt.Setup(pipeline)
247 | monitor.StartRecording()
248 | pipeline.Trigger(contextWithMonitor(monitor))
249 | if !monitor.Matches(tt.ExpectedEventTime) {
250 | t.Logf("log data:\n%s", logData.String())
251 | t.Fatalf("expected event time:\n%s\ngot:\n%s", tt.ExpectedEventTime, monitor.EventTime())
252 | return
253 | }
254 | if !shutdownCalled {
255 | t.Fatal("shutdown function did not get called")
256 | }
257 | }
258 | })
259 | }
260 | }
261 |
262 | func TestDefault(t *testing.T) {
263 | var d *shutdown.Manager
264 | wg := sync.WaitGroup{}
265 | wg.Add(10)
266 | broken := false
267 | for i := 0; i < 10; i++ {
268 | go func() {
269 | defer wg.Done()
270 | mgr := shutdown.Default()
271 | if d != nil && mgr != d {
272 | broken = true
273 | }
274 | d = mgr
275 | }()
276 | }
277 | wg.Wait()
278 | if broken {
279 | t.Fatal("default pointer changed")
280 | return
281 | }
282 |
283 | monitor := &SequenceMonitor{}
284 | monitor.StartRecording()
285 |
286 | shutdownCalled := false
287 | shutdown.AddSteps(ShutdownMarker{name: "1"}, ShutdownMarker{name: "2"})
288 | shutdown.AddSequence(ShutdownMarker{name: "3"}, ShutdownMarker{name: "4"})
289 | shutdown.AddParallelSequence(ShutdownMarker{name: "5"}, ShutdownMarker{name: "6"})
290 | shutdown.SetTimeout(350 * time.Millisecond)
291 | shutdown.SetLogger(nil)
292 | shutdown.SetCompletionFunc(func() {
293 | shutdownCalled = true
294 | })
295 | shutdown.Trigger(contextWithMonitor(monitor))
296 | if !shutdownCalled {
297 | t.Fatalf("shutdown did not get called")
298 | }
299 |
300 | expectation := MustParseEventTime("1:100ms,2:100ms,3:200ms,4:300ms")
301 | if !monitor.Matches(expectation) {
302 | t.Fatalf("expected:\n%s\ngot:\n%s", expectation, monitor.EventTime())
303 | }
304 | }
305 |
306 | func TestHandlerFunc(t *testing.T) {
307 | // just make sure the handle functions get compiled.
308 | manager := shutdown.New()
309 | manager.AddSteps(shutdown.HandlerFuncWithName("no-context-no-error", func() {}))
310 | manager.AddSteps(shutdown.HandlerFuncWithName("context-no-error", func(context.Context) {}))
311 | manager.AddSteps(shutdown.HandlerFuncWithName("no-context-error", func(context.Context) error { return nil }))
312 | manager.AddSteps(shutdown.HandlerFuncWithName("context-error", func(context.Context) error { return nil }))
313 | }
314 |
--------------------------------------------------------------------------------