74 lines
1.3 KiB
Go
74 lines
1.3 KiB
Go
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()
|
|
}()
|
|
}
|