gitee.com/quant1x/gox@v1.7.6/gls/gls.go (about)

     1  package gls
     2  
     3  import "sync"
     4  
     5  type globalMapType map[int64]map[interface{}]interface{}
     6  
     7  const shardsCount = 16
     8  
     9  var globalLocks []*sync.RWMutex
    10  var globalMaps []globalMapType
    11  
    12  type copiable interface {
    13  	Copy() interface{}
    14  }
    15  
    16  func init() {
    17  	globalMaps = make([]globalMapType, shardsCount)
    18  	globalLocks = make([]*sync.RWMutex, shardsCount)
    19  	for i := 0; i < shardsCount; i++ {
    20  		globalMaps[i] = make(globalMapType)
    21  		globalLocks[i] = &sync.RWMutex{}
    22  	}
    23  }
    24  
    25  // ResetGls reset the goroutine local storage for specified goroutine
    26  func ResetGls(goid int64, initialValue map[interface{}]interface{}) {
    27  	shardIndex := goid % shardsCount
    28  	lock := globalLocks[shardIndex]
    29  	lock.Lock()
    30  	globalMaps[shardIndex][goid] = initialValue
    31  	lock.Unlock()
    32  }
    33  
    34  // DeleteGls remove goroutine local storage for specified goroutine
    35  func DeleteGls(goid int64) {
    36  	shardIndex := goid % shardsCount
    37  	lock := globalLocks[shardIndex]
    38  	lock.Lock()
    39  	delete(globalMaps[shardIndex], goid)
    40  	lock.Unlock()
    41  }
    42  
    43  // GetGls get goroutine local storage for specified goroutine
    44  // if the goroutine did not set gls, it will return nil
    45  func GetGls(goid int64) map[interface{}]interface{} {
    46  	shardIndex := goid % shardsCount
    47  	lock := globalLocks[shardIndex]
    48  	lock.RLock()
    49  	gls, found := globalMaps[shardIndex][goid]
    50  	lock.RUnlock()
    51  	if found {
    52  		return gls
    53  	} else {
    54  		return nil
    55  	}
    56  }
    57  
    58  // WithGls setup and teardown the gls in the wrapper.
    59  // go WithGls(func(){}) will add gls for the new goroutine.
    60  // The gls will be removed once goroutine exit
    61  func WithGls(f func()) func() {
    62  	parentGls := GetGls(GoID())
    63  	// parentGls can not be used in other goroutine, otherwise not thread safe
    64  	// make a deep for child goroutine
    65  	childGls := map[interface{}]interface{}{}
    66  	for k, v := range parentGls {
    67  		asCopiable, ok := v.(copiable)
    68  		if ok {
    69  			childGls[k] = asCopiable.Copy()
    70  		} else {
    71  			childGls[k] = v
    72  		}
    73  	}
    74  	return func() {
    75  		goid := GoID()
    76  		ResetGls(goid, childGls)
    77  		defer DeleteGls(goid)
    78  		f()
    79  	}
    80  }
    81  
    82  // WithEmptyGls works like WithGls, but do not inherit gls from parent goroutine.
    83  func WithEmptyGls(f func()) func() {
    84  	// do not inherit from parent gls
    85  	return func() {
    86  		goid := GoID()
    87  		ResetGls(goid, make(map[interface{}]interface{}))
    88  		defer DeleteGls(goid)
    89  		f()
    90  	}
    91  }
    92  
    93  // Get key from goroutine local storage
    94  func Get(key interface{}) interface{} {
    95  	glsMap := GetGls(GoID())
    96  	if glsMap == nil {
    97  		return nil
    98  	}
    99  	return glsMap[key]
   100  }
   101  
   102  // Set key and element to goroutine local storage
   103  func Set(key interface{}, value interface{}) {
   104  	goid := GoID()
   105  	if !IsGlsEnabled(goid) {
   106  		ResetGls(goid, make(map[interface{}]interface{}))
   107  	}
   108  	glsMap := GetGls(goid)
   109  	if glsMap == nil {
   110  		panic("gls not enabled for this goroutine")
   111  	}
   112  	glsMap[key] = value
   113  }
   114  
   115  func Remove(key interface{}) {
   116  	goid := GoID()
   117  	glsMap := GetGls(goid)
   118  	if glsMap != nil {
   119  		// 删除key
   120  		delete(glsMap, key)
   121  	}
   122  	// 如果map为空, 则删除
   123  	if len(glsMap) == 0 {
   124  		DeleteGls(goid)
   125  	}
   126  }
   127  
   128  // IsGlsEnabled test if the gls is available for specified goroutine
   129  func IsGlsEnabled(goid int64) bool {
   130  	return GetGls(goid) != nil
   131  }