github.com/etecs-ru/ristretto@v0.9.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 struct {
    26  	expiration time.Time
    27  	value      interface{}
    28  	key        uint64
    29  	conflict   uint64
    30  }
    31  
    32  const numShards uint64 = 256
    33  
    34  type (
    35  	updateFn   func(prev, cur interface{}) bool
    36  	shardedMap struct {
    37  		expiryMap    *expirationMap
    38  		shouldUpdate func(prev, cur interface{}) bool
    39  		shards       []*lockedMap
    40  	}
    41  )
    42  
    43  // newShardedMap is safe for concurrent usage.
    44  func newShardedMap(fn updateFn) *shardedMap {
    45  	sm := &shardedMap{
    46  		shards:    make([]*lockedMap, int(numShards)),
    47  		expiryMap: newExpirationMap(),
    48  	}
    49  	if fn == nil {
    50  		fn = func(prev, cur interface{}) bool {
    51  			return true
    52  		}
    53  	}
    54  	for i := range sm.shards {
    55  		sm.shards[i] = newLockedMap(fn, sm.expiryMap)
    56  	}
    57  	return sm
    58  }
    59  
    60  func (sm *shardedMap) Get(key, conflict uint64) (interface{}, bool) {
    61  	return sm.shards[key%numShards].get(key, conflict)
    62  }
    63  
    64  func (sm *shardedMap) Expiration(key uint64) time.Time {
    65  	return sm.shards[key%numShards].Expiration(key)
    66  }
    67  
    68  func (sm *shardedMap) Set(i *Item) {
    69  	if i == nil {
    70  		// If item is nil make this Set a no-op.
    71  		return
    72  	}
    73  
    74  	sm.shards[i.Key%numShards].Set(i)
    75  }
    76  
    77  func (sm *shardedMap) Del(key, conflict uint64) (uint64, interface{}) {
    78  	return sm.shards[key%numShards].Del(key, conflict)
    79  }
    80  
    81  func (sm *shardedMap) Update(newItem *Item) (interface{}, bool) {
    82  	return sm.shards[newItem.Key%numShards].Update(newItem)
    83  }
    84  
    85  func (sm *shardedMap) Cleanup(policy *lfuPolicy, onEvict itemCallback) {
    86  	sm.expiryMap.cleanup(sm, policy, onEvict)
    87  }
    88  
    89  func (sm *shardedMap) Clear(onEvict itemCallback) {
    90  	for i := uint64(0); i < numShards; i++ {
    91  		sm.shards[i].Clear(onEvict)
    92  	}
    93  }
    94  
    95  type lockedMap struct {
    96  	data         map[uint64]storeItem
    97  	em           *expirationMap
    98  	shouldUpdate updateFn
    99  	sync.RWMutex
   100  }
   101  
   102  func newLockedMap(fn updateFn, em *expirationMap) *lockedMap {
   103  	return &lockedMap{
   104  		data:         make(map[uint64]storeItem),
   105  		em:           em,
   106  		shouldUpdate: fn,
   107  	}
   108  }
   109  
   110  func (m *lockedMap) get(key, conflict uint64) (interface{}, bool) {
   111  	m.RLock()
   112  	item, ok := m.data[key]
   113  	m.RUnlock()
   114  	if !ok {
   115  		return nil, false
   116  	}
   117  	if conflict != 0 && (conflict != item.conflict) {
   118  		return nil, false
   119  	}
   120  
   121  	// Handle expired items.
   122  	if !item.expiration.IsZero() && time.Now().After(item.expiration) {
   123  		return nil, false
   124  	}
   125  	return item.value, true
   126  }
   127  
   128  func (m *lockedMap) Expiration(key uint64) time.Time {
   129  	m.RLock()
   130  	defer m.RUnlock()
   131  	return m.data[key].expiration
   132  }
   133  
   134  func (m *lockedMap) Set(i *Item) {
   135  	if i == nil {
   136  		// If the item is nil make this Set a no-op.
   137  		return
   138  	}
   139  
   140  	m.Lock()
   141  	defer m.Unlock()
   142  	item, ok := m.data[i.Key]
   143  
   144  	if ok {
   145  		// The item existed already. We need to check the conflict key and reject the
   146  		// update if they do not match. Only after that the expiration map is updated.
   147  		if i.Conflict != 0 && (i.Conflict != item.conflict) {
   148  			return
   149  		}
   150  		if !m.shouldUpdate(item.value, i.Value) {
   151  			return
   152  		}
   153  		m.em.update(i.Key, i.Conflict, item.expiration, i.Expiration)
   154  	} else {
   155  		// The value is not in the map already. There's no need to return anything.
   156  		// Simply add the expiration map.
   157  		m.em.add(i.Key, i.Conflict, i.Expiration)
   158  	}
   159  
   160  	m.data[i.Key] = storeItem{
   161  		key:        i.Key,
   162  		conflict:   i.Conflict,
   163  		value:      i.Value,
   164  		expiration: i.Expiration,
   165  	}
   166  }
   167  
   168  func (m *lockedMap) Del(key, conflict uint64) (uint64, interface{}) {
   169  	m.Lock()
   170  	item, ok := m.data[key]
   171  	if !ok {
   172  		m.Unlock()
   173  		return 0, nil
   174  	}
   175  	if conflict != 0 && (conflict != item.conflict) {
   176  		m.Unlock()
   177  		return 0, nil
   178  	}
   179  
   180  	if !item.expiration.IsZero() {
   181  		m.em.del(key, item.expiration)
   182  	}
   183  
   184  	delete(m.data, key)
   185  	m.Unlock()
   186  	return item.conflict, item.value
   187  }
   188  
   189  func (m *lockedMap) Update(newItem *Item) (interface{}, bool) {
   190  	m.Lock()
   191  	defer m.Unlock()
   192  
   193  	item, ok := m.data[newItem.Key]
   194  	if !ok {
   195  		return nil, false
   196  	}
   197  	if newItem.Conflict != 0 && (newItem.Conflict != item.conflict) {
   198  		return nil, false
   199  	}
   200  	if !m.shouldUpdate(item.value, newItem.Value) {
   201  		return item.value, false
   202  	}
   203  
   204  	m.em.update(newItem.Key, newItem.Conflict, item.expiration, newItem.Expiration)
   205  	m.data[newItem.Key] = storeItem{
   206  		key:        newItem.Key,
   207  		conflict:   newItem.Conflict,
   208  		value:      newItem.Value,
   209  		expiration: newItem.Expiration,
   210  	}
   211  	return item.value, true
   212  }
   213  
   214  func (m *lockedMap) Clear(onEvict itemCallback) {
   215  	m.Lock()
   216  	i := &Item{}
   217  	if onEvict != nil {
   218  		for _, si := range m.data {
   219  			i.Key = si.key
   220  			i.Conflict = si.conflict
   221  			i.Value = si.value
   222  			onEvict(i)
   223  		}
   224  	}
   225  	m.data = make(map[uint64]storeItem)
   226  	m.Unlock()
   227  }