gitech

1
// Copyright 2019 The Gitea Authors. All rights reserved.
2
// SPDX-License-Identifier: MIT
3

4
package graceful
5

6
import (
7
	"context"
8
	"runtime/pprof"
9
	"sync"
10
	"time"
11

12
	"code.gitea.io/gitea/modules/log"
13
	"code.gitea.io/gitea/modules/process"
14
	"code.gitea.io/gitea/modules/setting"
15
)
16

17
type state uint8
18

19
const (
20
	stateInit state = iota
21
	stateRunning
22
	stateShuttingDown
23
	stateTerminate
24
)
25

26
type RunCanceler interface {
27
	Run()
28
	Cancel()
29
}
30

31
// There are some places that could inherit sockets:
32
//
33
// * HTTP or HTTPS main listener
34
// * HTTP or HTTPS install listener
35
// * HTTP redirection fallback
36
// * Builtin SSH listener
37
//
38
// If you add a new place you must increment this number
39
// and add a function to call manager.InformCleanup if it's not going to be used
40
const numberOfServersToCreate = 4
41

42
var (
43
	manager  *Manager
44
	initOnce sync.Once
45
)
46

47
// GetManager returns the Manager
48
func GetManager() *Manager {
49
	InitManager(context.Background())
50
	return manager
51
}
52

53
// InitManager creates the graceful manager in the provided context
54
func InitManager(ctx context.Context) {
55
	initOnce.Do(func() {
56
		manager = newGracefulManager(ctx)
57

58
		// Set the process default context to the HammerContext
59
		process.DefaultContext = manager.HammerContext()
60
	})
61
}
62

63
// RunWithCancel helps to run a function with a custom context, the Cancel function will be called at shutdown
64
// The Cancel function should stop the Run function in predictable time.
65
func (g *Manager) RunWithCancel(rc RunCanceler) {
66
	g.RunAtShutdown(context.Background(), rc.Cancel)
67
	g.runningServerWaitGroup.Add(1)
68
	defer g.runningServerWaitGroup.Done()
69
	defer func() {
70
		if err := recover(); err != nil {
71
			log.Critical("PANIC during RunWithCancel: %v\nStacktrace: %s", err, log.Stack(2))
72
			g.doShutdown()
73
		}
74
	}()
75
	rc.Run()
76
}
77

78
// RunWithShutdownContext takes a function that has a context to watch for shutdown.
79
// After the provided context is Done(), the main function must return once shutdown is complete.
80
// (Optionally the HammerContext may be obtained and waited for however, this should be avoided if possible.)
81
func (g *Manager) RunWithShutdownContext(run func(context.Context)) {
82
	g.runningServerWaitGroup.Add(1)
83
	defer g.runningServerWaitGroup.Done()
84
	defer func() {
85
		if err := recover(); err != nil {
86
			log.Critical("PANIC during RunWithShutdownContext: %v\nStacktrace: %s", err, log.Stack(2))
87
			g.doShutdown()
88
		}
89
	}()
90
	ctx := g.ShutdownContext()
91
	pprof.SetGoroutineLabels(ctx) // We don't have a label to restore back to but I think this is fine
92
	run(ctx)
93
}
94

95
// RunAtTerminate adds to the terminate wait group and creates a go-routine to run the provided function at termination
96
func (g *Manager) RunAtTerminate(terminate func()) {
97
	g.terminateWaitGroup.Add(1)
98
	g.lock.Lock()
99
	defer g.lock.Unlock()
100
	g.toRunAtTerminate = append(g.toRunAtTerminate,
101
		func() {
102
			defer g.terminateWaitGroup.Done()
103
			defer func() {
104
				if err := recover(); err != nil {
105
					log.Critical("PANIC during RunAtTerminate: %v\nStacktrace: %s", err, log.Stack(2))
106
				}
107
			}()
108
			terminate()
109
		})
110
}
111

112
// RunAtShutdown creates a go-routine to run the provided function at shutdown
113
func (g *Manager) RunAtShutdown(ctx context.Context, shutdown func()) {
114
	g.lock.Lock()
115
	defer g.lock.Unlock()
116
	g.toRunAtShutdown = append(g.toRunAtShutdown,
117
		func() {
118
			defer func() {
119
				if err := recover(); err != nil {
120
					log.Critical("PANIC during RunAtShutdown: %v\nStacktrace: %s", err, log.Stack(2))
121
				}
122
			}()
123
			select {
124
			case <-ctx.Done():
125
				return
126
			default:
127
				shutdown()
128
			}
129
		})
130
}
131

