github.com/pingcap/badger@v1.5.1-0.20230103063557-828f39b09b6d/cache/store.go (about)

     1  /*
     2   * Copyright 2019 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package cache
    18  
    19  import (
    20  	"sync"
    21  	"sync/atomic"
    22  	"unsafe"
    23  )
    24  
    25  const numShards = 256
    26  
    27  // item is passed to setBuf so items can eventually be added to the cache
    28  type item struct {
    29  	sync.Mutex
    30  	dead  bool
    31  	key   uint64
    32  	value atomic.Value
    33  }
    34  
    35  func (i *item) del(s *store) bool {
    36  	if i.dead {
    37  		return false
    38  	}
    39  
    40  	i.dead = true
    41  	s.Del(i.key)
    42  	return true
    43  }
    44  
    45  // TODO: implement a more powerful concurrent hash, such as cuckoo hash.
    46  type shardMap struct {
    47  	sync.RWMutex
    48  	data map[uint64]*item
    49  }
    50  
    51  type shard struct {
    52  	shardMap
    53  	pad [64 - unsafe.Sizeof(shardMap{})%64]byte
    54  }
    55  
    56  type store struct {
    57  	shards [numShards]shard
    58  }
    59  
    60  func newStore() *store {
    61  	s := new(store)
    62  	for i := range s.shards {
    63  		s.shards[i].data = make(map[uint64]*item)
    64  	}
    65  	return s
    66  }
    67  
    68  func (s *store) GetValue(key uint64) (interface{}, bool) {
    69  	i, ok := s.Get(key)
    70  	if !ok {
    71  		return nil, false
    72  	}
    73  	return i.value.Load(), true
    74  }
    75  
    76  func (s *store) Get(key uint64) (*item, bool) {
    77  	m := &s.shards[key%numShards]
    78  	m.RLock()
    79  	i, ok := m.data[key]
    80  	m.RUnlock()
    81  	return i, ok
    82  }
    83  
    84  func (s *store) GetOrNew(key uint64) *item {
    85  	m := &s.shards[key%numShards]
    86  	m.RLock()
    87  	if i, ok := m.data[key]; ok {
    88  		m.RUnlock()
    89  		return i
    90  	}
    91  	m.RUnlock()
    92  	m.Lock()
    93  	if i, ok := m.data[key]; ok {
    94  		m.Unlock()
    95  		return i
    96  	}
    97  	i := &item{key: key}
    98  	m.data[key] = i
    99  	m.Unlock()
   100  	return i
   101  }
   102  
   103  func (s *store) Set(key uint64, value *item) {
   104  }
   105  
   106  func (s *store) Del(key uint64) *item {
   107  	m := &s.shards[key%numShards]
   108  	m.Lock()
   109  	i := m.data[key]
   110  	delete(m.data, key)
   111  	m.Unlock()
   112  	return i
   113  }
   114  
   115  func (s *store) Clear() {
   116  	for i := uint64(0); i < numShards; i++ {
   117  		s.shards[i].Lock()
   118  		s.shards[i].data = make(map[uint64]*item)
   119  		s.shards[i].Unlock()
   120  	}
   121  }