feat: handling panic in mapreduce, panic in calling goroutine, not inside goroutines (#1490)

* feat: handle panic

* chore: update fuzz test

* chore: optimize square sum algorithm

core/mr/mapreduce_test.go

@@ -11,8 +11,6 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/stringx"
-	"github.com/zeromicro/go-zero/core/syncx"
 	"go.uber.org/goleak"
 )
 
@@ -124,84 +122,69 @@ func TestForEach(t *testing.T) {
 	t.Run("all", func(t *testing.T) {
 		defer goleak.VerifyNone(t)
 
-		ForEach(func(source chan<- interface{}) {
-			for i := 0; i < tasks; i++ {
-				source <- i
-			}
-		}, func(item interface{}) {
-			panic("foo")
+		assert.PanicsWithValue(t, "foo", func() {
+			ForEach(func(source chan<- interface{}) {
+				for i := 0; i < tasks; i++ {
+					source <- i
+				}
+			}, func(item interface{}) {
+				panic("foo")
+			})
 		})
 	})
 }
 
-func TestMap(t *testing.T) {
+func TestGeneratePanic(t *testing.T) {
 	defer goleak.VerifyNone(t)
 
-	tests := []struct {
-		mapper MapFunc
-		expect int
-	}{
-		{
-			mapper: func(item interface{}, writer Writer) {
-				v := item.(int)
-				writer.Write(v * v)
-			},
-			expect: 30,
-		},
-		{
-			mapper: func(item interface{}, writer Writer) {
-				v := item.(int)
-				if v%2 == 0 {
-					return
-				}
-				writer.Write(v * v)
-			},
-			expect: 10,
-		},
-		{
-			mapper: func(item interface{}, writer Writer) {
-				v := item.(int)
-				if v%2 == 0 {
-					panic(v)
-				}
-				writer.Write(v * v)
-			},
-			expect: 10,
-		},
-	}
+	t.Run("all", func(t *testing.T) {
+		assert.PanicsWithValue(t, "foo", func() {
+			ForEach(func(source chan<- interface{}) {
+				panic("foo")
+			}, func(item interface{}) {
+			})
+		})
+	})
+}
 
-	for _, test := range tests {
-		t.Run(stringx.Rand(), func(t *testing.T) {
-			channel := Map(func(source chan<- interface{}) {
-				for i := 1; i < 5; i++ {
+func TestMapperPanic(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	const tasks = 1000
+	var run int32
+	t.Run("all", func(t *testing.T) {
+		assert.PanicsWithValue(t, "foo", func() {
+			_, _ = MapReduce(func(source chan<- interface{}) {
+				for i := 0; i < tasks; i++ {
 					source <- i
 				}
-			}, test.mapper, WithWorkers(-1))
-
-			var result int
-			for v := range channel {
-				result += v.(int)
-			}
-
-			assert.Equal(t, test.expect, result)
+			}, func(item interface{}, writer Writer, cancel func(error)) {
+				atomic.AddInt32(&run, 1)
+				panic("foo")
+			}, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
+			})
 		})
-	}
+		assert.True(t, atomic.LoadInt32(&run) < tasks/2)
+	})
 }
 
 func TestMapReduce(t *testing.T) {
 	defer goleak.VerifyNone(t)
 
 	tests := []struct {
+		name        string
 		mapper      MapperFunc
 		reducer     ReducerFunc
 		expectErr   error
 		expectValue interface{}
 	}{
 		{
+			name:        "simple",
 			expectErr:   nil,
 			expectValue: 30,
 		},
 		{
+			name: "cancel with error",
 			mapper: func(item interface{}, writer Writer, cancel func(error)) {
 				v := item.(int)
 				if v%3 == 0 {
@@ -212,6 +195,7 @@ func TestMapReduce(t *testing.T) {
 			expectErr: errDummy,
 		},
 		{
+			name: "cancel with nil",
 			mapper: func(item interface{}, writer Writer, cancel func(error)) {
 				v := item.(int)
 				if v%3 == 0 {
@@ -223,6 +207,7 @@ func TestMapReduce(t *testing.T) {
 			expectValue: nil,
 		},
 		{
+			name: "cancel with more",
 			reducer: func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
 				var result int
 				for item := range pipe {
@@ -237,45 +222,68 @@ func TestMapReduce(t *testing.T) {
 		},
 	}
 
-	for _, test := range tests {
-		t.Run(stringx.Rand(), func(t *testing.T) {
-			if test.mapper == nil {
-				test.mapper = func(item interface{}, writer Writer, cancel func(error)) {
-					v := item.(int)
-					writer.Write(v * v)
-				}
-			}
-			if test.reducer == nil {
-				test.reducer = func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
-					var result int
-					for item := range pipe {
-						result += item.(int)
+	t.Run("MapReduce", func(t *testing.T) {
+		for _, test := range tests {
+			t.Run(test.name, func(t *testing.T) {
+				if test.mapper == nil {
+					test.mapper = func(item interface{}, writer Writer, cancel func(error)) {
+						v := item.(int)
+						writer.Write(v * v)
 					}
-					writer.Write(result)
 				}
-			}
-			value, err := MapReduce(func(source chan<- interface{}) {
-				for i := 1; i < 5; i++ {
-					source <- i
+				if test.reducer == nil {
+					test.reducer = func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
+						var result int
+						for item := range pipe {
+							result += item.(int)
+						}
+						writer.Write(result)
+					}
 				}
-			}, test.mapper, test.reducer, WithWorkers(runtime.NumCPU()))
+				value, err := MapReduce(func(source chan<- interface{}) {
+					for i := 1; i < 5; i++ {
+						source <- i
+					}
+				}, test.mapper, test.reducer, WithWorkers(runtime.NumCPU()))
 
-			assert.Equal(t, test.expectErr, err)
-			assert.Equal(t, test.expectValue, value)
-		})
-	}
-}
+				assert.Equal(t, test.expectErr, err)
+				assert.Equal(t, test.expectValue, value)
+			})
+		}
+	})
 
-func TestMapReducePanicBothMapperAndReducer(t *testing.T) {
-	defer goleak.VerifyNone(t)
+	t.Run("MapReduce", func(t *testing.T) {
+		for _, test := range tests {
+			t.Run(test.name, func(t *testing.T) {
+				if test.mapper == nil {
+					test.mapper = func(item interface{}, writer Writer, cancel func(error)) {
+						v := item.(int)
+						writer.Write(v * v)
+					}
+				}
+				if test.reducer == nil {
+					test.reducer = func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
+						var result int
+						for item := range pipe {
+							result += item.(int)
+						}
+						writer.Write(result)
+					}
+				}
 
-	_, _ = MapReduce(func(source chan<- interface{}) {
-		source <- 0
-		source <- 1
-	}, func(item interface{}, writer Writer, cancel func(error)) {
-		panic("foo")
-	}, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
-		panic("bar")
+				source := make(chan interface{})
+				go func() {
+					for i := 1; i < 5; i++ {
+						source <- i
+					}
+					close(source)
+				}()
+
+				value, err := MapReduceChan(source, test.mapper, test.reducer, WithWorkers(-1))
+				assert.Equal(t, test.expectErr, err)
+				assert.Equal(t, test.expectValue, value)
+			})
+		}
 	})
 }
 
@@ -302,16 +310,19 @@ func TestMapReduceVoid(t *testing.T) {
 
 	var value uint32
 	tests := []struct {
+		name        string
 		mapper      MapperFunc
 		reducer     VoidReducerFunc
 		expectValue uint32
 		expectErr   error
 	}{
 		{
+			name:        "simple",
 			expectValue: 30,
 			expectErr:   nil,
 		},
 		{
+			name: "cancel with error",
 			mapper: func(item interface{}, writer Writer, cancel func(error)) {
 				v := item.(int)
 				if v%3 == 0 {
@@ -322,6 +333,7 @@ func TestMapReduceVoid(t *testing.T) {
 			expectErr: errDummy,
 		},
 		{
+			name: "cancel with nil",
 			mapper: func(item interface{}, writer Writer, cancel func(error)) {
 				v := item.(int)
 				if v%3 == 0 {
@@ -332,6 +344,7 @@ func TestMapReduceVoid(t *testing.T) {
 			expectErr: ErrCancelWithNil,
 		},
 		{
+			name: "cancel with more",
 			reducer: func(pipe <-chan interface{}, cancel func(error)) {
 				for item := range pipe {
 					result := atomic.AddUint32(&value, uint32(item.(int)))
@@ -345,7 +358,7 @@ func TestMapReduceVoid(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		t.Run(stringx.Rand(), func(t *testing.T) {
+		t.Run(test.name, func(t *testing.T) {
 			atomic.StoreUint32(&value, 0)
 
 			if test.mapper == nil {
@@ -400,39 +413,59 @@ func TestMapReduceVoidWithDelay(t *testing.T) {
 	assert.Equal(t, 0, result[1])
 }
 
-func TestMapVoid(t *testing.T) {
-	defer goleak.VerifyNone(t)
-
-	const tasks = 1000
-	var count uint32
-	ForEach(func(source chan<- interface{}) {
-		for i := 0; i < tasks; i++ {
-			source <- i
-		}
-	}, func(item interface{}) {
-		atomic.AddUint32(&count, 1)
-	})
-
-	assert.Equal(t, tasks, int(count))
-}
-
 func TestMapReducePanic(t *testing.T) {
 	defer goleak.VerifyNone(t)
 
-	v, err := MapReduce(func(source chan<- interface{}) {
-		source <- 0
-		source <- 1
-	}, func(item interface{}, writer Writer, cancel func(error)) {
-		i := item.(int)
-		writer.Write(i)
-	}, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
-		for range pipe {
-			panic("panic")
-		}
+	assert.Panics(t, func() {
+		_, _ = MapReduce(func(source chan<- interface{}) {
+			source <- 0
+			source <- 1
+		}, func(item interface{}, writer Writer, cancel func(error)) {
+			i := item.(int)
+			writer.Write(i)
+		}, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
+			for range pipe {
+				panic("panic")
+			}
+		})
+	})
+}
+
+func TestMapReducePanicOnce(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	assert.Panics(t, func() {
+		_, _ = MapReduce(func(source chan<- interface{}) {
+			for i := 0; i < 100; i++ {
+				source <- i
+			}
+		}, func(item interface{}, writer Writer, cancel func(error)) {
+			i := item.(int)
+			if i == 0 {
+				panic("foo")
+			}
+			writer.Write(i)
+		}, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
+			for range pipe {
+				panic("bar")
+			}
+		})
+	})
+}
+
+func TestMapReducePanicBothMapperAndReducer(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	assert.Panics(t, func() {
+		_, _ = MapReduce(func(source chan<- interface{}) {
+			source <- 0
+			source <- 1
+		}, func(item interface{}, writer Writer, cancel func(error)) {
+			panic("foo")
+		}, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
+			panic("bar")
+		})
 	})
-	assert.Nil(t, v)
-	assert.NotNil(t, err)
-	assert.Equal(t, "panic", err.Error())
 }
 
 func TestMapReduceVoidCancel(t *testing.T) {
@@ -461,13 +494,13 @@ func TestMapReduceVoidCancel(t *testing.T) {
 func TestMapReduceVoidCancelWithRemains(t *testing.T) {
 	defer goleak.VerifyNone(t)
 
-	var done syncx.AtomicBool
+	var done int32
 	var result []int
 	err := MapReduceVoid(func(source chan<- interface{}) {
 		for i := 0; i < defaultWorkers*2; i++ {
 			source <- i
 		}
-		done.Set(true)
+		atomic.AddInt32(&done, 1)
 	}, func(item interface{}, writer Writer, cancel func(error)) {
 		i := item.(int)
 		if i == defaultWorkers/2 {
@@ -482,7 +515,7 @@ func TestMapReduceVoidCancelWithRemains(t *testing.T) {
 	})
 	assert.NotNil(t, err)
 	assert.Equal(t, "anything", err.Error())
-	assert.True(t, done.True())
+	assert.Equal(t, int32(1), done)
 }
 
 func TestMapReduceWithoutReducerWrite(t *testing.T) {
@@ -507,34 +540,51 @@ func TestMapReduceVoidPanicInReducer(t *testing.T) {
 	defer goleak.VerifyNone(t)
 
 	const message = "foo"
-	var done syncx.AtomicBool
-	err := MapReduceVoid(func(source chan<- interface{}) {
+	assert.Panics(t, func() {
+		var done int32
+		_ = MapReduceVoid(func(source chan<- interface{}) {
+			for i := 0; i < defaultWorkers*2; i++ {
+				source <- i
+			}
+			atomic.AddInt32(&done, 1)
+		}, func(item interface{}, writer Writer, cancel func(error)) {
+			i := item.(int)
+			writer.Write(i)
+		}, func(pipe <-chan interface{}, cancel func(error)) {
+			panic(message)
+		}, WithWorkers(1))
+	})
+}
+
+func TestForEachWithContext(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	var done int32
+	ctx, cancel := context.WithCancel(context.Background())
+	ForEach(func(source chan<- interface{}) {
 		for i := 0; i < defaultWorkers*2; i++ {
 			source <- i
 		}
-		done.Set(true)
-	}, func(item interface{}, writer Writer, cancel func(error)) {
+		atomic.AddInt32(&done, 1)
+	}, func(item interface{}) {
 		i := item.(int)
-		writer.Write(i)
-	}, func(pipe <-chan interface{}, cancel func(error)) {
-		panic(message)
-	}, WithWorkers(1))
-	assert.NotNil(t, err)
-	assert.Equal(t, message, err.Error())
-	assert.True(t, done.True())
+		if i == defaultWorkers/2 {
+			cancel()
+		}
+	}, WithContext(ctx))
 }
 
 func TestMapReduceWithContext(t *testing.T) {
 	defer goleak.VerifyNone(t)
 
-	var done syncx.AtomicBool
+	var done int32
 	var result []int
 	ctx, cancel := context.WithCancel(context.Background())
 	err := MapReduceVoid(func(source chan<- interface{}) {
 		for i := 0; i < defaultWorkers*2; i++ {
 			source <- i
 		}
-		done.Set(true)
+		atomic.AddInt32(&done, 1)
 	}, func(item interface{}, writer Writer, c func(error)) {
 		i := item.(int)
 		if i == defaultWorkers/2 {