Home | History | Annotate | Download | only in sync 
      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 }
    152