github.com/cyub/threadlocal@v0.0.0-20220816024249-5db4997a97f4/store.go (about)

     1  package threadlocal
     2  
     3  const (
     4  	INITIALIZE_THREADLOCALMAP_SIZE   = 16
     5  	INITIALIZE_THREADLOCALSTORE_SIZE = 128
     6  )
     7  
     8  var (
     9  	store         = make(threadLocalStore, INITIALIZE_THREADLOCALSTORE_SIZE)
    10  	expungeEntity = NewEntity(nil, nil)
    11  )
    12  
    13  type threadLocalStore map[uint32]*ThreadlocalMap
    14  
    15  func NewThreadlocalMap(capacity int) *ThreadlocalMap {
    16  	return &ThreadlocalMap{
    17  		size:     0,
    18  		capacity: capacity,
    19  		entities: make([]*Entity, capacity),
    20  	}
    21  }
    22  
    23  type ThreadlocalMap struct {
    24  	size     int
    25  	capacity int
    26  	entities []*Entity
    27  }
    28  
    29  func (tlm *ThreadlocalMap) Size() int {
    30  	return tlm.size
    31  }
    32  
    33  func (tlm *ThreadlocalMap) Set(key *Threadlocal, val interface{}) {
    34  	i := key.HashCode & (tlm.capacity - 1)
    35  	for e := tlm.entities[i]; e != nil && e != expungeEntity; i = tlm.nextIndex(i, tlm.capacity) {
    36  		if e.key == key {
    37  			e.val = val
    38  			return
    39  		}
    40  		e = tlm.entities[i]
    41  	}
    42  	tlm.entities[i] = NewEntity(key, val)
    43  	tlm.size++
    44  	tlm.rehash()
    45  }
    46  
    47  func (tlm *ThreadlocalMap) Get(key *Threadlocal) interface{} {
    48  	i := key.HashCode & (tlm.capacity - 1)
    49  	for e := tlm.entities[i]; e != nil && e != expungeEntity; i = tlm.nextIndex(i, tlm.capacity) {
    50  		if e.key == key {
    51  			return e.val
    52  		}
    53  		e = tlm.entities[i]
    54  	}
    55  	return nil
    56  }
    57  
    58  func (tlm *ThreadlocalMap) Remove(key *Threadlocal) {
    59  	i := key.HashCode & (tlm.capacity - 1)
    60  	for e := tlm.entities[i]; e != nil && e != expungeEntity; i = tlm.nextIndex(i, tlm.capacity) {
    61  		if e.key == key {
    62  			tlm.entities[i] = expungeEntity // mark deleted
    63  			tlm.size--
    64  			break
    65  		}
    66  		e = tlm.entities[i]
    67  	}
    68  }
    69  
    70  func (tlm *ThreadlocalMap) nextIndex(i, cap int) int {
    71  	if i+1 < cap {
    72  		return i + 1
    73  	}
    74  	return 0
    75  }
    76  
    77  func (tlm *ThreadlocalMap) rehash() {
    78  	if tlm.size < tlm.capacity*3/4 {
    79  		return
    80  	}
    81  	newLen := tlm.capacity << 1
    82  	newTab := make([]*Entity, newLen)
    83  	for _, e := range tlm.entities {
    84  		if e == nil || e == expungeEntity {
    85  			continue
    86  		}
    87  		i := e.key.HashCode & (newLen - 1)
    88  		for newTab[i] != nil {
    89  			i = tlm.nextIndex(i, newLen)
    90  		}
    91  		newTab[i] = e
    92  	}
    93  	tlm.capacity = newLen
    94  	tlm.entities = newTab
    95  }
    96  
    97  type Entity struct {
    98  	key *Threadlocal
    99  	val interface{}
   100  }
   101  
   102  func NewEntity(tl *Threadlocal, val interface{}) *Entity {
   103  	return &Entity{
   104  		key: tl,
   105  		val: val,
   106  	}
   107  }
   108  
   109  func currentThreadLocalMap() *ThreadlocalMap {
   110  	var tid = ThreadId()
   111  	if store[tid] == nil {
   112  		store[tid] = NewThreadlocalMap(INITIALIZE_THREADLOCALMAP_SIZE)
   113  	}
   114  	return store[tid]
   115  }