github.com/ActiveState/go@v0.0.0-20170614201249-0b81c023a722/src/sync/map_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 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  	defer m.mu.RUnlock()
    68  	for k, v := range m.dirty {
    69  		if !f(k, v) {
    70  			break
    71  		}
    72  	}
    73  }
    74  
    75  // DeepCopyMap is an implementation of mapInterface using a Mutex and
    76  // atomic.Value.  It makes deep copies of the map on every write to avoid
    77  // acquiring the Mutex in Load.
    78  type DeepCopyMap struct {
    79  	mu    sync.Mutex
    80  	clean atomic.Value
    81  }
    82  
    83  func (m *DeepCopyMap) Load(key interface{}) (value interface{}, ok bool) {
    84  	clean, _ := m.clean.Load().(map[interface{}]interface{})
    85  	value, ok = clean[key]
    86  	return value, ok
    87  }
    88  
    89  func (m *DeepCopyMap) Store(key, value interface{}) {
    90  	m.mu.Lock()
    91  	dirty := m.dirty()
    92  	dirty[key] = value
    93  	m.clean.Store(dirty)
    94  	m.mu.Unlock()
    95  }
    96  
    97  func (m *DeepCopyMap) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) {
    98  	clean, _ := m.clean.Load().(map[interface{}]interface{})
    99  	actual, loaded = clean[key]
   100  	if loaded {
   101  		return actual, loaded
   102  	}
   103  
   104  	m.mu.Lock()
   105  	// Reload clean in case it changed while we were waiting on m.mu.
   106  	clean, _ = m.clean.Load().(map[interface{}]interface{})
   107  	actual, loaded = clean[key]
   108  	if !loaded {
   109  		dirty := m.dirty()
   110  		dirty[key] = value
   111  		actual = value
   112  		m.clean.Store(dirty)
   113  	}
   114  	m.mu.Unlock()
   115  	return actual, loaded
   116  }
   117  
   118  func (m *DeepCopyMap) Delete(key interface{}) {
   119  	m.mu.Lock()
   120  	dirty := m.dirty()
   121  	delete(dirty, key)
   122  	m.clean.Store(dirty)
   123  	m.mu.Unlock()
   124  }
   125  
   126  func (m *DeepCopyMap) Range(f func(key, value interface{}) (shouldContinue bool)) {
   127  	clean, _ := m.clean.Load().(map[interface{}]interface{})
   128  	for k, v := range clean {
   129  		if !f(k, v) {
   130  			break
   131  		}
   132  	}
   133  }
   134  
   135  func (m *DeepCopyMap) dirty() map[interface{}]interface{} {
   136  	clean, _ := m.clean.Load().(map[interface{}]interface{})
   137  	dirty := make(map[interface{}]interface{}, len(clean)+1)
   138  	for k, v := range clean {
   139  		dirty[k] = v
   140  	}
   141  	return dirty
   142  }