...
Run Format

Source file src/sync/map_reference_test.go

Documentation: sync

  // Copyright 2016 The Go Authors. All rights reserved.
  // Use of this source code is governed by a BSD-style
  // license that can be found in the LICENSE file.
  
  package sync_test
  
  import (
  	"sync"
  	"sync/atomic"
  )
  
  // This file contains reference map implementations for unit-tests.
  
  // mapInterface is the interface Map implements.
  type mapInterface interface {
  	Load(interface{}) (interface{}, bool)
  	Store(key, value interface{})
  	LoadOrStore(key, value interface{}) (actual interface{}, loaded bool)
  	Delete(interface{})
  	Range(func(key, value interface{}) (shouldContinue bool))
  }
  
  // RWMutexMap is an implementation of mapInterface using a sync.RWMutex.
  type RWMutexMap struct {
  	mu    sync.RWMutex
  	dirty map[interface{}]interface{}
  }
  
  func (m *RWMutexMap) Load(key interface{}) (value interface{}, ok bool) {
  	m.mu.RLock()
  	value, ok = m.dirty[key]
  	m.mu.RUnlock()
  	return
  }
  
  func (m *RWMutexMap) Store(key, value interface{}) {
  	m.mu.Lock()
  	if m.dirty == nil {
  		m.dirty = make(map[interface{}]interface{})
  	}
  	m.dirty[key] = value
  	m.mu.Unlock()
  }
  
  func (m *RWMutexMap) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) {
  	m.mu.Lock()
  	actual, loaded = m.dirty[key]
  	if !loaded {
  		actual = value
  		if m.dirty == nil {
  			m.dirty = make(map[interface{}]interface{})
  		}
  		m.dirty[key] = value
  	}
  	m.mu.Unlock()
  	return actual, loaded
  }
  
  func (m *RWMutexMap) Delete(key interface{}) {
  	m.mu.Lock()
  	delete(m.dirty, key)
  	m.mu.Unlock()
  }
  
  func (m *RWMutexMap) Range(f func(key, value interface{}) (shouldContinue bool)) {
  	m.mu.RLock()
  	keys := make([]interface{}, 0, len(m.dirty))
  	for k := range m.dirty {
  		keys = append(keys, k)
  	}
  	m.mu.RUnlock()
  
  	for _, k := range keys {
  		v, ok := m.Load(k)
  		if !ok {
  			continue
  		}
  		if !f(k, v) {
  			break
  		}
  	}
  }
  
  // DeepCopyMap is an implementation of mapInterface using a Mutex and
  // atomic.Value.  It makes deep copies of the map on every write to avoid
  // acquiring the Mutex in Load.
  type DeepCopyMap struct {
  	mu    sync.Mutex
  	clean atomic.Value
  }
  
  func (m *DeepCopyMap) Load(key interface{}) (value interface{}, ok bool) {
  	clean, _ := m.clean.Load().(map[interface{}]interface{})
  	value, ok = clean[key]
  	return value, ok
  }
  
  func (m *DeepCopyMap) Store(key, value interface{}) {
  	m.mu.Lock()
  	dirty := m.dirty()
  	dirty[key] = value
  	m.clean.Store(dirty)
  	m.mu.Unlock()
  }
  
  func (m *DeepCopyMap) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) {
  	clean, _ := m.clean.Load().(map[interface{}]interface{})
  	actual, loaded = clean[key]
  	if loaded {
  		return actual, loaded
  	}
  
  	m.mu.Lock()
  	// Reload clean in case it changed while we were waiting on m.mu.
  	clean, _ = m.clean.Load().(map[interface{}]interface{})
  	actual, loaded = clean[key]
  	if !loaded {
  		dirty := m.dirty()
  		dirty[key] = value
  		actual = value
  		m.clean.Store(dirty)
  	}
  	m.mu.Unlock()
  	return actual, loaded
  }
  
  func (m *DeepCopyMap) Delete(key interface{}) {
  	m.mu.Lock()
  	dirty := m.dirty()
  	delete(dirty, key)
  	m.clean.Store(dirty)
  	m.mu.Unlock()
  }
  
  func (m *DeepCopyMap) Range(f func(key, value interface{}) (shouldContinue bool)) {
  	clean, _ := m.clean.Load().(map[interface{}]interface{})
  	for k, v := range clean {
  		if !f(k, v) {
  			break
  		}
  	}
  }
  
  func (m *DeepCopyMap) dirty() map[interface{}]interface{} {
  	clean, _ := m.clean.Load().(map[interface{}]interface{})
  	dirty := make(map[interface{}]interface{}, len(clean)+1)
  	for k, v := range clean {
  		dirty[k] = v
  	}
  	return dirty
  }
  

View as plain text