github.com/uber/kraken@v0.1.4/utils/lockermap/map.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package lockermap
    15  
    16  import (
    17  	"sync"
    18  
    19  	"golang.org/x/sync/syncmap"
    20  )
    21  
    22  // Map is a two-level locking map which synchronizes access to the map
    23  // in addition to synchronizing access to the values within the map. Useful
    24  // for mutating values in-place.
    25  //
    26  // The zero Map is valid and empty.
    27  type Map struct {
    28  	m syncmap.Map
    29  }
    30  
    31  // Load looks up the value of key k and executes f under the protection of
    32  // said value's lock. While f executes, it is guaranteed that k will not
    33  // be deleted from the map. Returns false if k was not found.
    34  func (m *Map) Load(k interface{}, f func(sync.Locker)) bool {
    35  	v, ok := m.m.Load(k)
    36  	if !ok {
    37  		return false
    38  	}
    39  
    40  	l := v.(sync.Locker)
    41  	l.Lock()
    42  	defer l.Unlock()
    43  
    44  	// Now that we have the value lock, make sure k was not deleted.
    45  	if _, ok = m.m.Load(k); !ok {
    46  		return false
    47  	}
    48  
    49  	f(l)
    50  
    51  	return true
    52  }
    53  
    54  // TryStore stores the given key / value pair within the Map. Returns
    55  // false if the key is already present.
    56  func (m *Map) TryStore(k interface{}, v sync.Locker) bool {
    57  	_, loaded := m.m.LoadOrStore(k, v)
    58  	return !loaded
    59  }
    60  
    61  // Delete deletes the given key from the Map.
    62  func (m *Map) Delete(k interface{}) {
    63  	v, ok := m.m.Load(k)
    64  	if !ok {
    65  		return
    66  	}
    67  
    68  	l := v.(sync.Locker)
    69  	l.Lock()
    70  	defer l.Unlock()
    71  
    72  	m.m.Delete(k)
    73  }
    74  
    75  // Range interates over the Map and execs f until f returns false.
    76  func (m *Map) Range(f func(k interface{}, v sync.Locker) bool) {
    77  	m.m.Range(func(k, v interface{}) bool {
    78  		l := v.(sync.Locker)
    79  		l.Lock()
    80  		defer l.Unlock()
    81  
    82  		if _, ok := m.m.Load(k); !ok {
    83  			// Skip this entry because it has been deleted.
    84  			return true
    85  		}
    86  
    87  		return f(k, l)
    88  	})
    89  }