github.com/fiatjaf/generic-ristretto@v0.0.1/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 ristretto
    18  
    19  import (
    20  	"sync"
    21  	"time"
    22  )
    23  
    24  // TODO: Do we need this to be a separate struct from Item?
    25  type storeItem[V any] struct {
    26  	key        uint64
    27  	conflict   uint64
    28  	value      V
    29  	expiration time.Time
    30  }
    31  
    32  // store is the interface fulfilled by all hash map implementations in this
    33  // file. Some hash map implementations are better suited for certain data
    34  // distributions than others, so this allows us to abstract that out for use
    35  // in Ristretto.
    36  //
    37  // Every store is safe for concurrent usage.
    38  type store[V any] interface {
    39  	// Get returns the value associated with the key parameter.
    40  	Get(uint64, uint64) (V, bool)
    41  	// Expiration returns the expiration time for this key.
    42  	Expiration(uint64) time.Time
    43  	// Set adds the key-value pair to the Map or updates the value if it's
    44  	// already present. The key-value pair is passed as a pointer to an
    45  	// item object.
    46  	Set(*Item[V])
    47  	// Del deletes the key-value pair from the Map.
    48  	Del(uint64, uint64) (uint64, V)
    49  	// Update attempxsts to update the key with a new value and returns true if
    50  	// successful.
    51  	Update(*Item[V]) (V, bool)
    52  	// Cleanup removes items that have an expired TTL.
    53  	Cleanup(policy policy[V], onEvict func(item *Item[V]))
    54  	// Clear clears all contents of the store.
    55  	Clear(onEvict func(item *Item[V]))
    56  }
    57  
    58  // newStore returns the default store implementation.
    59  func newStore[V any]() store[V] {
    60  	return newShardedMap[V]()
    61  }
    62  
    63  const numShards uint64 = 256
    64  
    65  type shardedMap[V any] struct {
    66  	shards    []*lockedMap[V]
    67  	expiryMap *expirationMap[V]
    68  }
    69  
    70  func newShardedMap[V any]() *shardedMap[V] {
    71  	sm := &shardedMap[V]{
    72  		shards:    make([]*lockedMap[V], int(numShards)),
    73  		expiryMap: newExpirationMap[V](),
    74  	}
    75  	for i := range sm.shards {
    76  		sm.shards[i] = newLockedMap[V](sm.expiryMap)
    77  	}
    78  	return sm
    79  }
    80  
    81  func (sm *shardedMap[V]) Get(key, conflict uint64) (V, bool) {
    82  	return sm.shards[key%numShards].get(key, conflict)
    83  }
    84  
    85  func (sm *shardedMap[V]) Expiration(key uint64) time.Time {
    86  	return sm.shards[key%numShards].Expiration(key)
    87  }
    88  
    89  func (sm *shardedMap[V]) Set(i *Item[V]) {
    90  	if i == nil {
    91  		// If item is nil make this Set a no-op.
    92  		return
    93  	}
    94  
    95  	sm.shards[i.Key%numShards].Set(i)
    96  }
    97  
    98  func (sm *shardedMap[V]) Del(key, conflict uint64) (uint64, V) {
    99  	return sm.shards[key%numShards].Del(key, conflict)
   100  }
   101  
   102  func (sm *shardedMap[V]) Update(newItem *Item[V]) (V, bool) {
   103  	return sm.shards[newItem.Key%numShards].Update(newItem)
   104  }
   105  
   106  func (sm *shardedMap[V]) Cleanup(policy policy[V], onEvict func(item *Item[V])) {
   107  	sm.expiryMap.cleanup(sm, policy, onEvict)
   108  }
   109  
   110  func (sm *shardedMap[V]) Clear(onEvict func(item *Item[V])) {
   111  	for i := uint64(0); i < numShards; i++ {
   112  		sm.shards[i].Clear(onEvict)
   113  	}
   114  }
   115  
   116  type lockedMap[V any] struct {
   117  	sync.RWMutex
   118  	data map[uint64]storeItem[V]
   119  	em   *expirationMap[V]
   120  }
   121  
   122  func newLockedMap[V any](em *expirationMap[V]) *lockedMap[V] {
   123  	return &lockedMap[V]{
   124  		data: make(map[uint64]storeItem[V]),
   125  		em:   em,
   126  	}
   127  }
   128  
   129  func (m *lockedMap[V]) get(key, conflict uint64) (V, bool) {
   130  	m.RLock()
   131  	item, ok := m.data[key]
   132  	m.RUnlock()
   133  	if !ok {
   134  		return zeroValue[V](), false
   135  	}
   136  	if conflict != 0 && (conflict != item.conflict) {
   137  		return zeroValue[V](), false
   138  	}
   139  
   140  	// Handle expired items.
   141  	if !item.expiration.IsZero() && time.Now().After(item.expiration) {
   142  		return zeroValue[V](), false
   143  	}
   144  	return item.value, true
   145  }
   146  
   147  func (m *lockedMap[V]) Expiration(key uint64) time.Time {
   148  	m.RLock()
   149  	defer m.RUnlock()
   150  	return m.data[key].expiration
   151  }
   152  
   153  func (m *lockedMap[V]) Set(i *Item[V]) {
   154  	if i == nil {
   155  		// If the item is nil make this Set a no-op.
   156  		return
   157  	}
   158  
   159  	m.Lock()
   160  	defer m.Unlock()
   161  	item, ok := m.data[i.Key]
   162  
   163  	if ok {
   164  		// The item existed already. We need to check the conflict key and reject the
   165  		// update if they do not match. Only after that the expiration map is updated.
   166  		if i.Conflict != 0 && (i.Conflict != item.conflict) {
   167  			return
   168  		}
   169  		m.em.update(i.Key, i.Conflict, item.expiration, i.Expiration)
   170  	} else {
   171  		// The value is not in the map already. There's no need to return anything.
   172  		// Simply add the expiration map.
   173  		m.em.add(i.Key, i.Conflict, i.Expiration)
   174  	}
   175  
   176  	m.data[i.Key] = storeItem[V]{
   177  		key:        i.Key,
   178  		conflict:   i.Conflict,
   179  		value:      i.Value,
   180  		expiration: i.Expiration,
   181  	}
   182  }
   183  
   184  func (m *lockedMap[V]) Del(key, conflict uint64) (uint64, V) {
   185  	m.Lock()
   186  	item, ok := m.data[key]
   187  	if !ok {
   188  		m.Unlock()
   189  		return 0, zeroValue[V]()
   190  	}
   191  	if conflict != 0 && (conflict != item.conflict) {
   192  		m.Unlock()
   193  		return 0, zeroValue[V]()
   194  	}
   195  
   196  	if !item.expiration.IsZero() {
   197  		m.em.del(key, item.expiration)
   198  	}
   199  
   200  	delete(m.data, key)
   201  	m.Unlock()
   202  	return item.conflict, item.value
   203  }
   204  
   205  func (m *lockedMap[V]) Update(newItem *Item[V]) (V, bool) {
   206  	m.Lock()
   207  	item, ok := m.data[newItem.Key]
   208  	if !ok {
   209  		m.Unlock()
   210  		return zeroValue[V](), false
   211  	}
   212  	if newItem.Conflict != 0 && (newItem.Conflict != item.conflict) {
   213  		m.Unlock()
   214  		return zeroValue[V](), false
   215  	}
   216  
   217  	m.em.update(newItem.Key, newItem.Conflict, item.expiration, newItem.Expiration)
   218  	m.data[newItem.Key] = storeItem[V]{
   219  		key:        newItem.Key,
   220  		conflict:   newItem.Conflict,
   221  		value:      newItem.Value,
   222  		expiration: newItem.Expiration,
   223  	}
   224  
   225  	m.Unlock()
   226  	return item.value, true
   227  }
   228  
   229  func (m *lockedMap[V]) Clear(onEvict func(item *Item[V])) {
   230  	m.Lock()
   231  	i := &Item[V]{}
   232  	if onEvict != nil {
   233  		for _, si := range m.data {
   234  			i.Key = si.key
   235  			i.Conflict = si.conflict
   236  			i.Value = si.value
   237  			onEvict(i)
   238  		}
   239  	}
   240  	m.data = make(map[uint64]storeItem[V])
   241  	m.Unlock()
   242  }