Source file src/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  	isync "internal/sync"
     9  	"sync"
    10  	"sync/atomic"
    11  )
    12  
    13  // This file contains reference map implementations for unit-tests.
    14  
    15  // mapInterface is the interface Map implements.
    16  type mapInterface interface {
    17  	Load(key any) (value any, ok bool)
    18  	Store(key, value any)
    19  	LoadOrStore(key, value any) (actual any, loaded bool)
    20  	LoadAndDelete(key any) (value any, loaded bool)
    21  	Delete(any)
    22  	Swap(key, value any) (previous any, loaded bool)
    23  	CompareAndSwap(key, old, new any) (swapped bool)
    24  	CompareAndDelete(key, old any) (deleted bool)
    25  	Range(func(key, value any) (shouldContinue bool))
    26  	Clear()
    27  }
    28  
    29  var (
    30  	_ mapInterface = &RWMutexMap{}
    31  	_ mapInterface = &DeepCopyMap{}
    32  	_ mapInterface = &isync.HashTrieMap[any, any]{}
    33  )
    34  
    35  // RWMutexMap is an implementation of mapInterface using a sync.RWMutex.
    36  type RWMutexMap struct {
    37  	mu    sync.RWMutex
    38  	dirty map[any]any
    39  }
    40  
    41  func (m *RWMutexMap) Load(key any) (value any, ok bool) {
    42  	m.mu.RLock()
    43  	value, ok = m.dirty[key]
    44  	m.mu.RUnlock()
    45  	return
    46  }
    47  
    48  func (m *RWMutexMap) Store(key, value any) {
    49  	m.mu.Lock()
    50  	if m.dirty == nil {
    51  		m.dirty = make(map[any]any)
    52  	}
    53  	m.dirty[key] = value
    54  	m.mu.Unlock()
    55  }
    56  
    57  func (m *RWMutexMap) LoadOrStore(key, value any) (actual any, loaded bool) {
    58  	m.mu.Lock()
    59  	actual, loaded = m.dirty[key]
    60  	if !loaded {
    61  		actual = value
    62  		if m.dirty == nil {
    63  			m.dirty = make(map[any]any)
    64  		}
    65  		m.dirty[key] = value
    66  	}
    67  	m.mu.Unlock()
    68  	return actual, loaded
    69  }
    70  
    71  func (m *RWMutexMap) Swap(key, value any) (previous any, loaded bool) {
    72  	m.mu.Lock()
    73  	if m.dirty == nil {
    74  		m.dirty = make(map[any]any)
    75  	}
    76  
    77  	previous, loaded = m.dirty[key]
    78  	m.dirty[key] = value
    79  	m.mu.Unlock()
    80  	return
    81  }
    82  
    83  func (m *RWMutexMap) LoadAndDelete(key any) (value any, loaded bool) {
    84  	m.mu.Lock()
    85  	value, loaded = m.dirty[key]
    86  	if !loaded {
    87  		m.mu.Unlock()
    88  		return nil, false
    89  	}
    90  	delete(m.dirty, key)
    91  	m.mu.Unlock()
    92  	return value, loaded
    93  }
    94  
    95  func (m *RWMutexMap) Delete(key any) {
    96  	m.mu.Lock()
    97  	delete(m.dirty, key)
    98  	m.mu.Unlock()
    99  }
   100  
   101  func (m *RWMutexMap) CompareAndSwap(key, old, new any) (swapped bool) {
   102  	m.mu.Lock()
   103  	defer m.mu.Unlock()
   104  	if m.dirty == nil {
   105  		return false
   106  	}
   107  
   108  	value, loaded := m.dirty[key]
   109  	if loaded && value == old {
   110  		m.dirty[key] = new
   111  		return true
   112  	}
   113  	return false
   114  }
   115  
   116  func (m *RWMutexMap) CompareAndDelete(key, old any) (deleted bool) {
   117  	m.mu.Lock()
   118  	defer m.mu.Unlock()
   119  	if m.dirty == nil {
   120  		return false
   121  	}
   122  
   123  	value, loaded := m.dirty[key]
   124  	if loaded && value == old {
   125  		delete(m.dirty, key)
   126  		return true
   127  	}
   128  	return false
   129  }
   130  
   131  func (m *RWMutexMap) Range(f func(key, value any) (shouldContinue bool)) {
   132  	m.mu.RLock()
   133  	keys := make([]any, 0, len(m.dirty))
   134  	for k := range m.dirty {
   135  		keys = append(keys, k)
   136  	}
   137  	m.mu.RUnlock()
   138  
   139  	for _, k := range keys {
   140  		v, ok := m.Load(k)
   141  		if !ok {
   142  			continue
   143  		}
   144  		if !f(k, v) {
   145  			break
   146  		}
   147  	}
   148  }
   149  
   150  func (m *RWMutexMap) Clear() {
   151  	m.mu.Lock()
   152  	defer m.mu.Unlock()
   153  
   154  	clear(m.dirty)
   155  }
   156  
   157  // DeepCopyMap is an implementation of mapInterface using a Mutex and
   158  // atomic.Value.  It makes deep copies of the map on every write to avoid
   159  // acquiring the Mutex in Load.
   160  type DeepCopyMap struct {
   161  	mu    sync.Mutex
   162  	clean atomic.Value
   163  }
   164  
   165  func (m *DeepCopyMap) Load(key any) (value any, ok bool) {
   166  	clean, _ := m.clean.Load().(map[any]any)
   167  	value, ok = clean[key]
   168  	return value, ok
   169  }
   170  
   171  func (m *DeepCopyMap) Store(key, value any) {
   172  	m.mu.Lock()
   173  	dirty := m.dirty()
   174  	dirty[key] = value
   175  	m.clean.Store(dirty)
   176  	m.mu.Unlock()
   177  }
   178  
   179  func (m *DeepCopyMap) LoadOrStore(key, value any) (actual any, loaded bool) {
   180  	clean, _ := m.clean.Load().(map[any]any)
   181  	actual, loaded = clean[key]
   182  	if loaded {
   183  		return actual, loaded
   184  	}
   185  
   186  	m.mu.Lock()
   187  	// Reload clean in case it changed while we were waiting on m.mu.
   188  	clean, _ = m.clean.Load().(map[any]any)
   189  	actual, loaded = clean[key]
   190  	if !loaded {
   191  		dirty := m.dirty()
   192  		dirty[key] = value
   193  		actual = value
   194  		m.clean.Store(dirty)
   195  	}
   196  	m.mu.Unlock()
   197  	return actual, loaded
   198  }
   199  
   200  func (m *DeepCopyMap) Swap(key, value any) (previous any, loaded bool) {
   201  	m.mu.Lock()
   202  	dirty := m.dirty()
   203  	previous, loaded = dirty[key]
   204  	dirty[key] = value
   205  	m.clean.Store(dirty)
   206  	m.mu.Unlock()
   207  	return
   208  }
   209  
   210  func (m *DeepCopyMap) LoadAndDelete(key any) (value any, loaded bool) {
   211  	m.mu.Lock()
   212  	dirty := m.dirty()
   213  	value, loaded = dirty[key]
   214  	delete(dirty, key)
   215  	m.clean.Store(dirty)
   216  	m.mu.Unlock()
   217  	return
   218  }
   219  
   220  func (m *DeepCopyMap) Delete(key any) {
   221  	m.mu.Lock()
   222  	dirty := m.dirty()
   223  	delete(dirty, key)
   224  	m.clean.Store(dirty)
   225  	m.mu.Unlock()
   226  }
   227  
   228  func (m *DeepCopyMap) CompareAndSwap(key, old, new any) (swapped bool) {
   229  	clean, _ := m.clean.Load().(map[any]any)
   230  	if previous, ok := clean[key]; !ok || previous != old {
   231  		return false
   232  	}
   233  
   234  	m.mu.Lock()
   235  	defer m.mu.Unlock()
   236  	dirty := m.dirty()
   237  	value, loaded := dirty[key]
   238  	if loaded && value == old {
   239  		dirty[key] = new
   240  		m.clean.Store(dirty)
   241  		return true
   242  	}
   243  	return false
   244  }
   245  
   246  func (m *DeepCopyMap) CompareAndDelete(key, old any) (deleted bool) {
   247  	clean, _ := m.clean.Load().(map[any]any)
   248  	if previous, ok := clean[key]; !ok || previous != old {
   249  		return false
   250  	}
   251  
   252  	m.mu.Lock()
   253  	defer m.mu.Unlock()
   254  
   255  	dirty := m.dirty()
   256  	value, loaded := dirty[key]
   257  	if loaded && value == old {
   258  		delete(dirty, key)
   259  		m.clean.Store(dirty)
   260  		return true
   261  	}
   262  	return false
   263  }
   264  
   265  func (m *DeepCopyMap) Range(f func(key, value any) (shouldContinue bool)) {
   266  	clean, _ := m.clean.Load().(map[any]any)
   267  	for k, v := range clean {
   268  		if !f(k, v) {
   269  			break
   270  		}
   271  	}
   272  }
   273  
   274  func (m *DeepCopyMap) dirty() map[any]any {
   275  	clean, _ := m.clean.Load().(map[any]any)
   276  	dirty := make(map[any]any, len(clean)+1)
   277  	for k, v := range clean {
   278  		dirty[k] = v
   279  	}
   280  	return dirty
   281  }
   282  
   283  func (m *DeepCopyMap) Clear() {
   284  	m.mu.Lock()
   285  	defer m.mu.Unlock()
   286  
   287  	m.clean.Store((map[any]any)(nil))
   288  }
   289  

View as plain text