132
func (g *Manager) doShutdown() {
133
	if !g.setStateTransition(stateRunning, stateShuttingDown) {
134
		g.DoImmediateHammer()
135
		return
136
	}
137
	g.lock.Lock()
138
	g.shutdownCtxCancel()
139
	atShutdownCtx := pprof.WithLabels(g.hammerCtx, pprof.Labels("graceful-lifecycle", "post-shutdown"))
140
	pprof.SetGoroutineLabels(atShutdownCtx)
141
	for _, fn := range g.toRunAtShutdown {
142
		go fn()
143
	}
144
	g.lock.Unlock()
145

146
	if setting.GracefulHammerTime >= 0 {
147
		go g.doHammerTime(setting.GracefulHammerTime)
148
	}
149
	go func() {
150
		g.runningServerWaitGroup.Wait()
151
		// Mop up any remaining unclosed events.
152
		g.doHammerTime(0)
153
		<-time.After(1 * time.Second)
154
		g.doTerminate()
155
		g.terminateWaitGroup.Wait()
156
		g.lock.Lock()
157
		g.managerCtxCancel()
158
		g.lock.Unlock()
159
	}()
160
}
161

162
func (g *Manager) doHammerTime(d time.Duration) {
163
	time.Sleep(d)
164
	g.lock.Lock()
165
	select {
166
	case <-g.hammerCtx.Done():
167
	default:
168
		log.Warn("Setting Hammer condition")
169
		g.hammerCtxCancel()
170
		atHammerCtx := pprof.WithLabels(g.terminateCtx, pprof.Labels("graceful-lifecycle", "post-hammer"))
171
		pprof.SetGoroutineLabels(atHammerCtx)
172
	}
173
	g.lock.Unlock()
174
}
175

176
func (g *Manager) doTerminate() {
177
	if !g.setStateTransition(stateShuttingDown, stateTerminate) {
178
		return
179
	}
180
	g.lock.Lock()
181
	select {
182
	case <-g.terminateCtx.Done():
183
	default:
184
		log.Warn("Terminating")
185
		g.terminateCtxCancel()
186
		atTerminateCtx := pprof.WithLabels(g.managerCtx, pprof.Labels("graceful-lifecycle", "post-terminate"))
187
		pprof.SetGoroutineLabels(atTerminateCtx)
188

189
		for _, fn := range g.toRunAtTerminate {
190
			go fn()
191
		}
192
	}
193
	g.lock.Unlock()
194
}
195

196
// IsChild returns if the current process is a child of previous Gitea process
197
func (g *Manager) IsChild() bool {
198
	return g.isChild
199
}
200

201
// IsShutdown returns a channel which will be closed at shutdown.
202
// The order of closure is shutdown, hammer (potentially), terminate
203
func (g *Manager) IsShutdown() <-chan struct{} {
204
	return g.shutdownCtx.Done()
205
}
206

207
// IsHammer returns a channel which will be closed at hammer.
208
// Servers running within the running server wait group should respond to IsHammer
209
// if not shutdown already
210
func (g *Manager) IsHammer() <-chan struct{} {
211
	return g.hammerCtx.Done()
212
}
213

214
// ServerDone declares a running server done and subtracts one from the
215
// running server wait group. Users probably do not want to call this
216
// and should use one of the RunWithShutdown* functions
217
func (g *Manager) ServerDone() {
218
	g.runningServerWaitGroup.Done()
219
}
220

221
func (g *Manager) setStateTransition(old, new state) bool {
222
	g.lock.Lock()
223
	if g.state != old {
224
		g.lock.Unlock()
225
		return false
226
	}
227
	g.state = new
228
	g.lock.Unlock()
229
	return true
230
}
231

232
// InformCleanup tells the cleanup wait group that we have either taken a listener or will not be taking a listener.
233
// At the moment the total number of servers (numberOfServersToCreate) are pre-defined as a const before global init,
234
// so this function MUST be called if a server is not used.
235
func (g *Manager) InformCleanup() {
236
	g.createServerCond.L.Lock()
237
	defer g.createServerCond.L.Unlock()
238
	g.createdServer++
239
	g.createServerCond.Signal()
240
}
241

242
// Done allows the manager to be viewed as a context.Context, it returns a channel that is closed when the server is finished terminating
243
func (g *Manager) Done() <-chan struct{} {
244
	return g.managerCtx.Done()
245
}
246

247
// Err allows the manager to be viewed as a context.Context done at Terminate
248
func (g *Manager) Err() error {
249
	return g.managerCtx.Err()
250
}
251

252
// Value allows the manager to be viewed as a context.Context done at Terminate
253
func (g *Manager) Value(key any) any {
254
	return g.managerCtx.Value(key)
255
}
256

257
// Deadline returns nil as there is no fixed Deadline for the manager, it allows the manager to be viewed as a context.Context
258
func (g *Manager) Deadline() (deadline time.Time, ok bool) {
259
	return g.managerCtx.Deadline()
260
}
261