init

pkg/threading/routinegroup.go

@@ -0,0 +1,44 @@
+package threading
+
+import (
+	"sync"
+)
+
+// A RoutineGroup is used to group goroutines together and all wait all goroutines to be done.
+type RoutineGroup struct {
+	waitGroup sync.WaitGroup
+}
+
+// NewRoutineGroup returns a RoutineGroup.
+func NewRoutineGroup() *RoutineGroup {
+	return new(RoutineGroup)
+}
+
+// Run runs the given fn in RoutineGroup.
+// Don't reference the variables from outside,
+// because outside variables can be changed by other goroutines
+func (g *RoutineGroup) Run(fn func()) {
+	g.waitGroup.Add(1)
+
+	go func() {
+		defer g.waitGroup.Done()
+		fn()
+	}()
+}
+
+// RunSafe runs the given fn in RoutineGroup, and avoid panics.
+// Don't reference the variables from outside,
+// because outside variables can be changed by other goroutines
+func (g *RoutineGroup) RunSafe(fn func()) {
+	g.waitGroup.Add(1)
+
+	GoSafe(func() {
+		defer g.waitGroup.Done()
+		fn()
+	})
+}
+
+// Wait waits all running functions to be done.
+func (g *RoutineGroup) Wait() {
+	g.waitGroup.Wait()
+}

pkg/threading/routinegroup_test.go

@@ -0,0 +1,45 @@
+package threading
+
+import (
+	"io"
+	"log"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRoutineGroupRun(t *testing.T) {
+	var count int32
+	group := NewRoutineGroup()
+	for i := 0; i < 3; i++ {
+		group.Run(func() {
+			atomic.AddInt32(&count, 1)
+		})
+	}
+
+	group.Wait()
+
+	assert.Equal(t, int32(3), count)
+}
+
+func TestRoutingGroupRunSafe(t *testing.T) {
+	log.SetOutput(io.Discard)
+
+	var count int32
+	group := NewRoutineGroup()
+	var once sync.Once
+	for i := 0; i < 3; i++ {
+		group.RunSafe(func() {
+			once.Do(func() {
+				panic("")
+			})
+			atomic.AddInt32(&count, 1)
+		})
+	}
+
+	group.Wait()
+
+	assert.Equal(t, int32(2), count)
+}

pkg/threading/routines.go

@@ -0,0 +1,46 @@
+package threading
+
+import (
+	"bytes"
+	"context"
+	"runtime"
+	"strconv"
+
+	"github.com/perfect-panel/ppanel-server/pkg/rescue"
+)
+
+// GoSafe runs the given fn using another goroutine, recovers if fn panics.
+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)
+	b = b[:runtime.Stack(b, false)]
+	b = bytes.TrimPrefix(b, []byte("goroutine "))
+	b = b[:bytes.IndexByte(b, ' ')]
+	// if error, just return 0
+	n, _ := strconv.ParseUint(string(b), 10, 64)
+
+	return n
+}
+
+// RunSafe runs the given fn, recovers if fn panics.
+func RunSafe(fn func()) {
+	defer rescue.Recover()
+
+	fn()
+}
+
+// RunSafeCtx runs the given fn, recovers if fn panics with ctx.
+func RunSafeCtx(ctx context.Context, fn func()) {
+	defer rescue.RecoverCtx(ctx)
+
+	fn()
+}

pkg/threading/routines_test.go

@@ -0,0 +1,11 @@
+package threading
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRoutineId(t *testing.T) {
+	assert.True(t, RoutineId() > 0)
+}