github.com/isyscore/isc-gobase@v1.5.3-0.20231218061332-cbc7451899e9/goid/rl/routinelocal.go (about)

     1  package rl
     2  
     3  import (
     4  	"context"
     5  	"runtime/pprof"
     6  	"strconv"
     7  	"sync/atomic"
     8  	"unsafe"
     9  )
    10  
    11  const labelTag = "github.com/isyscore/isc-gobase"
    12  
    13  var lastID uintptr
    14  
    15  type labelMapAndContext struct {
    16  	lm  labelMap
    17  	ctx context.Context
    18  }
    19  
    20  func Set(ctx context.Context) {
    21  	var gctx labelMapAndContext
    22  	pprof.SetGoroutineLabels(ctx)
    23  	lm := (*labelMap)(runtimeGetProfLabel())
    24  	if lm != nil && len(*lm) > 0 {
    25  		gctx.lm = *lm
    26  	} else {
    27  		gctx.lm = make(labelMap)
    28  	}
    29  	id := atomic.AddUintptr(&lastID, 1)
    30  	gctx.lm[labelTag] = strconv.FormatUint(uint64(id), 16)
    31  	gctx.ctx = ctx
    32  	runtimeSetProfLabel(unsafe.Pointer(&gctx))
    33  }
    34  
    35  func Get() context.Context {
    36  	lm := (*labelMap)(runtimeGetProfLabel())
    37  	if lm != nil {
    38  		if _, ok := (*lm)[labelTag]; ok {
    39  			gctx := (*labelMapAndContext)(unsafe.Pointer(lm))
    40  			return gctx.ctx
    41  		}
    42  	}
    43  	return nil
    44  }