github.com/v2pro/plz@v0.0.0-20221028024117-e5f9aec5b631/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  func ResetGls(goid int64, initialValue map[interface{}]interface{}) {
    26  	shardIndex := goid % shardsCount
    27  	lock := globalLocks[shardIndex]
    28  	lock.Lock()
    29  	globalMaps[shardIndex][goid] = initialValue
    30  	lock.Unlock()
    31  }
    32  
    33  func DeleteGls(goid int64) {
    34  	shardIndex := goid % shardsCount
    35  	lock := globalLocks[shardIndex]
    36  	lock.Lock()
    37  	delete(globalMaps[shardIndex], goid)
    38  	lock.Unlock()
    39  }
    40  
    41  func GetGls(goid int64) map[interface{}]interface{} {
    42  	shardIndex := goid % shardsCount
    43  	lock := globalLocks[shardIndex]
    44  	lock.RLock()
    45  	gls, found := globalMaps[shardIndex][goid]
    46  	lock.RUnlock()
    47  	if found {
    48  		return gls
    49  	} else {
    50  		return nil
    51  	}
    52  }
    53  
    54  func WithGls(f func()) func() {
    55  	parentGls := GetGls(GoID())
    56  	// parentGls can not be used in other goroutine, otherwise not thread safe
    57  	// make a deep for child goroutine
    58  	childGls := map[interface{}]interface{}{}
    59  	for k, v := range parentGls {
    60  		asCopiable, ok := v.(copiable)
    61  		if ok {
    62  			childGls[k] = asCopiable.Copy()
    63  		} else {
    64  			childGls[k] = v
    65  		}
    66  	}
    67  	return func() {
    68  		goid := GoID()
    69  		ResetGls(goid, childGls)
    70  		defer DeleteGls(goid)
    71  		f()
    72  	}
    73  }
    74  
    75  func WithEmptyGls(f func()) func() {
    76  	// do not inherit from parent gls
    77  	return func() {
    78  		goid := GoID()
    79  		ResetGls(goid, make(map[interface{}]interface{}))
    80  		defer DeleteGls(goid)
    81  		f()
    82  	}
    83  }
    84  
    85  func Get(key interface{}) interface{} {
    86  	glsMap := GetGls(GoID())
    87  	if glsMap == nil {
    88  		return nil
    89  	}
    90  	return glsMap[key]
    91  }
    92  
    93  func Set(key interface{}, value interface{}) {
    94  	glsMap := GetGls(GoID())
    95  	if glsMap == nil {
    96  		panic("gls not enabled for this goroutine")
    97  	}
    98  	glsMap[key] = value
    99  }
   100  
   101  func IsGlsEnabled(goid int64) bool {
   102  	return GetGls(goid) != nil
   103  }
   104