github.com/lrita/cmap@v0.0.0-20231108122212-cb084a67f554/cmap_reference_test.go (about) 1 // Copyright 2016 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 package cmap_test 6 7 import ( 8 "sync" 9 "sync/atomic" 10 ) 11 12 // This file contains reference map implementations for unit-tests. 13 14 // mapInterface is the interface Map implements. 15 type mapInterface interface { 16 Load(interface{}) (interface{}, bool) 17 Store(key, value interface{}) 18 LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) 19 Delete(interface{}) 20 Range(func(key, value interface{}) (shouldContinue bool)) 21 } 22 23 // RWMutexMap is an implementation of mapInterface using a sync.RWMutex. 24 type RWMutexMap struct { 25 mu sync.RWMutex 26 dirty map[interface{}]interface{} 27 } 28 29 func (m *RWMutexMap) Load(key interface{}) (value interface{}, ok bool) { 30 m.mu.RLock() 31 value, ok = m.dirty[key] 32 m.mu.RUnlock() 33 return 34 } 35 36 func (m *RWMutexMap) Store(key, value interface{}) { 37 m.mu.Lock() 38 if m.dirty == nil { 39 m.dirty = make(map[interface{}]interface{}) 40 } 41 m.dirty[key] = value 42 m.mu.Unlock() 43 } 44 45 func (m *RWMutexMap) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) { 46 m.mu.Lock() 47 actual, loaded = m.dirty[key] 48 if !loaded { 49 actual = value 50 if m.dirty == nil { 51 m.dirty = make(map[interface{}]interface{}) 52 } 53 m.dirty[key] = value 54 } 55 m.mu.Unlock() 56 return actual, loaded 57 } 58 59 func (m *RWMutexMap) Delete(key interface{}) { 60 m.mu.Lock() 61 delete(m.dirty, key) 62 m.mu.Unlock() 63 } 64 65 func (m *RWMutexMap) Range(f func(key, value interface{}) (shouldContinue bool)) { 66 m.mu.RLock() 67 keys := make([]interface{}, 0, len(m.dirty)) 68 for k := range m.dirty { 69 keys = append(keys, k) 70 } 71 m.mu.RUnlock() 72 73 for _, k := range keys { 74 v, ok := m.Load(k) 75 if !ok { 76 continue 77 } 78 if !f(k, v) { 79 break 80 } 81 } 82 } 83 84 // DeepCopyMap is an implementation of mapInterface using a Mutex and 85 // atomic.Value. It makes deep copies of the map on every write to avoid 86 // acquiring the Mutex in Load. 87 type DeepCopyMap struct { 88 mu sync.Mutex 89 clean atomic.Value 90 } 91 92 func (m *DeepCopyMap) Load(key interface{}) (value interface{}, ok bool) { 93 clean, _ := m.clean.Load().(map[interface{}]interface{}) 94 value, ok = clean[key] 95 return value, ok 96 } 97 98 func (m *DeepCopyMap) Store(key, value interface{}) { 99 m.mu.Lock() 100 dirty := m.dirty() 101 dirty[key] = value 102 m.clean.Store(dirty) 103 m.mu.Unlock() 104 } 105 106 func (m *DeepCopyMap) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) { 107 clean, _ := m.clean.Load().(map[interface{}]interface{}) 108 actual, loaded = clean[key] 109 if loaded { 110 return actual, loaded 111 } 112 113 m.mu.Lock() 114 // Reload clean in case it changed while we were waiting on m.mu. 115 clean, _ = m.clean.Load().(map[interface{}]interface{}) 116 actual, loaded = clean[key] 117 if !loaded { 118 dirty := m.dirty() 119 dirty[key] = value 120 actual = value 121 m.clean.Store(dirty) 122 } 123 m.mu.Unlock() 124 return actual, loaded 125 } 126 127 func (m *DeepCopyMap) Delete(key interface{}) { 128 m.mu.Lock() 129 dirty := m.dirty() 130 delete(dirty, key) 131 m.clean.Store(dirty) 132 m.mu.Unlock() 133 } 134 135 func (m *DeepCopyMap) Range(f func(key, value interface{}) (shouldContinue bool)) { 136 clean, _ := m.clean.Load().(map[interface{}]interface{}) 137 for k, v := range clean { 138 if !f(k, v) { 139 break 140 } 141 } 142 } 143 144 func (m *DeepCopyMap) dirty() map[interface{}]interface{} { 145 clean, _ := m.clean.Load().(map[interface{}]interface{}) 146 dirty := make(map[interface{}]interface{}, len(clean)+1) 147 for k, v := range clean { 148 dirty[k] = v 149 } 150 return dirty 151 }