github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/interlock/concurrent_map.go (about)

     1  // Copyright 2020 WHTCORPS INC, 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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package interlock
    15  
    16  import (
    17  	"sync"
    18  )
    19  
    20  // ShardCount controls the shard maps within the concurrent map
    21  var ShardCount = 320
    22  
    23  // A "thread" safe map of type string:Anything.
    24  // To avoid dagger bottlenecks this map is dived to several (ShardCount) map shards.
    25  type concurrentMap []*concurrentMapShared
    26  
    27  // A "thread" safe string to anything map.
    28  type concurrentMapShared struct {
    29  	items        map[uint64]*entry
    30  	sync.RWMutex // Read Write mutex, guards access to internal map.
    31  }
    32  
    33  // newConcurrentMap creates a new concurrent map.
    34  func newConcurrentMap() concurrentMap {
    35  	m := make(concurrentMap, ShardCount)
    36  	for i := 0; i < ShardCount; i++ {
    37  		m[i] = &concurrentMapShared{items: make(map[uint64]*entry)}
    38  	}
    39  	return m
    40  }
    41  
    42  // getShard returns shard under given key
    43  func (m concurrentMap) getShard(hashKey uint64) *concurrentMapShared {
    44  	return m[hashKey%uint64(ShardCount)]
    45  }
    46  
    47  // Insert inserts a value in a shard safely
    48  func (m concurrentMap) Insert(key uint64, value *entry) {
    49  	shard := m.getShard(key)
    50  	shard.Lock()
    51  	v, ok := shard.items[key]
    52  	if !ok {
    53  		shard.items[key] = value
    54  	} else {
    55  		value.next = v
    56  		shard.items[key] = value
    57  	}
    58  	shard.Unlock()
    59  	return
    60  }
    61  
    62  // UpsertCb : Callback to return new element to be inserted into the map
    63  // It is called while dagger is held, therefore it MUST NOT
    64  // try to access other keys in same map, as it can lead to deadlock since
    65  // Go sync.RWLock is not reentrant
    66  type UpsertCb func(exist bool, valueInMap, newValue *entry) *entry
    67  
    68  // Upsert: Insert or UFIDelate - uFIDelates existing element or inserts a new one using UpsertCb
    69  func (m concurrentMap) Upsert(key uint64, value *entry, cb UpsertCb) (res *entry) {
    70  	shard := m.getShard(key)
    71  	shard.Lock()
    72  	v, ok := shard.items[key]
    73  	res = cb(ok, v, value)
    74  	shard.items[key] = res
    75  	shard.Unlock()
    76  	return res
    77  }
    78  
    79  // Get retrieves an element from map under given key.
    80  // Note that in hash joins, reading proceeds after all writes, so we ignore RLock() here.
    81  // Otherwise, we should use RLock() for concurrent reads and writes.
    82  func (m concurrentMap) Get(key uint64) (*entry, bool) {
    83  	// Get shard
    84  	shard := m.getShard(key)
    85  	// shard.RLock()
    86  	// Get item from shard.
    87  	val, ok := shard.items[key]
    88  	// shard.RUnlock()
    89  	return val, ok
    90  }
    91  
    92  // IterCb :Iterator callback,called for every key,value found in
    93  // maps. RLock is held for all calls for a given shard
    94  // therefore callback sess consistent view of a shard,
    95  // but not across the shards
    96  type IterCb func(key uint64, e *entry)
    97  
    98  // IterCb iterates the map using a callback, cheapest way to read
    99  // all elements in a map.
   100  func (m concurrentMap) IterCb(fn IterCb) {
   101  	for idx := range m {
   102  		shard := (m)[idx]
   103  		shard.RLock()
   104  		for key, value := range shard.items {
   105  			fn(key, value)
   106  		}
   107  		shard.RUnlock()
   108  	}
   109  }