2
2
// Use of this source code is governed by a BSD-style
3
3
// license that can be found in the LICENSE file.
4
4
5
- // Package singleflight provides a duplicate function call suppression
6
- // mechanism.
7
5
package singleflight
8
6
9
- import "sync"
7
+ import (
8
+ "bytes"
9
+ "errors"
10
+ "fmt"
11
+ "runtime"
12
+ "runtime/debug"
13
+ "sync"
14
+ )
15
+
16
+ // errGoexit indicates the runtime.Goexit was called in
17
+ // the user given function.
18
+ var errGoexit = errors .New ("runtime.Goexit was called" )
19
+
20
+ // A panicError is an arbitrary value recovered from a panic
21
+ // with the stack trace during the execution of given function.
22
+ type panicError struct {
23
+ value interface {}
24
+ stack []byte
25
+ }
26
+
27
+ // Error implements error interface.
28
+ func (p * panicError ) Error () string {
29
+ return fmt .Sprintf ("%v\n \n %s" , p .value , p .stack )
30
+ }
31
+
32
+ func newPanicError (v interface {}) error {
33
+ stack := debug .Stack ()
34
+
35
+ // The first line of the stack trace is of the form "goroutine N [status]:"
36
+ // but by the time the panic reaches Do the goroutine may no longer exist
37
+ // and its status will have changed. Trim out the misleading line.
38
+ if line := bytes .IndexByte (stack [:], '\n' ); line >= 0 {
39
+ stack = stack [line + 1 :]
40
+ }
41
+ return & panicError {value : v , stack : stack }
42
+ }
10
43
11
44
// call is an in-flight or completed singleflight.Do call
12
45
type call struct {
@@ -57,6 +90,12 @@ func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, e
57
90
c .dups ++
58
91
g .mu .Unlock ()
59
92
c .wg .Wait ()
93
+
94
+ if e , ok := c .err .(* panicError ); ok {
95
+ panic (e )
96
+ } else if c .err == errGoexit {
97
+ runtime .Goexit ()
98
+ }
60
99
return c .val , c .err , true
61
100
}
62
101
c := new (call )
@@ -70,6 +109,8 @@ func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, e
70
109
71
110
// DoChan is like Do but returns a channel that will receive the
72
111
// results when they are ready.
112
+ //
113
+ // The returned channel will not be closed.
73
114
func (g * Group ) DoChan (key string , fn func () (interface {}, error )) <- chan Result {
74
115
ch := make (chan Result , 1 )
75
116
g .mu .Lock ()
@@ -94,17 +135,66 @@ func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result
94
135
95
136
// doCall handles the single call for a key.
96
137
func (g * Group ) doCall (c * call , key string , fn func () (interface {}, error )) {
97
- c .val , c .err = fn ()
98
- c .wg .Done ()
99
-
100
- g .mu .Lock ()
101
- if ! c .forgotten {
102
- delete (g .m , key )
103
- }
104
- for _ , ch := range c .chans {
105
- ch <- Result {c .val , c .err , c .dups > 0 }
138
+ normalReturn := false
139
+ recovered := false
140
+
141
+ // use double-defer to distinguish panic from runtime.Goexit,
142
+ // more details see https://golang.org/cl/134395
143
+ defer func () {
144
+ // the given function invoked runtime.Goexit
145
+ if ! normalReturn && ! recovered {
146
+ c .err = errGoexit
147
+ }
148
+
149
+ c .wg .Done ()
150
+ g .mu .Lock ()
151
+ defer g .mu .Unlock ()
152
+ if ! c .forgotten {
153
+ delete (g .m , key )
154
+ }
155
+
156
+ if e , ok := c .err .(* panicError ); ok {
157
+ // In order to prevent the waiting channels from being blocked forever,
158
+ // needs to ensure that this panic cannot be recovered.
159
+ if len (c .chans ) > 0 {
160
+ go panic (e )
161
+ select {} // Keep this goroutine around so that it will appear in the crash dump.
162
+ } else {
163
+ panic (e )
164
+ }
165
+ } else if c .err == errGoexit {
166
+ // Already in the process of goexit, no need to call again
167
+ } else {
168
+ // Normal return
169
+ for _ , ch := range c .chans {
170
+ ch <- Result {c .val , c .err , c .dups > 0 }
171
+ }
172
+ }
173
+ }()
174
+
175
+ func () {
176
+ defer func () {
177
+ if ! normalReturn {
178
+ // Ideally, we would wait to take a stack trace until we've determined
179
+ // whether this is a panic or a runtime.Goexit.
180
+ //
181
+ // Unfortunately, the only way we can distinguish the two is to see
182
+ // whether the recover stopped the goroutine from terminating, and by
183
+ // the time we know that, the part of the stack trace relevant to the
184
+ // panic has been discarded.
185
+ if r := recover (); r != nil {
186
+ c .err = newPanicError (r )
187
+ }
188
+ }
189
+ }()
190
+
191
+ c .val , c .err = fn ()
192
+ normalReturn = true
193
+ }()
194
+
195
+ if ! normalReturn {
196
+ recovered = true
106
197
}
107
- g .mu .Unlock ()
108
198
}
109
199
110
200
// Forget tells the singleflight to forget about a key. Future calls
0 commit comments