22// Use of this source code is governed by a BSD-style
33// license that can be found in the LICENSE file.
44
5- // Package singleflight provides a duplicate function call suppression
6- // mechanism.
75package singleflight
86
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+ }
1043
1144// call is an in-flight or completed singleflight.Do call
1245type call struct {
@@ -57,6 +90,12 @@ func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, e
5790 c .dups ++
5891 g .mu .Unlock ()
5992 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+ }
6099 return c .val , c .err , true
61100 }
62101 c := new (call )
@@ -70,6 +109,8 @@ func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, e
70109
71110// DoChan is like Do but returns a channel that will receive the
72111// results when they are ready.
112+ //
113+ // The returned channel will not be closed.
73114func (g * Group ) DoChan (key string , fn func () (interface {}, error )) <- chan Result {
74115 ch := make (chan Result , 1 )
75116 g .mu .Lock ()
@@ -94,17 +135,66 @@ func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result
94135
95136// doCall handles the single call for a key.
96137func (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
106197 }
107- g .mu .Unlock ()
108198}
109199
110200// Forget tells the singleflight to forget about a key. Future calls
0 commit comments