github.com/saintwish/kv@v1.0.4/kvmap/kvmap.go (about)

     1  package kvmap
     2  
     3  import (
     4  	"fmt"
     5  	
     6  	"github.com/dolthub/maphash"
     7  )
     8  
     9  type Cache[K comparable, V any] struct {
    10  	shards []shard[K, V]
    11  	shardCount uint64
    12  	hash maphash.Hasher[K]
    13  
    14  	OnEvicted func(K, V) //function that's called when cached item is deleted by the system
    15  }
    16  
    17  func New[K comparable, V any](sc uint64) *Cache[K, V] {
    18  	cache := Cache[K, V] {}
    19  	cache.shards = make([]shard[K, V], sc)
    20  	cache.hash = maphash.NewHasher[K]()
    21  	cache.shardCount = sc
    22  
    23  	for i := 0; i < int(sc); i++ {
    24  		cache.shards[i] = newShard[K, V]()
    25  	}
    26  
    27  	return &cache
    28  }
    29  
    30  func (c *Cache[K, V]) getShardIndex(key K) uint64 {
    31  	sum := c.hash.Hash(key)
    32  
    33  	return sum % c.shardCount
    34  }
    35  
    36  func (c *Cache[K, V]) getShard(key K) shard[K, V] {
    37  	sum := c.hash.Hash(key)
    38  	return c.shards[sum%c.shardCount]
    39  }
    40  
    41  func (c *Cache[K, V]) SetOnEvicted(f func(K, V)) {
    42  	c.OnEvicted = f
    43  }
    44  
    45  func (c *Cache[K, V]) Set(key K, val V) {
    46  	shard := c.getShard(key)
    47  	shard.set(key, val)
    48  }
    49  
    50  func (c *Cache[K, V]) Get(key K) V {
    51  	shard := c.getShard(key)
    52  	return shard.get(key)
    53  }
    54  
    55  func (c *Cache[K, V]) Has(key K) bool {
    56  	shard := c.getShard(key)
    57  	shard.RLock()
    58  	defer shard.RUnlock()
    59  
    60  	if _, ok := shard.Map[key]; ok {
    61  		return true
    62  	}
    63  
    64  	return false
    65  }
    66  
    67  func (c *Cache[K, V]) Add(key K, val V) (err error) {
    68  	shard := c.getShard(key)
    69  	shard.Lock()
    70  
    71  	if _, ok := shard.Map[key]; ok {
    72  		err = fmt.Errorf("kvmap: Data already exists with given key %T", key)
    73  		shard.Unlock()
    74  		return
    75  	}
    76  	
    77  	c.Set(key, val)
    78  
    79  	shard.Unlock()
    80  	return nil
    81  }
    82  
    83  func (c *Cache[K, V]) Update(key K, val V) (err error) {
    84  	shard := c.getShard(key)
    85  	shard.Lock()
    86  
    87  	if _, ok := shard.Map[key]; !ok {
    88  		return fmt.Errorf("kvmap: Data doesn't exists with given key %T", key)
    89  		shard.Unlock()
    90  		return
    91  	}
    92  
    93  	c.Set(key, val)
    94  
    95  	shard.Unlock()
    96  	return nil
    97  }
    98  
    99  func (c *Cache[K, V]) Delete(key K) {
   100  	if c.Has(key) {
   101  		shard := c.getShard(key)
   102  		delete(shard.Map, key)
   103  	}
   104  }
   105  
   106  func (c *Cache[K, V]) Flush() {
   107  	for i := 0; i < len(c.shards); i++ {
   108  		shard := c.shards[i]
   109  
   110  		for k,v := range shard.Map {
   111  			c.OnEvicted(k, v)
   112  			delete(shard.Map, k)
   113  		}
   114  	}
   115  }