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

     1  package goid
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"time"
     7  	"unsafe"
     8  )
     9  
    10  var (
    11  	storages          atomic.Value       // The global storage map (map[int64]*store)
    12  	storageLock       sync.Mutex         // The Lock to control accessing of storages
    13  	storageGCTimer    *time.Timer        // The timer of storage's garbage collector
    14  	storageGCInterval = time.Second * 30 // The pre-defined gc interval
    15  )
    16  
    17  func init() {
    18  	storages.Store(map[int64]*store{})
    19  }
    20  
    21  func gcRunning() bool {
    22  	storageLock.Lock()
    23  	defer storageLock.Unlock()
    24  	return storageGCTimer != nil
    25  }
    26  
    27  type store struct {
    28  	gid    int64
    29  	count  uint32
    30  	values map[uintptr]any
    31  }
    32  
    33  type storage struct {
    34  	ignore string
    35  }
    36  
    37  func (t *storage) Get() (v any) {
    38  	s := loadCurrentStore()
    39  	id := uintptr(unsafe.Pointer(t))
    40  	return s.values[id]
    41  }
    42  
    43  func (t *storage) Set(v any) (oldValue any) {
    44  	s := loadCurrentStore()
    45  	id := uintptr(unsafe.Pointer(t))
    46  	oldValue = s.values[id]
    47  	s.values[id] = v
    48  	atomic.StoreUint32(&s.count, uint32(len(s.values)))
    49  
    50  	// try restart gc timer if Set for the first time
    51  	if oldValue == nil {
    52  		storageLock.Lock()
    53  		if storageGCTimer == nil {
    54  			storageGCTimer = time.AfterFunc(storageGCInterval, clearDeadStore)
    55  		}
    56  		storageLock.Unlock()
    57  	}
    58  	return
    59  }
    60  
    61  func (t *storage) Del() (v any) {
    62  	s := loadCurrentStore()
    63  	id := uintptr(unsafe.Pointer(t))
    64  	v = s.values[id]
    65  	delete(s.values, id)
    66  	atomic.StoreUint32(&s.count, uint32(len(s.values)))
    67  	return
    68  }
    69  
    70  func (t *storage) Clear() {
    71  	s := loadCurrentStore()
    72  	s.values = make(map[uintptr]any)
    73  	atomic.StoreUint32(&s.count, 0)
    74  }
    75  
    76  // loadCurrentStore load the store of current goroutine.
    77  func loadCurrentStore() (s *store) {
    78  	gid := Goid()
    79  	storeMap := storages.Load().(map[int64]*store)
    80  	if s = storeMap[gid]; s == nil {
    81  		storageLock.Lock()
    82  		oldStoreMap := storages.Load().(map[int64]*store)
    83  		if s = oldStoreMap[gid]; s == nil {
    84  			s = &store{
    85  				gid:    gid,
    86  				values: make(map[uintptr]any),
    87  			}
    88  			newStoreMap := make(map[int64]*store, len(oldStoreMap)+1)
    89  			for k, v := range oldStoreMap {
    90  				newStoreMap[k] = v
    91  			}
    92  			newStoreMap[gid] = s
    93  			storages.Store(newStoreMap)
    94  		}
    95  		storageLock.Unlock()
    96  	}
    97  	return
    98  }
    99  
   100  // clearDeadStore clear all data of dead goroutine.
   101  func clearDeadStore() {
   102  	storageLock.Lock()
   103  	defer storageLock.Unlock()
   104  
   105  	// load all alive goids
   106  	gids := AllGoids()
   107  	gidMap := make(map[int64]struct{}, len(gids))
   108  	for _, gid := range gids {
   109  		gidMap[gid] = struct{}{}
   110  	}
   111  
   112  	// scan global storeMap check the dead and live store count.
   113  	var storeMap = storages.Load().(map[int64]*store)
   114  	var deadCnt, liveCnt int
   115  	for id, s := range storeMap {
   116  		if _, ok := gidMap[id]; ok {
   117  			if atomic.LoadUint32(&s.count) > 0 {
   118  				liveCnt++
   119  			}
   120  			liveCnt++
   121  		} else {
   122  			deadCnt++
   123  		}
   124  	}
   125  
   126  	// clean dead store of dead goroutine if need.
   127  	if deadCnt > 0 {
   128  		newStoreMap := make(map[int64]*store, len(storeMap)-deadCnt)
   129  		for id, s := range storeMap {
   130  			if _, ok := gidMap[id]; ok {
   131  				newStoreMap[id] = s
   132  			}
   133  		}
   134  		storages.Store(newStoreMap)
   135  	}
   136  
   137  	if liveCnt > 0 {
   138  		storageGCTimer.Reset(storageGCInterval)
   139  	} else {
   140  		storageGCTimer = nil
   141  	}
   142  }