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(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 }