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

     1  package kv2
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/saintwish/kv/swiss"
     7  	"github.com/saintwish/kv/stack"
     8  )
     9  
    10  type item[V any] struct {
    11  	Object V
    12  	Index int
    13  }
    14  
    15  //used internally
    16  type shardCapacity[K comparable, V any] struct {
    17  	Map *swiss.Map[K, item[V]]
    18  	Stack *stack.Stack[K]
    19  	*sync.RWMutex //mutex
    20  }
    21  
    22  func newShardCapacity[K comparable, V any](size uint64, count uint64) shardCapacity[K, V] {
    23  	return shardCapacity[K, V] {
    24  		Map: swiss.NewMap[K, item[V]]( uint32(size/count) ),
    25  		Stack: stack.New[K](int(size/count)),
    26  		RWMutex: &sync.RWMutex{},
    27  	}
    28  }
    29  
    30  func (m shardCapacity[K, V]) has(key K) bool {
    31  	m.RLock()
    32  
    33  	ok := m.Map.Has(key)
    34  
    35  	m.RUnlock()
    36  
    37  	return ok
    38  }
    39  
    40  func (m shardCapacity[K, V]) getHasRenew(key K) (V, bool) {
    41  	m.Lock()
    42  
    43  	ok,val := m.Map.GetHas(key)
    44  	val.Index = m.Stack.MoveToBack(val.Index)
    45  	m.Map.Set(key, val)
    46  
    47  	m.Unlock()
    48  
    49  	return val.Object, ok
    50  }
    51  
    52  func (m shardCapacity[K, V]) getHas(key K) (V, bool) {
    53  	m.RLock()
    54  
    55  	ok,val := m.Map.GetHas(key)
    56  
    57  	m.RUnlock()
    58  
    59  	return val.Object, ok
    60  }
    61  
    62  func (m shardCapacity[K, V]) getRenew(key K) V {
    63  	m.Lock()
    64  
    65  	val := m.Map.Get(key)
    66  	val.Index = m.Stack.MoveToBack(val.Index)
    67  	m.Map.Set(key, val)
    68  
    69  	m.Unlock()
    70  
    71  	return val.Object
    72  }
    73  
    74  func (m shardCapacity[K, V]) get(key K) V {
    75  	m.RLock()
    76  
    77  	val := m.Map.Get(key)
    78  
    79  	m.RUnlock()
    80  
    81  	return val.Object
    82  }
    83  
    84  /*--------
    85  	Other functions
    86  ----------*/
    87  func (m shardCapacity[K, V]) set(key K, val V, callback func(K, V)) {
    88  	itm := item[V]{
    89  		Object: val,
    90  	}
    91  
    92  	m.Lock()
    93  
    94  	
    95  	if m.Map.Capacity() > 0 {
    96  		itm.Index = m.Stack.Push(key)
    97  		m.Map.Set(key, itm)
    98  	}
    99  
   100  	if m.Map.Capacity() == 0 {
   101  		_, oldKey := m.Stack.Pop()
   102  		_,v := m.Map.Delete(oldKey)
   103  		callback(oldKey, v.Object)
   104  
   105  		itm.Index = m.Stack.Push(key)
   106  		m.Map.Set(key, itm)
   107  	}
   108  
   109  	m.Unlock()
   110  }
   111  
   112  func (m shardCapacity[K, V]) update(key K, val V) {
   113  	m.Lock()
   114  
   115  	if ok,v := m.Map.GetHas(key); ok {
   116  		v.Object = val
   117  		v.Index = m.Stack.MoveToBack(v.Index)
   118  		m.Map.Set(key, v)
   119  	}
   120  
   121  	m.Unlock()
   122  }
   123  
   124  func (m shardCapacity[K, V]) delete(key K) bool {
   125  	m.Lock()
   126  
   127  	ok, val := m.Map.Delete(key)
   128  	m.Stack.Remove(val.Index)
   129  
   130  	m.Unlock()
   131  
   132  	return ok
   133  }
   134  
   135  func (m shardCapacity[K, V]) deleteCallback(key K, callback func(K, V)) bool {
   136  	m.Lock()
   137  
   138  	ok, val := m.Map.Delete(key)
   139  	m.Stack.Remove(val.Index)
   140  	callback(key, val.Object)
   141  
   142  	m.Unlock()
   143  
   144  	return ok
   145  }
   146  
   147  func (m shardCapacity[K, V]) clear() {
   148  	m.Map.Clear()
   149  }
   150  
   151  func (m shardCapacity[K, V]) flush(callback func(K, V)) {
   152  	m.Lock()
   153  
   154  	m.Stack.Clear()
   155  	m.Map.Iter(func(key K, val item[V]) (stop bool) {
   156  		callback(key, val.Object)
   157  		m.Map.Delete(key)
   158  		
   159  		if stop {
   160  			m.Unlock()
   161  			return
   162  		}
   163  
   164  		return
   165  	})
   166  
   167  	m.Unlock()
   168  }