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
5package sync_test
6
7import (
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.
15type 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.
24type RWMutexMap struct {
25	mu    sync.RWMutex
26	dirty map[interface{}]interface{}
27}
28
29func (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
36func (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
45func (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
59func (m *RWMutexMap) Delete(key interface{}) {
60	m.mu.Lock()
61	delete(m.dirty, key)
62	m.mu.Unlock()
63}
64
65func (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.
87type DeepCopyMap struct {
88	mu    sync.Mutex
89	clean atomic.Value
90}
91
92func (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
98func (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
106func (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
127func (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
135func (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
144func (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