Add Clear() function to sync.Map (#142)

Signed-off-by: kpango <[email protected]>

gache.go

@@ -376,15 +376,14 @@ func (g *gache[V]) Stop() {
 // Clear deletes all key and value present in the Gache.
 func (g *gache[V]) Clear() {
 	for i := range g.shards {
-		g.shards[i] = newMap[V]()
+		if g.shards[i] == nil {
+			g.shards[i] = newMap[V]()
+		} else {
+			g.shards[i].Clear()
+		}
 	}
 }
 
-func (v *value[V]) Size() uintptr {
-	var size uintptr
-
-	size += unsafe.Sizeof(v.expire) // int64
-	size += unsafe.Sizeof(v.val)    // V size
-
-	return size
+func (v *value[V]) Size() (size uintptr) {
+	return unsafe.Sizeof(v.expire) + unsafe.Sizeof(v.val)
 }

map.go

@@ -97,6 +97,26 @@ func (m *Map[K, V]) Store(key K, value V) {
 	_, _ = m.Swap(key, value)
 }
 
+func (m *Map[K, V]) Clear() {
+	read := m.loadReadOnly()
+	if len(read.m) == 0 && !read.amended {
+		// Avoid allocating a new readOnly when the map is already clear.
+		return
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	read = m.loadReadOnly()
+	if len(read.m) > 0 || read.amended {
+		m.read.Store(&readOnly[K, V]{})
+	}
+
+	clear(m.dirty)
+	// Don't immediately promote the newly-cleared dirty map on the next operation.
+	m.misses = 0
+}
+
 func (e *entry[V]) tryCompareAndSwap(old, new V) (ok bool) {
 	p := e.p.Load()
 	if p == nil || p == e.expunged.Load() || *p != old {
@@ -136,6 +156,9 @@ func (m *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
 	read = m.loadReadOnly()
 	if e, ok := read.m[key]; ok {
 		if e.unexpungeLocked() {
+			if m.dirty == nil {
+				m.initDirty(len(read.m))
+			}
 			m.dirty[key] = e
 		}
 		actual, loaded, _ = e.tryLoadOrStore(value)
@@ -147,6 +170,9 @@ func (m *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
 			m.dirtyLocked()
 			m.read.Store(&readOnly[K, V]{m: read.m, amended: true})
 		}
+		if m.dirty == nil {
+			m.initDirty(len(read.m))
+		}
 		m.dirty[key] = newEntry(value)
 		actual, loaded = value, false
 	}
@@ -241,6 +267,9 @@ func (m *Map[K, V]) Swap(key K, value V) (previous V, loaded bool) {
 	read = m.loadReadOnly()
 	if e, ok := read.m[key]; ok {
 		if e.unexpungeLocked() {
+			if m.dirty == nil {
+				m.initDirty(len(read.m))
+			}
 			m.dirty[key] = e
 		}
 		if v := e.swapLocked(&value); v != nil {
@@ -257,6 +286,9 @@ func (m *Map[K, V]) Swap(key K, value V) (previous V, loaded bool) {
 			m.dirtyLocked()
 			m.read.Store(&readOnly[K, V]{m: read.m, amended: true})
 		}
+		if m.dirty == nil {
+			m.initDirty(len(read.m))
+		}
 		m.dirty[key] = newEntry(value)
 	}
 	m.mu.Unlock()
@@ -316,7 +348,8 @@ func (m *Map[K, V]) Range(f func(key K, value V) bool) {
 		read = m.loadReadOnly()
 		if read.amended {
 			read = readOnly[K, V]{m: m.dirty}
-			m.read.Store(&read)
+			copyRead := read
+			m.read.Store(&copyRead)
 			m.dirty = nil
 			m.misses = 0
 		}
@@ -334,23 +367,6 @@ func (m *Map[K, V]) Range(f func(key K, value V) bool) {
 	}
 }
 
-func (m *Map[K, V]) Len() int {
-	read := m.loadReadOnly()
-	if read.amended {
-		m.mu.Lock()
-		read = m.loadReadOnly()
-		if read.amended {
-			read = readOnly[K, V]{m: m.dirty}
-			m.read.Store(&read)
-			m.dirty = nil
-			m.misses = 0
-		}
-		m.mu.Unlock()
-	}
-
-	return len(read.m)
-}
-
 func (m *Map[K, V]) missLocked() {
 	m.misses++
 	if m.misses < len(m.dirty) {
@@ -367,14 +383,18 @@ func (m *Map[K, V]) dirtyLocked() {
 	}
 
 	read := m.loadReadOnly()
-	m.dirty = make(map[K]*entry[V], len(read.m))
+	m.initDirty(len(read.m))
 	for k, e := range read.m {
 		if !e.tryExpungeLocked() {
 			m.dirty[k] = e
 		}
 	}
 }
 
+func (m *Map[K, V]) initDirty(size int) {
+	m.dirty = make(map[K]*entry[V], size)
+}
+
 func (e *entry[V]) tryExpungeLocked() (isExpunged bool) {
 	p := e.p.Load()
 	for p == nil {
@@ -386,6 +406,23 @@ func (e *entry[V]) tryExpungeLocked() (isExpunged bool) {
 	return p == e.expunged.Load()
 }
 
+func (m *Map[K, V]) Len() int {
+	read := m.loadReadOnly()
+	if read.amended {
+		m.mu.Lock()
+		read = m.loadReadOnly()
+		if read.amended {
+			read = readOnly[K, V]{m: m.dirty}
+			m.read.Store(&read)
+			m.dirty = nil
+			m.misses = 0
+		}
+		m.mu.Unlock()
+	}
+
+	return len(read.m)
+}
+
 func (m *Map[K, V]) Size() (size uintptr) {
 	if m == nil {
 		return 0