store.go

Documentation: github.com/dgraph-io/ristretto

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