github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/mapp/safemap.go (about)

     1  package mapp
     2  
     3  import "sync"
     4  
     5  const (
     6  	copyThreshold = 1000
     7  	maxDeletion   = 10000
     8  )
     9  
    10  // SafeMap provides a map alternative to avoid memory leak.
    11  // This implementation is not needed until issue below fixed.
    12  // https://github.com/golang/go/issues/20135
    13  type SafeMap struct {
    14  	lock        sync.RWMutex
    15  	deletionOld int
    16  	deletionNew int
    17  	dirtyOld    map[interface{}]interface{}
    18  	dirtyNew    map[interface{}]interface{}
    19  }
    20  
    21  // NewSafeMap returns a SafeMap.
    22  func NewSafeMap() *SafeMap {
    23  	return &SafeMap{
    24  		dirtyOld: make(map[interface{}]interface{}),
    25  		dirtyNew: make(map[interface{}]interface{}),
    26  	}
    27  }
    28  
    29  // Del deletes the value with the given key from m.
    30  func (m *SafeMap) Del(key interface{}) {
    31  	m.lock.Lock()
    32  	if _, ok := m.dirtyOld[key]; ok {
    33  		delete(m.dirtyOld, key)
    34  		m.deletionOld++
    35  	} else if _, ok := m.dirtyNew[key]; ok {
    36  		delete(m.dirtyNew, key)
    37  		m.deletionNew++
    38  	}
    39  	if m.deletionOld >= maxDeletion && len(m.dirtyOld) < copyThreshold {
    40  		for k, v := range m.dirtyOld {
    41  			m.dirtyNew[k] = v
    42  		}
    43  		m.dirtyOld = m.dirtyNew
    44  		m.deletionOld = m.deletionNew
    45  		m.dirtyNew = make(map[interface{}]interface{})
    46  		m.deletionNew = 0
    47  	}
    48  	if m.deletionNew >= maxDeletion && len(m.dirtyNew) < copyThreshold {
    49  		for k, v := range m.dirtyNew {
    50  			m.dirtyOld[k] = v
    51  		}
    52  		m.dirtyNew = make(map[interface{}]interface{})
    53  		m.deletionNew = 0
    54  	}
    55  	m.lock.Unlock()
    56  }
    57  
    58  // Get gets the value with the given key from m.
    59  func (m *SafeMap) Get(key interface{}) (interface{}, bool) {
    60  	m.lock.RLock()
    61  	defer m.lock.RUnlock()
    62  
    63  	if val, ok := m.dirtyOld[key]; ok {
    64  		return val, true
    65  	}
    66  
    67  	val, ok := m.dirtyNew[key]
    68  	return val, ok
    69  }
    70  
    71  // Range calls f sequentially for each key and value present in the map.
    72  // If f returns false, range stops the iteration.
    73  func (m *SafeMap) Range(f func(key, val interface{}) bool) {
    74  	m.lock.RLock()
    75  	defer m.lock.RUnlock()
    76  
    77  	for k, v := range m.dirtyOld {
    78  		if !f(k, v) {
    79  			return
    80  		}
    81  	}
    82  	for k, v := range m.dirtyNew {
    83  		if !f(k, v) {
    84  			return
    85  		}
    86  	}
    87  }
    88  
    89  // Set sets the value into m with the given key.
    90  func (m *SafeMap) Set(key, value interface{}) {
    91  	m.lock.Lock()
    92  	if m.deletionOld <= maxDeletion {
    93  		if _, ok := m.dirtyNew[key]; ok {
    94  			delete(m.dirtyNew, key)
    95  			m.deletionNew++
    96  		}
    97  		m.dirtyOld[key] = value
    98  	} else {
    99  		if _, ok := m.dirtyOld[key]; ok {
   100  			delete(m.dirtyOld, key)
   101  			m.deletionOld++
   102  		}
   103  		m.dirtyNew[key] = value
   104  	}
   105  	m.lock.Unlock()
   106  }
   107  
   108  // Size returns the size of m.
   109  func (m *SafeMap) Size() int {
   110  	m.lock.RLock()
   111  	size := len(m.dirtyOld) + len(m.dirtyNew)
   112  	m.lock.RUnlock()
   113  	return size
   114  }