Source file src/sync/map.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
     6  
     7  import (
     8  	"sync/atomic"
     9  )
    10  
    11  // Map is like a Go map[any]any but is safe for concurrent use
    12  // by multiple goroutines without additional locking or coordination.
    13  // Loads, stores, and deletes run in amortized constant time.
    14  //
    15  // The Map type is specialized. Most code should use a plain Go map instead,
    16  // with separate locking or coordination, for better type safety and to make it
    17  // easier to maintain other invariants along with the map content.
    18  //
    19  // The Map type is optimized for two common use cases: (1) when the entry for a given
    20  // key is only ever written once but read many times, as in caches that only grow,
    21  // or (2) when multiple goroutines read, write, and overwrite entries for disjoint
    22  // sets of keys. In these two cases, use of a Map may significantly reduce lock
    23  // contention compared to a Go map paired with a separate [Mutex] or [RWMutex].
    24  //
    25  // The zero Map is empty and ready for use. A Map must not be copied after first use.
    26  //
    27  // In the terminology of the Go memory model, Map arranges that a write operation
    28  // “synchronizes before” any read operation that observes the effect of the write, where
    29  // read and write operations are defined as follows.
    30  // [Map.Load], [Map.LoadAndDelete], [Map.LoadOrStore], [Map.Swap], [Map.CompareAndSwap],
    31  // and [Map.CompareAndDelete] are read operations;
    32  // [Map.Delete], [Map.LoadAndDelete], [Map.Store], and [Map.Swap] are write operations;
    33  // [Map.LoadOrStore] is a write operation when it returns loaded set to false;
    34  // [Map.CompareAndSwap] is a write operation when it returns swapped set to true;
    35  // and [Map.CompareAndDelete] is a write operation when it returns deleted set to true.
    36  type Map struct {
    37  	mu Mutex
    38  
    39  	// read contains the portion of the map's contents that are safe for
    40  	// concurrent access (with or without mu held).
    41  	//
    42  	// The read field itself is always safe to load, but must only be stored with
    43  	// mu held.
    44  	//
    45  	// Entries stored in read may be updated concurrently without mu, but updating
    46  	// a previously-expunged entry requires that the entry be copied to the dirty
    47  	// map and unexpunged with mu held.
    48  	read atomic.Pointer[readOnly]
    49  
    50  	// dirty contains the portion of the map's contents that require mu to be
    51  	// held. To ensure that the dirty map can be promoted to the read map quickly,
    52  	// it also includes all of the non-expunged entries in the read map.
    53  	//
    54  	// Expunged entries are not stored in the dirty map. An expunged entry in the
    55  	// clean map must be unexpunged and added to the dirty map before a new value
    56  	// can be stored to it.
    57  	//
    58  	// If the dirty map is nil, the next write to the map will initialize it by
    59  	// making a shallow copy of the clean map, omitting stale entries.
    60  	dirty map[any]*entry
    61  
    62  	// misses counts the number of loads since the read map was last updated that
    63  	// needed to lock mu to determine whether the key was present.
    64  	//
    65  	// Once enough misses have occurred to cover the cost of copying the dirty
    66  	// map, the dirty map will be promoted to the read map (in the unamended
    67  	// state) and the next store to the map will make a new dirty copy.
    68  	misses int
    69  }
    70  
    71  // readOnly is an immutable struct stored atomically in the Map.read field.
    72  type readOnly struct {
    73  	m       map[any]*entry
    74  	amended bool // true if the dirty map contains some key not in m.
    75  }
    76  
    77  // expunged is an arbitrary pointer that marks entries which have been deleted
    78  // from the dirty map.
    79  var expunged = new(any)
    80  
    81  // An entry is a slot in the map corresponding to a particular key.
    82  type entry struct {
    83  	// p points to the interface{} value stored for the entry.
    84  	//
    85  	// If p == nil, the entry has been deleted, and either m.dirty == nil or
    86  	// m.dirty[key] is e.
    87  	//
    88  	// If p == expunged, the entry has been deleted, m.dirty != nil, and the entry
    89  	// is missing from m.dirty.
    90  	//
    91  	// Otherwise, the entry is valid and recorded in m.read.m[key] and, if m.dirty
    92  	// != nil, in m.dirty[key].
    93  	//
    94  	// An entry can be deleted by atomic replacement with nil: when m.dirty is
    95  	// next created, it will atomically replace nil with expunged and leave
    96  	// m.dirty[key] unset.
    97  	//
    98  	// An entry's associated value can be updated by atomic replacement, provided
    99  	// p != expunged. If p == expunged, an entry's associated value can be updated
   100  	// only after first setting m.dirty[key] = e so that lookups using the dirty
   101  	// map find the entry.
   102  	p atomic.Pointer[any]
   103  }
   104  
   105  func newEntry(i any) *entry {
   106  	e := &entry{}
   107  	e.p.Store(&i)
   108  	return e
   109  }
   110  
   111  func (m *Map) loadReadOnly() readOnly {
   112  	if p := m.read.Load(); p != nil {
   113  		return *p
   114  	}
   115  	return readOnly{}
   116  }
   117  
   118  // Load returns the value stored in the map for a key, or nil if no
   119  // value is present.
   120  // The ok result indicates whether value was found in the map.
   121  func (m *Map) Load(key any) (value any, ok bool) {
   122  	read := m.loadReadOnly()
   123  	e, ok := read.m[key]
   124  	if !ok && read.amended {
   125  		m.mu.Lock()
   126  		// Avoid reporting a spurious miss if m.dirty got promoted while we were
   127  		// blocked on m.mu. (If further loads of the same key will not miss, it's
   128  		// not worth copying the dirty map for this key.)
   129  		read = m.loadReadOnly()
   130  		e, ok = read.m[key]
   131  		if !ok && read.amended {
   132  			e, ok = m.dirty[key]
   133  			// Regardless of whether the entry was present, record a miss: this key
   134  			// will take the slow path until the dirty map is promoted to the read
   135  			// map.
   136  			m.missLocked()
   137  		}
   138  		m.mu.Unlock()
   139  	}
   140  	if !ok {
   141  		return nil, false
   142  	}
   143  	return e.load()
   144  }
   145  
   146  func (e *entry) load() (value any, ok bool) {
   147  	p := e.p.Load()
   148  	if p == nil || p == expunged {
   149  		return nil, false
   150  	}
   151  	return *p, true
   152  }
   153  
   154  // Store sets the value for a key.
   155  func (m *Map) Store(key, value any) {
   156  	_, _ = m.Swap(key, value)
   157  }
   158  
   159  // Clear deletes all the entries, resulting in an empty Map.
   160  func (m *Map) Clear() {
   161  	read := m.loadReadOnly()
   162  	if len(read.m) == 0 && !read.amended {
   163  		// Avoid allocating a new readOnly when the map is already clear.
   164  		return
   165  	}
   166  
   167  	m.mu.Lock()
   168  	defer m.mu.Unlock()
   169  
   170  	read = m.loadReadOnly()
   171  	if len(read.m) > 0 || read.amended {
   172  		m.read.Store(&readOnly{})
   173  	}
   174  
   175  	clear(m.dirty)
   176  	// Don't immediately promote the newly-cleared dirty map on the next operation.
   177  	m.misses = 0
   178  }
   179  
   180  // tryCompareAndSwap compare the entry with the given old value and swaps
   181  // it with a new value if the entry is equal to the old value, and the entry
   182  // has not been expunged.
   183  //
   184  // If the entry is expunged, tryCompareAndSwap returns false and leaves
   185  // the entry unchanged.
   186  func (e *entry) tryCompareAndSwap(old, new any) bool {
   187  	p := e.p.Load()
   188  	if p == nil || p == expunged || *p != old {
   189  		return false
   190  	}
   191  
   192  	// Copy the interface after the first load to make this method more amenable
   193  	// to escape analysis: if the comparison fails from the start, we shouldn't
   194  	// bother heap-allocating an interface value to store.
   195  	nc := new
   196  	for {
   197  		if e.p.CompareAndSwap(p, &nc) {
   198  			return true
   199  		}
   200  		p = e.p.Load()
   201  		if p == nil || p == expunged || *p != old {
   202  			return false
   203  		}
   204  	}
   205  }
   206  
   207  // unexpungeLocked ensures that the entry is not marked as expunged.
   208  //
   209  // If the entry was previously expunged, it must be added to the dirty map
   210  // before m.mu is unlocked.
   211  func (e *entry) unexpungeLocked() (wasExpunged bool) {
   212  	return e.p.CompareAndSwap(expunged, nil)
   213  }
   214  
   215  // swapLocked unconditionally swaps a value into the entry.
   216  //
   217  // The entry must be known not to be expunged.
   218  func (e *entry) swapLocked(i *any) *any {
   219  	return e.p.Swap(i)
   220  }
   221  
   222  // LoadOrStore returns the existing value for the key if present.
   223  // Otherwise, it stores and returns the given value.
   224  // The loaded result is true if the value was loaded, false if stored.
   225  func (m *Map) LoadOrStore(key, value any) (actual any, loaded bool) {
   226  	// Avoid locking if it's a clean hit.
   227  	read := m.loadReadOnly()
   228  	if e, ok := read.m[key]; ok {
   229  		actual, loaded, ok := e.tryLoadOrStore(value)
   230  		if ok {
   231  			return actual, loaded
   232  		}
   233  	}
   234  
   235  	m.mu.Lock()
   236  	read = m.loadReadOnly()
   237  	if e, ok := read.m[key]; ok {
   238  		if e.unexpungeLocked() {
   239  			m.dirty[key] = e
   240  		}
   241  		actual, loaded, _ = e.tryLoadOrStore(value)
   242  	} else if e, ok := m.dirty[key]; ok {
   243  		actual, loaded, _ = e.tryLoadOrStore(value)
   244  		m.missLocked()
   245  	} else {
   246  		if !read.amended {
   247  			// We're adding the first new key to the dirty map.
   248  			// Make sure it is allocated and mark the read-only map as incomplete.
   249  			m.dirtyLocked()
   250  			m.read.Store(&readOnly{m: read.m, amended: true})
   251  		}
   252  		m.dirty[key] = newEntry(value)
   253  		actual, loaded = value, false
   254  	}
   255  	m.mu.Unlock()
   256  
   257  	return actual, loaded
   258  }
   259  
   260  // tryLoadOrStore atomically loads or stores a value if the entry is not
   261  // expunged.
   262  //
   263  // If the entry is expunged, tryLoadOrStore leaves the entry unchanged and
   264  // returns with ok==false.
   265  func (e *entry) tryLoadOrStore(i any) (actual any, loaded, ok bool) {
   266  	p := e.p.Load()
   267  	if p == expunged {
   268  		return nil, false, false
   269  	}
   270  	if p != nil {
   271  		return *p, true, true
   272  	}
   273  
   274  	// Copy the interface after the first load to make this method more amenable
   275  	// to escape analysis: if we hit the "load" path or the entry is expunged, we
   276  	// shouldn't bother heap-allocating.
   277  	ic := i
   278  	for {
   279  		if e.p.CompareAndSwap(nil, &ic) {
   280  			return i, false, true
   281  		}
   282  		p = e.p.Load()
   283  		if p == expunged {
   284  			return nil, false, false
   285  		}
   286  		if p != nil {
   287  			return *p, true, true
   288  		}
   289  	}
   290  }
   291  
   292  // LoadAndDelete deletes the value for a key, returning the previous value if any.
   293  // The loaded result reports whether the key was present.
   294  func (m *Map) LoadAndDelete(key any) (value any, loaded bool) {
   295  	read := m.loadReadOnly()
   296  	e, ok := read.m[key]
   297  	if !ok && read.amended {
   298  		m.mu.Lock()
   299  		read = m.loadReadOnly()
   300  		e, ok = read.m[key]
   301  		if !ok && read.amended {
   302  			e, ok = m.dirty[key]
   303  			delete(m.dirty, key)
   304  			// Regardless of whether the entry was present, record a miss: this key
   305  			// will take the slow path until the dirty map is promoted to the read
   306  			// map.
   307  			m.missLocked()
   308  		}
   309  		m.mu.Unlock()
   310  	}
   311  	if ok {
   312  		return e.delete()
   313  	}
   314  	return nil, false
   315  }
   316  
   317  // Delete deletes the value for a key.
   318  func (m *Map) Delete(key any) {
   319  	m.LoadAndDelete(key)
   320  }
   321  
   322  func (e *entry) delete() (value any, ok bool) {
   323  	for {
   324  		p := e.p.Load()
   325  		if p == nil || p == expunged {
   326  			return nil, false
   327  		}
   328  		if e.p.CompareAndSwap(p, nil) {
   329  			return *p, true
   330  		}
   331  	}
   332  }
   333  
   334  // trySwap swaps a value if the entry has not been expunged.
   335  //
   336  // If the entry is expunged, trySwap returns false and leaves the entry
   337  // unchanged.
   338  func (e *entry) trySwap(i *any) (*any, bool) {
   339  	for {
   340  		p := e.p.Load()
   341  		if p == expunged {
   342  			return nil, false
   343  		}
   344  		if e.p.CompareAndSwap(p, i) {
   345  			return p, true
   346  		}
   347  	}
   348  }
   349  
   350  // Swap swaps the value for a key and returns the previous value if any.
   351  // The loaded result reports whether the key was present.
   352  func (m *Map) Swap(key, value any) (previous any, loaded bool) {
   353  	read := m.loadReadOnly()
   354  	if e, ok := read.m[key]; ok {
   355  		if v, ok := e.trySwap(&value); ok {
   356  			if v == nil {
   357  				return nil, false
   358  			}
   359  			return *v, true
   360  		}
   361  	}
   362  
   363  	m.mu.Lock()
   364  	read = m.loadReadOnly()
   365  	if e, ok := read.m[key]; ok {
   366  		if e.unexpungeLocked() {
   367  			// The entry was previously expunged, which implies that there is a
   368  			// non-nil dirty map and this entry is not in it.
   369  			m.dirty[key] = e
   370  		}
   371  		if v := e.swapLocked(&value); v != nil {
   372  			loaded = true
   373  			previous = *v
   374  		}
   375  	} else if e, ok := m.dirty[key]; ok {
   376  		if v := e.swapLocked(&value); v != nil {
   377  			loaded = true
   378  			previous = *v
   379  		}
   380  	} else {
   381  		if !read.amended {
   382  			// We're adding the first new key to the dirty map.
   383  			// Make sure it is allocated and mark the read-only map as incomplete.
   384  			m.dirtyLocked()
   385  			m.read.Store(&readOnly{m: read.m, amended: true})
   386  		}
   387  		m.dirty[key] = newEntry(value)
   388  	}
   389  	m.mu.Unlock()
   390  	return previous, loaded
   391  }
   392  
   393  // CompareAndSwap swaps the old and new values for key
   394  // if the value stored in the map is equal to old.
   395  // The old value must be of a comparable type.
   396  func (m *Map) CompareAndSwap(key, old, new any) (swapped bool) {
   397  	read := m.loadReadOnly()
   398  	if e, ok := read.m[key]; ok {
   399  		return e.tryCompareAndSwap(old, new)
   400  	} else if !read.amended {
   401  		return false // No existing value for key.
   402  	}
   403  
   404  	m.mu.Lock()
   405  	defer m.mu.Unlock()
   406  	read = m.loadReadOnly()
   407  	swapped = false
   408  	if e, ok := read.m[key]; ok {
   409  		swapped = e.tryCompareAndSwap(old, new)
   410  	} else if e, ok := m.dirty[key]; ok {
   411  		swapped = e.tryCompareAndSwap(old, new)
   412  		// We needed to lock mu in order to load the entry for key,
   413  		// and the operation didn't change the set of keys in the map
   414  		// (so it would be made more efficient by promoting the dirty
   415  		// map to read-only).
   416  		// Count it as a miss so that we will eventually switch to the
   417  		// more efficient steady state.
   418  		m.missLocked()
   419  	}
   420  	return swapped
   421  }
   422  
   423  // CompareAndDelete deletes the entry for key if its value is equal to old.
   424  // The old value must be of a comparable type.
   425  //
   426  // If there is no current value for key in the map, CompareAndDelete
   427  // returns false (even if the old value is the nil interface value).
   428  func (m *Map) CompareAndDelete(key, old any) (deleted bool) {
   429  	read := m.loadReadOnly()
   430  	e, ok := read.m[key]
   431  	if !ok && read.amended {
   432  		m.mu.Lock()
   433  		read = m.loadReadOnly()
   434  		e, ok = read.m[key]
   435  		if !ok && read.amended {
   436  			e, ok = m.dirty[key]
   437  			// Don't delete key from m.dirty: we still need to do the “compare” part
   438  			// of the operation. The entry will eventually be expunged when the
   439  			// dirty map is promoted to the read map.
   440  			//
   441  			// Regardless of whether the entry was present, record a miss: this key
   442  			// will take the slow path until the dirty map is promoted to the read
   443  			// map.
   444  			m.missLocked()
   445  		}
   446  		m.mu.Unlock()
   447  	}
   448  	for ok {
   449  		p := e.p.Load()
   450  		if p == nil || p == expunged || *p != old {
   451  			return false
   452  		}
   453  		if e.p.CompareAndSwap(p, nil) {
   454  			return true
   455  		}
   456  	}
   457  	return false
   458  }
   459  
   460  // Range calls f sequentially for each key and value present in the map.
   461  // If f returns false, range stops the iteration.
   462  //
   463  // Range does not necessarily correspond to any consistent snapshot of the Map's
   464  // contents: no key will be visited more than once, but if the value for any key
   465  // is stored or deleted concurrently (including by f), Range may reflect any
   466  // mapping for that key from any point during the Range call. Range does not
   467  // block other methods on the receiver; even f itself may call any method on m.
   468  //
   469  // Range may be O(N) with the number of elements in the map even if f returns
   470  // false after a constant number of calls.
   471  func (m *Map) Range(f func(key, value any) bool) {
   472  	// We need to be able to iterate over all of the keys that were already
   473  	// present at the start of the call to Range.
   474  	// If read.amended is false, then read.m satisfies that property without
   475  	// requiring us to hold m.mu for a long time.
   476  	read := m.loadReadOnly()
   477  	if read.amended {
   478  		// m.dirty contains keys not in read.m. Fortunately, Range is already O(N)
   479  		// (assuming the caller does not break out early), so a call to Range
   480  		// amortizes an entire copy of the map: we can promote the dirty copy
   481  		// immediately!
   482  		m.mu.Lock()
   483  		read = m.loadReadOnly()
   484  		if read.amended {
   485  			read = readOnly{m: m.dirty}
   486  			copyRead := read
   487  			m.read.Store(&copyRead)
   488  			m.dirty = nil
   489  			m.misses = 0
   490  		}
   491  		m.mu.Unlock()
   492  	}
   493  
   494  	for k, e := range read.m {
   495  		v, ok := e.load()
   496  		if !ok {
   497  			continue
   498  		}
   499  		if !f(k, v) {
   500  			break
   501  		}
   502  	}
   503  }
   504  
   505  func (m *Map) missLocked() {
   506  	m.misses++
   507  	if m.misses < len(m.dirty) {
   508  		return
   509  	}
   510  	m.read.Store(&readOnly{m: m.dirty})
   511  	m.dirty = nil
   512  	m.misses = 0
   513  }
   514  
   515  func (m *Map) dirtyLocked() {
   516  	if m.dirty != nil {
   517  		return
   518  	}
   519  
   520  	read := m.loadReadOnly()
   521  	m.dirty = make(map[any]*entry, len(read.m))
   522  	for k, e := range read.m {
   523  		if !e.tryExpungeLocked() {
   524  			m.dirty[k] = e
   525  		}
   526  	}
   527  }
   528  
   529  func (e *entry) tryExpungeLocked() (isExpunged bool) {
   530  	p := e.p.Load()
   531  	for p == nil {
   532  		if e.p.CompareAndSwap(nil, expunged) {
   533  			return true
   534  		}
   535  		p = e.p.Load()
   536  	}
   537  	return p == expunged
   538  }
   539  

View as plain text