initial import

core/collection/safemap.go

@@ -0,0 +1,91 @@
+package collection
+
+import "sync"
+
+const (
+	copyThreshold = 1000
+	maxDeletion   = 10000
+)
+
+// SafeMap provides a map alternative to avoid memory leak.
+// This implementation is not needed until issue below fixed.
+// https://github.com/golang/go/issues/20135
+type SafeMap struct {
+	lock        sync.RWMutex
+	deletionOld int
+	deletionNew int
+	dirtyOld    map[interface{}]interface{}
+	dirtyNew    map[interface{}]interface{}
+}
+
+func NewSafeMap() *SafeMap {
+	return &SafeMap{
+		dirtyOld: make(map[interface{}]interface{}),
+		dirtyNew: make(map[interface{}]interface{}),
+	}
+}
+
+func (m *SafeMap) Del(key interface{}) {
+	m.lock.Lock()
+	if _, ok := m.dirtyOld[key]; ok {
+		delete(m.dirtyOld, key)
+		m.deletionOld++
+	} else if _, ok := m.dirtyNew[key]; ok {
+		delete(m.dirtyNew, key)
+		m.deletionNew++
+	}
+	if m.deletionOld >= maxDeletion && len(m.dirtyOld) < copyThreshold {
+		for k, v := range m.dirtyOld {
+			m.dirtyNew[k] = v
+		}
+		m.dirtyOld = m.dirtyNew
+		m.deletionOld = m.deletionNew
+		m.dirtyNew = make(map[interface{}]interface{})
+		m.deletionNew = 0
+	}
+	if m.deletionNew >= maxDeletion && len(m.dirtyNew) < copyThreshold {
+		for k, v := range m.dirtyNew {
+			m.dirtyOld[k] = v
+		}
+		m.dirtyNew = make(map[interface{}]interface{})
+		m.deletionNew = 0
+	}
+	m.lock.Unlock()
+}
+
+func (m *SafeMap) Get(key interface{}) (interface{}, bool) {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+
+	if val, ok := m.dirtyOld[key]; ok {
+		return val, true
+	} else {
+		val, ok := m.dirtyNew[key]
+		return val, ok
+	}
+}
+
+func (m *SafeMap) Set(key, value interface{}) {
+	m.lock.Lock()
+	if m.deletionOld <= maxDeletion {
+		if _, ok := m.dirtyNew[key]; ok {
+			delete(m.dirtyNew, key)
+			m.deletionNew++
+		}
+		m.dirtyOld[key] = value
+	} else {
+		if _, ok := m.dirtyOld[key]; ok {
+			delete(m.dirtyOld, key)
+			m.deletionOld++
+		}
+		m.dirtyNew[key] = value
+	}
+	m.lock.Unlock()
+}
+
+func (m *SafeMap) Size() int {
+	m.lock.RLock()
+	size := len(m.dirtyOld) + len(m.dirtyNew)
+	m.lock.RUnlock()
+	return size
+}