u

code/base-project/base-go-v3/utils/routinex/routinex.go

@@ -0,0 +1,73 @@
+package routinex
+
+import (
+	"runtime"
+	"strconv"
+	"strings"
+	"sync"
+)
+
+// CurGID 获取 routine 的id
+func CurGID() uint64 {
+	var buf [64]byte
+	n := runtime.Stack(buf[:], false)
+	line := strings.Fields(strings.TrimPrefix(string(buf[:n]), "goroutine "))[0]
+	gid, _ := strconv.ParseUint(line, 10, 64)
+	return gid
+}
+
+type RoutineLocal struct {
+	data sync.Map // gid -> map[key]any
+}
+
+var (
+	instance     *RoutineLocal
+	instanceOnce sync.Once
+)
+
+// GetInstance 获取当前示例
+func GetInstance() *RoutineLocal {
+	instanceOnce.Do(func() {
+		instance = &RoutineLocal{}
+	})
+	return instance
+}
+
+// Set 存储单条数据
+func Set(key string, value any) {
+	gid := CurGID()
+	v, _ := GetInstance().data.LoadOrStore(gid, &sync.Map{})
+	m := v.(*sync.Map)
+	m.Store(key, value)
+}
+
+// Get 获取单条数据
+func Get(key string) any {
+	gid := CurGID()
+	if v, ok := GetInstance().data.Load(gid); ok {
+		m := v.(*sync.Map)
+		if val, ok := m.Load(key); ok {
+			return val
+		}
+	}
+	return nil
+}
+
+// Clear 删除当前 goroutine 所有数据
+func Clear() {
+	GetInstance().data.Delete(CurGID())
+}
+
+// Go 自动继承父 goroutine 数据，多 routine 的时候需要使用
+func Go(f func()) {
+	parentGID := CurGID()
+	parentData, _ := GetInstance().data.Load(parentGID)
+
+	go func() {
+		if parentData != nil {
+			GetInstance().data.Store(CurGID(), parentData)
+			defer Clear()
+		}
+		f()
+	}()
+}