forked from pingcap/tidb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
service_manager.go
121 lines (108 loc) · 3.46 KB
/
service_manager.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package sync2
import (
"sync"
)
// These are the three predefined states of a service.
const (
SERVICE_STOPPED = iota
SERVICE_RUNNING
SERVICE_SHUTTING_DOWN
)
var stateNames = []string{
"Stopped",
"Running",
"ShuttingDown",
}
// ServiceManager manages the state of a service through its lifecycle.
type ServiceManager struct {
mu sync.Mutex
wg sync.WaitGroup
err error // err is the error returned from the service function.
state AtomicInt64
// shutdown is created when the service starts and is closed when the service
// enters the SERVICE_SHUTTING_DOWN state.
shutdown chan struct{}
}
// Go tries to change the state from SERVICE_STOPPED to SERVICE_RUNNING.
//
// If the current state is not SERVICE_STOPPED (already running), it returns
// false immediately.
//
// On successful transition, it launches the service as a goroutine and returns
// true. The service function is responsible for returning on its own when
// requested, either by regularly checking svc.IsRunning(), or by waiting for
// the svc.ShuttingDown channel to be closed.
//
// When the service func returns, the state is reverted to SERVICE_STOPPED.
func (svm *ServiceManager) Go(service func(svc *ServiceContext) error) bool {
svm.mu.Lock()
defer svm.mu.Unlock()
if !svm.state.CompareAndSwap(SERVICE_STOPPED, SERVICE_RUNNING) {
return false
}
svm.wg.Add(1)
svm.err = nil
svm.shutdown = make(chan struct{})
go func() {
svm.err = service(&ServiceContext{ShuttingDown: svm.shutdown})
svm.state.Set(SERVICE_STOPPED)
svm.wg.Done()
}()
return true
}
// Stop tries to change the state from SERVICE_RUNNING to SERVICE_SHUTTING_DOWN.
// If the current state is not SERVICE_RUNNING, it returns false immediately.
// On successul transition, it waits for the service to finish, and returns true.
// You are allowed to Go() again after a Stop().
func (svm *ServiceManager) Stop() bool {
svm.mu.Lock()
defer svm.mu.Unlock()
if !svm.state.CompareAndSwap(SERVICE_RUNNING, SERVICE_SHUTTING_DOWN) {
return false
}
// Signal the service that we've transitioned to SERVICE_SHUTTING_DOWN.
close(svm.shutdown)
svm.shutdown = nil
svm.wg.Wait()
return true
}
// Wait waits for the service to terminate if it's currently running.
func (svm *ServiceManager) Wait() {
svm.wg.Wait()
}
// Join waits for the service to terminate and returns the value returned by the
// service function.
func (svm *ServiceManager) Join() error {
svm.wg.Wait()
return svm.err
}
// State returns the current state of the service.
// This should only be used to report the current state.
func (svm *ServiceManager) State() int64 {
return svm.state.Get()
}
// StateName returns the name of the current state.
func (svm *ServiceManager) StateName() string {
return stateNames[svm.State()]
}
// ServiceContext is passed into the service function to give it access to
// information about the running service.
type ServiceContext struct {
// ShuttingDown is a channel that the service can select on to be notified
// when it should shut down. The channel is closed when the state transitions
// from SERVICE_RUNNING to SERVICE_SHUTTING_DOWN.
ShuttingDown chan struct{}
}
// IsRunning returns true if the ServiceContext.ShuttingDown channel has not
// been closed yet.
func (svc *ServiceContext) IsRunning() bool {
select {
case <-svc.ShuttingDown:
return false
default:
return true
}
}