libgo/go/sync/map_reference_test.go

   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 sync_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(any) (any, bool)
  17         Store(key, value any)
  18         LoadOrStore(key, value any) (actual any, loaded bool)
  19         LoadAndDelete(key any) (value any, loaded bool)
  20         Delete(any)
  21         Range(func(key, value any) (shouldContinue bool))
  22 }
  23
  24 // RWMutexMap is an implementation of mapInterface using a sync.RWMutex.
  25 type RWMutexMap struct {
  26         mu    sync.RWMutex
  27         dirty map[any]any
  28 }
  29
  30 func (m *RWMutexMap) Load(key any) (value any, ok bool) {
  31         m.mu.RLock()
  32         value, ok = m.dirty[key]
  33         m.mu.RUnlock()
  34         return
  35 }
  36
  37 func (m *RWMutexMap) Store(key, value any) {
  38         m.mu.Lock()
  39         if m.dirty == nil {
  40                 m.dirty = make(map[any]any)
  41         }
  42         m.dirty[key] = value
  43         m.mu.Unlock()
  44 }
  45
  46 func (m *RWMutexMap) LoadOrStore(key, value any) (actual any, loaded bool) {
  47         m.mu.Lock()
  48         actual, loaded = m.dirty[key]
  49         if !loaded {
  50                 actual = value
  51                 if m.dirty == nil {
  52                         m.dirty = make(map[any]any)
  53                 }
  54                 m.dirty[key] = value
  55         }
  56         m.mu.Unlock()
  57         return actual, loaded
  58 }
  59
  60 func (m *RWMutexMap) LoadAndDelete(key any) (value any, loaded bool) {
  61         m.mu.Lock()
  62         value, loaded = m.dirty[key]
  63         if !loaded {
  64                 m.mu.Unlock()
  65                 return nil, false
  66         }
  67         delete(m.dirty, key)
  68         m.mu.Unlock()
  69         return value, loaded
  70 }
  71
  72 func (m *RWMutexMap) Delete(key any) {
  73         m.mu.Lock()
  74         delete(m.dirty, key)
  75         m.mu.Unlock()
  76 }
  77
  78 func (m *RWMutexMap) Range(f func(key, value any) (shouldContinue bool)) {
  79         m.mu.RLock()
  80         keys := make([]any, 0, len(m.dirty))
  81         for k := range m.dirty {
  82                 keys = append(keys, k)
  83         }
  84         m.mu.RUnlock()
  85
  86         for _, k := range keys {
  87                 v, ok := m.Load(k)
  88                 if !ok {
  89                         continue
  90                 }
  91                 if !f(k, v) {
  92                         break
  93                 }
  94         }
  95 }
  96
  97 // DeepCopyMap is an implementation of mapInterface using a Mutex and
  98 // atomic.Value.  It makes deep copies of the map on every write to avoid
  99 // acquiring the Mutex in Load.
 100 type DeepCopyMap struct {
 101         mu    sync.Mutex
 102         clean atomic.Value
 103 }
 104
 105 func (m *DeepCopyMap) Load(key any) (value any, ok bool) {
 106         clean, _ := m.clean.Load().(map[any]any)
 107         value, ok = clean[key]
 108         return value, ok
 109 }
 110
 111 func (m *DeepCopyMap) Store(key, value any) {
 112         m.mu.Lock()
 113         dirty := m.dirty()
 114         dirty[key] = value
 115         m.clean.Store(dirty)
 116         m.mu.Unlock()
 117 }
 118
 119 func (m *DeepCopyMap) LoadOrStore(key, value any) (actual any, loaded bool) {
 120         clean, _ := m.clean.Load().(map[any]any)
 121         actual, loaded = clean[key]
 122         if loaded {
 123                 return actual, loaded
 124         }
 125
 126         m.mu.Lock()
 127         // Reload clean in case it changed while we were waiting on m.mu.
 128         clean, _ = m.clean.Load().(map[any]any)
 129         actual, loaded = clean[key]
 130         if !loaded {
 131                 dirty := m.dirty()
 132                 dirty[key] = value
 133                 actual = value
 134                 m.clean.Store(dirty)
 135         }
 136         m.mu.Unlock()
 137         return actual, loaded
 138 }
 139
 140 func (m *DeepCopyMap) LoadAndDelete(key any) (value any, loaded bool) {
 141         m.mu.Lock()
 142         dirty := m.dirty()
 143         value, loaded = dirty[key]
 144         delete(dirty, key)
 145         m.clean.Store(dirty)
 146         m.mu.Unlock()
 147         return
 148 }
 149
 150 func (m *DeepCopyMap) Delete(key any) {
 151         m.mu.Lock()
 152         dirty := m.dirty()
 153         delete(dirty, key)
 154         m.clean.Store(dirty)
 155         m.mu.Unlock()
 156 }
 157
 158 func (m *DeepCopyMap) Range(f func(key, value any) (shouldContinue bool)) {
 159         clean, _ := m.clean.Load().(map[any]any)
 160         for k, v := range clean {
 161                 if !f(k, v) {
 162                         break
 163                 }
 164         }
 165 }
 166
 167 func (m *DeepCopyMap) dirty() map[any]any {
 168         clean, _ := m.clean.Load().(map[any]any)
 169         dirty := make(map[any]any, len(clean)+1)
 170         for k, v := range clean {
 171                 dirty[k] = v
 172         }
 173         return dirty
 174 }