gitee.com/quant1x/gox@v1.21.2/threadlocal/gid.go (about)

     1  package threadlocal
     2  
     3  import (
     4  	"fmt"
     5  	"runtime"
     6  	"sync"
     7  )
     8  
     9  // Getgid returns the ID of the current Go routine
    10  //
    11  // The solution here is not the fastest in the world but it gets the job done
    12  // in a portable way. For a faster ways to do this that involves assembler code
    13  // for different platforms, please look at: https://github.com/huandu/go-tls
    14  func getg() int64 {
    15  	const prefixLen = 10 // Length of prefix "goroutine "
    16  	var buf [64]byte
    17  
    18  	l := runtime.Stack(buf[:64], false)
    19  	n := int64(0)
    20  	for i := prefixLen; i < l; i++ {
    21  		d := buf[i]
    22  		if d < 0x30 || d > 0x39 {
    23  			break
    24  		}
    25  		n = n*10 + int64(d-0x30)
    26  	}
    27  	if n == 0 {
    28  		panic(fmt.Errorf(`unable to retrieve id of current go routine`))
    29  	}
    30  	return n
    31  }
    32  
    33  var tlsLock sync.RWMutex
    34  var tls = make(map[int64]map[string]interface{}, 7)
    35  
    36  // Init initializes a go routine local storage for the current go routine
    37  func Init() {
    38  	gid := getg()
    39  	ls := make(map[string]interface{})
    40  	tlsLock.Lock()
    41  	tls[gid] = ls
    42  	tlsLock.Unlock()
    43  }
    44  
    45  // Cleanup deletes the local storage for the current go routine
    46  func Cleanup() {
    47  	gid := getg()
    48  	tlsLock.Lock()
    49  	delete(tls, gid)
    50  	tlsLock.Unlock()
    51  }
    52  
    53  // Get returns a variable from the local storage of the current go routine
    54  func Get(key string) (interface{}, bool) {
    55  	gid := getg()
    56  	tlsLock.RLock()
    57  	ls, ok := tls[gid]
    58  	tlsLock.RUnlock()
    59  	var found interface{}
    60  	if ok {
    61  		found, ok = ls[key]
    62  	}
    63  	return found, ok
    64  }
    65  
    66  // Go executes the given function in a go routine and ensures that the local
    67  // storage is initialized before the function is called and deleted before
    68  // after the function returns or panics
    69  func Go(f func()) {
    70  	go func() {
    71  		defer Cleanup()
    72  		Init()
    73  		f()
    74  	}()
    75  }
    76  
    77  // Delete deletes a variable from the local storage of the current go routine
    78  func Delete(key string) {
    79  	gid := getg()
    80  	tlsLock.RLock()
    81  	ls, ok := tls[gid]
    82  	tlsLock.RUnlock()
    83  	if ok {
    84  		delete(ls, key)
    85  	}
    86  }
    87  
    88  // Set adds or replaces a variable to the local storage of the current go routine
    89  func Set(key string, value interface{}) {
    90  	gid := getg()
    91  	tlsLock.RLock()
    92  	ls, ok := tls[gid]
    93  	tlsLock.RUnlock()
    94  	if !ok {
    95  		panic(`thread local not initialized for current go routine`)
    96  	}
    97  	ls[key] = value
    98  }