@@ -3,7 +3,6 @@ package mapreduce
33import (
44 "context"
55 "errors"
6- "fmt"
76 "sync"
87 "sync/atomic"
98)
@@ -21,12 +20,12 @@ var (
2120)
2221
2322type (
23+ // ForEachFunc is used to do element processing, but no output.
24+ ForEachFunc func (item interface {})
2425 // GenerateFunc is used to let callers send elements into source.
2526 GenerateFunc func (source chan <- interface {})
2627 // MapFunc is used to do element processing and write the output to writer.
2728 MapFunc func (item interface {}, writer Writer )
28- // VoidMapFunc is used to do element processing, but no output.
29- VoidMapFunc func (item interface {})
3029 // MapperFunc is used to do element processing and write the output to writer,
3130 // use cancel func to cancel the processing.
3231 MapperFunc func (item interface {}, writer Writer , cancel func (error ))
@@ -39,6 +38,16 @@ type (
3938 // Option defines the method to customize the mapreduce.
4039 Option func (opts * mapReduceOptions )
4140
41+ mapperContext struct {
42+ ctx context.Context
43+ mapper MapFunc
44+ source <- chan interface {}
45+ panicChan * onceChan
46+ collector chan <- interface {}
47+ doneChan <- chan struct {}
48+ workers int
49+ }
50+
4251 mapReduceOptions struct {
4352 ctx context.Context
4453 workers int
@@ -76,7 +85,7 @@ func FinishVoid(fns ...func()) {
7685 return
7786 }
7887
79- MapVoid (func (source chan <- interface {}) {
88+ ForEach (func (source chan <- interface {}) {
8089 for _ , fn := range fns {
8190 source <- fn
8291 }
@@ -86,38 +95,72 @@ func FinishVoid(fns ...func()) {
8695 }, WithWorkers (len (fns )))
8796}
8897
89- // Map maps all elements generated from given generate func, and returns an output channel .
90- func Map (generate GenerateFunc , mapper MapFunc , opts ... Option ) chan interface {} {
98+ // ForEach maps all elements from given generate but no output.
99+ func ForEach (generate GenerateFunc , mapper ForEachFunc , opts ... Option ) {
91100 options := buildOptions (opts ... )
92- source := buildSource (generate )
101+ panicChan := make (chan interface {})
102+ source := buildSource (generate , panicChan )
93103 collector := make (chan interface {}, options .workers )
94104 done := make (chan struct {})
95105
96- go executeMappers (options .ctx , mapper , source , collector , done , options .workers )
106+ go executeMappers (mapperContext {
107+ ctx : options .ctx ,
108+ mapper : func (item interface {}, writer Writer ) {
109+ mapper (item )
110+ },
111+ source : source ,
112+ panicChan : & onceChan {
113+ channel : panicChan ,
114+ },
115+ collector : collector ,
116+ doneChan : done ,
117+ workers : options .workers ,
118+ })
97119
98- return collector
120+ for {
121+ select {
122+ case v := <- panicChan :
123+ panic (v )
124+ case _ , ok := <- collector :
125+ if ! ok {
126+ return
127+ }
128+ }
129+ }
99130}
100131
101132// MapReduce maps all elements generated from given generate func,
102133// and reduces the output elements with given reducer.
103134func MapReduce (generate GenerateFunc , mapper MapperFunc , reducer ReducerFunc ,
104135 opts ... Option ) (interface {}, error ) {
105- source := buildSource (generate )
106- return MapReduceWithSource (source , mapper , reducer , opts ... )
136+ panicChan := make (chan interface {})
137+ source := buildSource (generate , panicChan )
138+ return mapReduceWithPanicChan (source , panicChan , mapper , reducer , opts ... )
107139}
108140
109- // MapReduceWithSource maps all elements from source, and reduce the output elements with given reducer.
110- func MapReduceWithSource (source <- chan interface {}, mapper MapperFunc , reducer ReducerFunc ,
141+ // MapReduceChan maps all elements from source, and reduce the output elements with given reducer.
142+ func MapReduceChan (source <- chan interface {}, mapper MapperFunc , reducer ReducerFunc ,
111143 opts ... Option ) (interface {}, error ) {
144+ panicChan := make (chan interface {})
145+ return mapReduceWithPanicChan (source , panicChan , mapper , reducer , opts ... )
146+ }
147+
148+ // MapReduceChan maps all elements from source, and reduce the output elements with given reducer.
149+ func mapReduceWithPanicChan (source <- chan interface {}, panicChan chan interface {}, mapper MapperFunc ,
150+ reducer ReducerFunc , opts ... Option ) (interface {}, error ) {
112151 options := buildOptions (opts ... )
152+ // output is used to write the final result
113153 output := make (chan interface {})
114154 defer func () {
155+ // reducer can only write once, if more, panic
115156 for range output {
116157 panic ("more than one element written in reducer" )
117158 }
118159 }()
119160
161+ // collector is used to collect data from mapper, and consume in reducer
120162 collector := make (chan interface {}, options .workers )
163+ // if done is closed, all mappers and reducer should stop processing
121164 done := make (chan struct {})
122165 writer := newGuardedWriter (options .ctx , output , done )
123166 var closeOnce sync.Once
@@ -142,30 +185,47 @@ func MapReduceWithSource(source <-chan interface{}, mapper MapperFunc, reducer R
142185
143186 go func () {
144187 defer func () {
145- drain (collector )
146-
147188 if r := recover (); r != nil {
148- cancel ( fmt . Errorf ( "%v" , r ))
189+ panicChan <- r
149190 } else {
150191 finish ()
151192 }
193+
194+ drain (collector )
152195 }()
153196
154- // callers need to make sure reducer not panic
155197 reducer (collector , writer , cancel )
156198 }()
157199
158- go executeMappers (options .ctx , func (item interface {}, w Writer ) {
159- mapper (item , w , cancel )
160- }, source , collector , done , options .workers )
161-
162- value , ok := <- output
163- if err := retErr .Load (); err != nil {
164- return nil , err .(error )
165- } else if ok {
166- return value , nil
167- } else {
168- return nil , ErrReduceNoOutput
200+ go executeMappers (mapperContext {
201+ ctx : options .ctx ,
202+ mapper : func (item interface {}, w Writer ) {
203+ mapper (item , w , cancel )
204+ },
205+ source : source ,
206+ panicChan : & onceChan {
207+ channel : panicChan ,
208+ },
209+ collector : collector ,
210+ doneChan : done ,
211+ workers : options .workers ,
212+ })
213+
214+ select {
215+ case <- options .ctx .Done ():
216+ cancel (context .DeadlineExceeded )
217+ return nil , context .DeadlineExceeded
218+ case v := <- panicChan :
219+ finish ()
220+ panic (v )
221+ case v , ok := <- output :
222+ if err := retErr .Load (); err != nil {
223+ return nil , err .(error )
224+ } else if ok {
225+ return v , nil
226+ } else {
227+ return nil , ErrReduceNoOutput
228+ }
169229 }
170230}
171231
@@ -175,20 +235,12 @@ func MapReduceVoid(generate GenerateFunc, mapper MapperFunc, reducer VoidReducer
175235 _ , err := MapReduce (generate , mapper , func (input <- chan interface {}, writer Writer , cancel func (error )) {
176236 reducer (input , cancel )
177237 // We need to write a placeholder to let MapReduce to continue on reducer done,
178- // otherwise, all goroutines are waiting.
179- // The placeholder will be discarded by MapReduce.
238+ // otherwise, all goroutines are waiting. The placeholder will be discarded by MapReduce.
180239 writer .Write (struct {}{})
181240 }, opts ... )
182241 return err
183242}
184243
185- // MapVoid maps all elements from given generate but no output.
186- func MapVoid (generate GenerateFunc , mapper VoidMapFunc , opts ... Option ) {
187- drain (Map (generate , func (item interface {}, writer Writer ) {
188- mapper (item )
189- }, opts ... ))
190- }
191-
192244// WithContext customizes a mapreduce processing accepts a given ctx.
193245func WithContext (ctx context.Context ) Option {
194246 return func (opts * mapReduceOptions ) {
@@ -216,10 +268,16 @@ func buildOptions(opts ...Option) *mapReduceOptions {
216268 return options
217269}
218270
219- func buildSource (generate GenerateFunc ) chan interface {} {
271+ func buildSource (generate GenerateFunc , panicChan chan interface {} ) chan interface {} {
220272 source := make (chan interface {})
221273 go func () {
222- defer close (source )
274+ defer func () {
275+ if r := recover (); r != nil {
276+ panicChan <- r
277+ }
278+ close (source )
279+ }()
280+
223281 generate (source )
224282 }()
225283
@@ -233,24 +291,31 @@ func drain(channel <-chan interface{}) {
233291 }
234292}
235293
236- func executeMappers (ctx context.Context , mapper MapFunc , input <- chan interface {},
237- collector chan <- interface {}, done <- chan struct {}, workers int ) {
294+ func executeMappers (mCtx mapperContext ) {
238295 var wg sync.WaitGroup
239296 defer func () {
240297 wg .Wait ()
241- close (collector )
298+ close (mCtx .collector )
299+ drain (mCtx .source )
242300 }()
243301
244- pool := make (chan struct {}, workers )
245- writer := newGuardedWriter (ctx , collector , done )
302+ pool := make (chan struct {}, mCtx .workers )
303+ innerPanicChan := make (chan interface {})
304+ onceInnerChan := & onceChan {
305+ channel : innerPanicChan ,
306+ }
307+ writer := newGuardedWriter (mCtx .ctx , mCtx .collector , mCtx .doneChan )
246308 for {
247309 select {
248- case <- ctx .Done ():
310+ case <- mCtx . ctx .Done ():
249311 return
250- case <- done :
312+ case <- mCtx .doneChan :
313+ return
314+ case v := <- innerPanicChan :
315+ mCtx .panicChan .write (v )
251316 return
252317 case pool <- struct {}{}:
253- item , ok := <- input
318+ item , ok := <- mCtx . source
254319 if ! ok {
255320 <- pool
256321 return
@@ -259,12 +324,14 @@ func executeMappers(ctx context.Context, mapper MapFunc, input <-chan interface{
259324 wg .Add (1 )
260325 go func () {
261326 defer func () {
327+ if r := recover (); r != nil {
328+ onceInnerChan .write (r )
329+ }
262330 wg .Done ()
263331 <- pool
264332 }()
265333
266- // callers need to make sure mapper won't panic
267- mapper (item , writer )
334+ mCtx .mapper (item , writer )
268335 }()
269336 }
270337 }
@@ -311,3 +378,16 @@ func (gw guardedWriter) Write(v interface{}) {
311378 gw .channel <- v
312379 }
313380}
381+
382+ type onceChan struct {
383+ channel chan <- interface {}
384+ wrote int32
385+ }
386+
387+ func (oc * onceChan ) write (val interface {}) {
388+ if atomic .AddInt32 (& oc .wrote , 1 ) > 1 {
389+ return
390+ }
391+
392+ oc .channel <- val
393+ }
0 commit comments