github.com/scottcagno/storage@v1.8.0/pkg/hashmap/chained/sharded.go (about)

     1  package chained
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"math/bits"
     7  	"runtime"
     8  	"sync"
     9  )
    10  
    11  type shard struct {
    12  	mu sync.RWMutex
    13  	hm *HashMap // chained
    14  }
    15  
    16  type ShardedHashMap struct {
    17  	mask   uint64
    18  	hash   hashFunc
    19  	shards []*shard
    20  }
    21  
    22  // NewShardedHashMap returns a new hashMap instantiated with the specified size or
    23  // the defaultMapSize, whichever is larger
    24  func NewShardedHashMap(size uint) *ShardedHashMap {
    25  	return newShardedHashMap(size, defaultHashFunc)
    26  }
    27  
    28  func newShardedHashMap(size uint, fn hashFunc) *ShardedHashMap {
    29  	shCount := alignShardCount(size)
    30  	if fn == nil {
    31  		fn = defaultHashFunc
    32  	}
    33  	shm := &ShardedHashMap{
    34  		mask:   shCount - 1,
    35  		hash:   fn,
    36  		shards: make([]*shard, shCount),
    37  	}
    38  	hmSize := initialMapShardSize(uint16(shCount))
    39  	log.Printf("new sharded hashmap with %d shards, each shard init with %d buckets\n", shCount, hmSize)
    40  	for i := range shm.shards {
    41  		shm.shards[i] = &shard{
    42  			hm: newHashMap(hmSize, fn),
    43  		}
    44  	}
    45  	return shm
    46  }
    47  
    48  func alignShardCount(size uint) uint64 {
    49  	count := uint(16)
    50  	for count < size {
    51  		count *= 2
    52  	}
    53  	return uint64(count)
    54  }
    55  
    56  func initialMapShardSize(x uint16) uint {
    57  	return uint(bits.Reverse16(x)) / 2
    58  }
    59  
    60  func (s *ShardedHashMap) getShard(key keyType) (uint64, uint64) {
    61  	// calculate the hashkey value
    62  	hashkey := s.hash(key)
    63  	// mask the hashkey to get the initial index
    64  	i := hashkey & s.mask
    65  	return i, hashkey
    66  }
    67  
    68  func (s *ShardedHashMap) Put(key keyType, val valType) (valType, bool) {
    69  	return s.insert(key, val)
    70  }
    71  
    72  func (s *ShardedHashMap) insert(key keyType, val valType) (valType, bool) {
    73  	buk, hashkey := s.getShard(key)
    74  	s.shards[buk].mu.Lock()
    75  	pv, ok := s.shards[buk].hm.insert(hashkey, key, val)
    76  	s.shards[buk].mu.Unlock()
    77  	return pv, ok
    78  }
    79  
    80  func (s *ShardedHashMap) Get(key keyType) (valType, bool) {
    81  	return s.lookup(key)
    82  }
    83  
    84  func (s *ShardedHashMap) lookup(key keyType) (valType, bool) {
    85  	buk, hashkey := s.getShard(key)
    86  	s.shards[buk].mu.RLock()
    87  	pv, ok := s.shards[buk].hm.lookup(hashkey, key)
    88  	s.shards[buk].mu.RUnlock()
    89  	return pv, ok
    90  }
    91  
    92  func (s *ShardedHashMap) Del(key keyType) (valType, bool) {
    93  	return s.delete(key)
    94  }
    95  
    96  func (s *ShardedHashMap) delete(key keyType) (valType, bool) {
    97  	buk, hashkey := s.getShard(key)
    98  	s.shards[buk].mu.Lock()
    99  	pv, ok := s.shards[buk].hm.delete(hashkey, key)
   100  	s.shards[buk].mu.Unlock()
   101  	return pv, ok
   102  }
   103  
   104  func (s *ShardedHashMap) Len() int {
   105  	var length int
   106  	for i := range s.shards {
   107  		s.shards[i].mu.Lock()
   108  		length += s.shards[i].hm.Len()
   109  		s.shards[i].mu.Unlock()
   110  	}
   111  	return length
   112  }
   113  
   114  func (s *ShardedHashMap) Range(it Iterator) {
   115  	for i := range s.shards {
   116  		s.shards[i].mu.Lock()
   117  		s.shards[i].hm.Range(it)
   118  		s.shards[i].mu.Unlock()
   119  	}
   120  }
   121  
   122  func (s *ShardedHashMap) Stats() {
   123  	for i := range s.shards {
   124  		s.shards[i].mu.Lock()
   125  		if pf := s.shards[i].hm.PercentFull(); pf > 0 {
   126  			fmt.Printf("shard %d, fill percent: %.4f\n", i, pf)
   127  		}
   128  		s.shards[i].mu.Unlock()
   129  	}
   130  }
   131  
   132  func (s *ShardedHashMap) Close() {
   133  	for i := range s.shards {
   134  		s.shards[i].mu.Lock()
   135  		destroy(s.shards[i].hm)
   136  		s.shards[i].mu.Unlock()
   137  	}
   138  	s.shards = nil
   139  	runtime.GC()
   140  }