Add RunSafe with context (#3224)

Co-authored-by: yangtao <mrynag8614@163.com>
Co-authored-by: Kevin Wan <wanjunfeng@gmail.com>

core/threading/routines.go

@@ -2,6 +2,7 @@ package threading
 
 import (
 	"bytes"
+	"context"
 	"runtime"
 	"strconv"
 
@@ -13,6 +14,11 @@ func GoSafe(fn func()) {
 	go RunSafe(fn)
 }
 
+// GoSafeCtx runs the given fn using another goroutine, recovers if fn panics with ctx.
+func GoSafeCtx(ctx context.Context, fn func()) {
+	go RunSafeCtx(ctx, fn)
+}
+
 // RoutineId is only for debug, never use it in production.
 func RoutineId() uint64 {
 	b := make([]byte, 64)
@@ -31,3 +37,10 @@ func RunSafe(fn func()) {
 
 	fn()
 }
+
+// RunSafeCtx runs the given fn, recovers if fn panics with ctx.
+func RunSafeCtx(ctx context.Context, fn func()) {
+	defer rescue.RecoverCtx(ctx)
+
+	fn()
+}

core/threading/routines_test.go

@@ -1,12 +1,15 @@
 package threading
 
 import (
+	"bytes"
+	"context"
 	"io"
 	"log"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/logx"
 )
 
 func TestRoutineId(t *testing.T) {
@@ -34,3 +37,51 @@ func TestRunSafe(t *testing.T) {
 	<-ch
 	i++
 }
+
+func TestRunSafeCtx(t *testing.T) {
+	var buf bytes.Buffer
+	logx.SetWriter(logx.NewWriter(&buf))
+	ctx := context.Background()
+	ch := make(chan lang.PlaceholderType)
+
+	i := 0
+
+	defer func() {
+		assert.Equal(t, 1, i)
+	}()
+
+	go RunSafeCtx(ctx, func() {
+		defer func() {
+			ch <- lang.Placeholder
+		}()
+
+		panic("panic")
+	})
+
+	<-ch
+	i++
+}
+
+func TestGoSafeCtx(t *testing.T) {
+	var buf bytes.Buffer
+	logx.SetWriter(logx.NewWriter(&buf))
+	ctx := context.Background()
+	ch := make(chan lang.PlaceholderType)
+
+	i := 0
+
+	defer func() {
+		assert.Equal(t, 1, i)
+	}()
+
+	GoSafeCtx(ctx, func() {
+		defer func() {
+			ch <- lang.Placeholder
+		}()
+
+		panic("panic")
+	})
+
+	<-ch
+	i++
+}