/
ctxgroup.go
257 lines (218 loc) · 7.71 KB
/
ctxgroup.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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
// package ctxgroup provides the ContextGroup, a hybrid between the
// context.Context and sync.WaitGroup, which models process trees.
package ctxgroup
import (
"io"
"sync"
context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
)
// TeardownFunc is a function used to cleanup state at the end of the
// lifecycle of a process.
type TeardownFunc func() error
// ChildFunc is a function to register as a child. It will be automatically
// tracked.
type ChildFunc func(parent ContextGroup)
var nilTeardownFunc = func() error { return nil }
// ContextGroup is an interface for services able to be opened and closed.
// It has a parent Context, and Children. But ContextGroup is not a proper
// "tree" like the Context tree. It is more like a Context-WaitGroup hybrid.
// It models a main object with a few children objects -- and, unlike the
// context -- concerns itself with the parent-child closing semantics:
//
// - Can define an optional TeardownFunc (func() error) to be run at Close time.
// - Children call Children().Add(1) to be waited upon
// - Children can select on <-Closing() to know when they should shut down.
// - Close() will wait until all children call Children().Done()
// - <-Closed() signals when the service is completely closed.
//
// ContextGroup can be embedded into the main object itself. In that case,
// the teardownFunc (if a member function) has to be set after the struct
// is intialized:
//
// type service struct {
// ContextGroup
// net.Conn
// }
//
// func (s *service) close() error {
// return s.Conn.Close()
// }
//
// func newService(ctx context.Context, c net.Conn) *service {
// s := &service{c}
// s.ContextGroup = NewContextGroup(ctx, s.close)
// return s
// }
//
type ContextGroup interface {
// Context is the context of this ContextGroup. It is "sort of" a parent.
Context() context.Context
// SetTeardown assigns the teardown function.
// It is called exactly _once_ when the ContextGroup is Closed.
SetTeardown(tf TeardownFunc)
// Children is a sync.Waitgroup for all children goroutines that should
// shut down completely before this service is said to be "closed".
// Follows the semantics of WaitGroup:
//
// Children().Add(1) // add one more dependent child
// Children().Done() // child signals it is done
//
// WARNING: this is deprecated and will go away soon.
Children() *sync.WaitGroup
// AddChild gives ownership of a child io.Closer. The child will be closed
// when the context group is closed.
AddChild(io.Closer)
// AddChildFunc registers a dependent ChildFund. The child will receive
// its parent ContextGroup, and can wait on its signals. Child references
// tracked automatically. It equivalent to the following:
//
// go func(parent, child ContextGroup) {
//
// <-parent.Closing() // wait until parent is closing
// child.Close() // signal child to close
// parent.Children().Done() // child signals it is done
// }(a, b)
//
AddChildFunc(c ChildFunc)
// Close is a method to call when you wish to stop this ContextGroup
Close() error
// Closing is a signal to wait upon, like Context.Done().
// It fires when the object should be closing (but hasn't yet fully closed).
// The primary use case is for child goroutines who need to know when
// they should shut down. (equivalent to Context().Done())
Closing() <-chan struct{}
// Closed is a method to wait upon, like Context.Done().
// It fires when the entire object is fully closed.
// The primary use case is for external listeners who need to know when
// this object is completly done, and all its children closed.
Closed() <-chan struct{}
}
// contextGroup is a Closer with a cancellable context
type contextGroup struct {
ctx context.Context
cancel context.CancelFunc
// called to run the teardown logic.
teardownFunc TeardownFunc
// closed is released once the close function is done.
closed chan struct{}
// wait group for child goroutines
children sync.WaitGroup
// sync primitive to ensure the close logic is only called once.
closeOnce sync.Once
// error to return to clients of Close().
closeErr error
}
// newContextGroup constructs and returns a ContextGroup. It will call
// cf TeardownFunc before its Done() Wait signals fire.
func newContextGroup(ctx context.Context, cf TeardownFunc) ContextGroup {
ctx, cancel := context.WithCancel(ctx)
c := &contextGroup{
ctx: ctx,
cancel: cancel,
closed: make(chan struct{}),
}
c.SetTeardown(cf)
c.Children().Add(1) // initialize with 1. calling Close will decrement it.
go c.closeOnContextDone()
return c
}
// SetTeardown assigns the teardown function.
func (c *contextGroup) SetTeardown(cf TeardownFunc) {
if cf == nil {
cf = nilTeardownFunc
}
c.teardownFunc = cf
}
func (c *contextGroup) Context() context.Context {
return c.ctx
}
func (c *contextGroup) Children() *sync.WaitGroup {
return &c.children
}
func (c *contextGroup) AddChild(child io.Closer) {
c.children.Add(1)
go func(parent ContextGroup, child io.Closer) {
<-parent.Closing() // wait until parent is closing
child.Close() // signal child to close
parent.Children().Done() // child signals it is done
}(c, child)
}
func (c *contextGroup) AddChildFunc(child ChildFunc) {
c.children.Add(1)
go func(parent ContextGroup, child ChildFunc) {
child(parent)
parent.Children().Done() // child signals it is done
}(c, child)
}
// Close is the external close function. it's a wrapper around internalClose
// that waits on Closed()
func (c *contextGroup) Close() error {
c.internalClose()
<-c.Closed() // wait until we're totally done.
return c.closeErr
}
func (c *contextGroup) Closing() <-chan struct{} {
return c.Context().Done()
}
func (c *contextGroup) Closed() <-chan struct{} {
return c.closed
}
func (c *contextGroup) internalClose() {
go c.closeOnce.Do(c.closeLogic)
}
// the _actual_ close process.
func (c *contextGroup) closeLogic() {
// this function should only be called once (hence the sync.Once).
// and it will panic at the bottom (on close(c.closed)) otherwise.
c.cancel() // signal that we're shutting down (Closing)
c.closeErr = c.teardownFunc() // actually run the close logic
c.children.Wait() // wait till all children are done.
close(c.closed) // signal that we're shut down (Closed)
}
// if parent context is shut down before we call Close explicitly,
// we need to go through the Close motions anyway. Hence all the sync
// stuff all over the place...
func (c *contextGroup) closeOnContextDone() {
<-c.Context().Done() // wait until parent (context) is done.
c.internalClose()
c.Children().Done()
}
// WithTeardown constructs and returns a ContextGroup with
// cf TeardownFunc (and context.Background)
func WithTeardown(cf TeardownFunc) ContextGroup {
if cf == nil {
panic("nil TeardownFunc")
}
return newContextGroup(context.Background(), cf)
}
// WithContext constructs and returns a ContextGroup with given context
func WithContext(ctx context.Context) ContextGroup {
if ctx == nil {
panic("nil Context")
}
return newContextGroup(ctx, nil)
}
// WithContextAndTeardown constructs and returns a ContextGroup with
// cf TeardownFunc (and context.Background)
func WithContextAndTeardown(ctx context.Context, cf TeardownFunc) ContextGroup {
if ctx == nil {
panic("nil Context")
}
if cf == nil {
panic("nil TeardownFunc")
}
return newContextGroup(ctx, cf)
}
// WithParent constructs and returns a ContextGroup with given parent
func WithParent(p ContextGroup) ContextGroup {
if p == nil {
panic("nil ContextGroup")
}
c := newContextGroup(p.Context(), nil)
p.AddChild(c)
return c
}
// WithBackground returns a ContextGroup with context.Background()
func WithBackground() ContextGroup {
return newContextGroup(context.Background(), nil)
}