github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/syncutil/int_map_reference_test.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  // Copyright 2016 The Go Authors. All rights reserved.
    12  // Use of this source code is governed by a BSD-style
    13  // license that can be found in licenses/BSD-golang.txt.
    14  
    15  // This code originated in Go's sync package.
    16  
    17  package syncutil
    18  
    19  import (
    20  	"sync/atomic"
    21  	"unsafe"
    22  )
    23  
    24  // This file contains reference map implementations for unit-tests.
    25  
    26  // mapInterface is the interface Map implements.
    27  type mapInterface interface {
    28  	Load(int64) (unsafe.Pointer, bool)
    29  	Store(key int64, value unsafe.Pointer)
    30  	LoadOrStore(key int64, value unsafe.Pointer) (actual unsafe.Pointer, loaded bool)
    31  	Delete(int64)
    32  	Range(func(key int64, value unsafe.Pointer) (shouldContinue bool))
    33  }
    34  
    35  // RWMutexMap is an implementation of mapInterface using a RWMutex.
    36  type RWMutexMap struct {
    37  	mu    RWMutex
    38  	dirty map[int64]unsafe.Pointer
    39  }
    40  
    41  func (m *RWMutexMap) Load(key int64) (value unsafe.Pointer, 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 int64, value unsafe.Pointer) {
    49  	m.mu.Lock()
    50  	if m.dirty == nil {
    51  		m.dirty = make(map[int64]unsafe.Pointer)
    52  	}
    53  	m.dirty[key] = value
    54  	m.mu.Unlock()
    55  }
    56  
    57  func (m *RWMutexMap) LoadOrStore(
    58  	key int64, value unsafe.Pointer,
    59  ) (actual unsafe.Pointer, loaded bool) {
    60  	m.mu.Lock()
    61  	actual, loaded = m.dirty[key]
    62  	if !loaded {
    63  		actual = value
    64  		if m.dirty == nil {
    65  			m.dirty = make(map[int64]unsafe.Pointer)
    66  		}
    67  		m.dirty[key] = value
    68  	}
    69  	m.mu.Unlock()
    70  	return actual, loaded
    71  }
    72  
    73  func (m *RWMutexMap) Delete(key int64) {
    74  	m.mu.Lock()
    75  	delete(m.dirty, key)
    76  	m.mu.Unlock()
    77  }
    78  
    79  func (m *RWMutexMap) Range(f func(key int64, value unsafe.Pointer) (shouldContinue bool)) {
    80  	m.mu.RLock()
    81  	keys := make([]int64, 0, len(m.dirty))
    82  	for k := range m.dirty {
    83  		keys = append(keys, k)
    84  	}
    85  	m.mu.RUnlock()
    86  
    87  	for _, k := range keys {
    88  		v, ok := m.Load(k)
    89  		if !ok {
    90  			continue
    91  		}
    92  		if !f(k, v) {
    93  			break
    94  		}
    95  	}
    96  }
    97  
    98  // DeepCopyMap is an implementation of mapInterface using a Mutex and
    99  // atomic.Value.  It makes deep copies of the map on every write to avoid
   100  // acquiring the Mutex in Load.
   101  type DeepCopyMap struct {
   102  	mu    Mutex
   103  	clean atomic.Value
   104  }
   105  
   106  func (m *DeepCopyMap) Load(key int64) (value unsafe.Pointer, ok bool) {
   107  	clean, _ := m.clean.Load().(map[int64]unsafe.Pointer)
   108  	value, ok = clean[key]
   109  	return value, ok
   110  }
   111  
   112  func (m *DeepCopyMap) Store(key int64, value unsafe.Pointer) {
   113  	m.mu.Lock()
   114  	dirty := m.dirty()
   115  	dirty[key] = value
   116  	m.clean.Store(dirty)
   117  	m.mu.Unlock()
   118  }
   119  
   120  func (m *DeepCopyMap) LoadOrStore(
   121  	key int64, value unsafe.Pointer,
   122  ) (actual unsafe.Pointer, loaded bool) {
   123  	clean, _ := m.clean.Load().(map[int64]unsafe.Pointer)
   124  	actual, loaded = clean[key]
   125  	if loaded {
   126  		return actual, loaded
   127  	}
   128  
   129  	m.mu.Lock()
   130  	// Reload clean in case it changed while we were waiting on m.mu.
   131  	clean, _ = m.clean.Load().(map[int64]unsafe.Pointer)
   132  	actual, loaded = clean[key]
   133  	if !loaded {
   134  		dirty := m.dirty()
   135  		dirty[key] = value
   136  		actual = value
   137  		m.clean.Store(dirty)
   138  	}
   139  	m.mu.Unlock()
   140  	return actual, loaded
   141  }
   142  
   143  func (m *DeepCopyMap) Delete(key int64) {
   144  	m.mu.Lock()
   145  	dirty := m.dirty()
   146  	delete(dirty, key)
   147  	m.clean.Store(dirty)
   148  	m.mu.Unlock()
   149  }
   150  
   151  func (m *DeepCopyMap) Range(f func(key int64, value unsafe.Pointer) (shouldContinue bool)) {
   152  	clean, _ := m.clean.Load().(map[int64]unsafe.Pointer)
   153  	for k, v := range clean {
   154  		if !f(k, v) {
   155  			break
   156  		}
   157  	}
   158  }
   159  
   160  func (m *DeepCopyMap) dirty() map[int64]unsafe.Pointer {
   161  	clean, _ := m.clean.Load().(map[int64]unsafe.Pointer)
   162  	dirty := make(map[int64]unsafe.Pointer, len(clean)+1)
   163  	for k, v := range clean {
   164  		dirty[k] = v
   165  	}
   166  	return dirty
   167  }