github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/store/util/sizecache/size_cache.go (about) 1 // Copyright 2019 Dolthub, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 // 15 // This file incorporates work covered by the following copyright and 16 // permission notice: 17 // 18 // Copyright 2016 Attic Labs, Inc. All rights reserved. 19 // Licensed under the Apache License, version 2.0: 20 // http://www.apache.org/licenses/LICENSE-2.0 21 22 package sizecache 23 24 // SizeCache implements a simple LRU cache of interface{}-typed key-value pairs. 25 // When items are added, the "size" of the item must be provided. LRU items will 26 // be expired until the total of all items is below the specified size for the 27 // SizeCache 28 import ( 29 "container/list" 30 "sync" 31 32 "github.com/dolthub/dolt/go/store/d" 33 ) 34 35 type sizeCacheEntry struct { 36 size uint64 37 lruEntry *list.Element 38 value interface{} 39 } 40 41 type SizeCache struct { 42 totalSize uint64 43 maxSize uint64 44 mu sync.Mutex 45 lru list.List 46 cache map[interface{}]sizeCacheEntry 47 expireCb func(elm interface{}) 48 } 49 50 type ExpireCallback func(key interface{}) 51 52 // New creates a SizeCache that will hold up to |maxSize| item data. 53 func New(maxSize uint64) *SizeCache { 54 return NewWithExpireCallback(maxSize, nil) 55 } 56 57 // NewWithExpireCallback creates a SizeCache that will hold up to |maxSize| 58 // item data, and will call cb(key) when the item corresponding with that key 59 // expires. 60 func NewWithExpireCallback(maxSize uint64, cb ExpireCallback) *SizeCache { 61 return &SizeCache{ 62 maxSize: maxSize, 63 cache: map[interface{}]sizeCacheEntry{}, 64 expireCb: cb, 65 } 66 } 67 68 // entry() checks if the value is in the cache. If not in the cache, it returns an 69 // empty sizeCacheEntry and false. It it is in the cache, it moves it to 70 // to the back of lru and returns the entry and true. 71 // Callers should have locked down the |c| with a call to c.mu.Lock() before 72 // calling this entry(). 73 func (c *SizeCache) entry(key interface{}) (sizeCacheEntry, bool) { 74 entry, ok := c.cache[key] 75 if !ok { 76 return sizeCacheEntry{}, false 77 } 78 c.lru.MoveToBack(entry.lruEntry) 79 return entry, true 80 } 81 82 // Get checks the searches the cache for an entry. If it exists, it moves it's 83 // lru entry to the back of the queue and returns (value, true). Otherwise, it 84 // returns (nil, false). 85 func (c *SizeCache) Get(key interface{}) (interface{}, bool) { 86 c.mu.Lock() 87 defer c.mu.Unlock() 88 89 if entry, ok := c.entry(key); ok { 90 return entry.value, true 91 } 92 return nil, false 93 } 94 95 // Add will add this element to the cache at the back of the queue as long it's 96 // size does not exceed maxSize. If the addition of this entry causes the size of 97 // the cache to exceed maxSize, the necessary entries at the front of the queue 98 // will be deleted in order to keep the total cache size below maxSize. 99 func (c *SizeCache) Add(key interface{}, size uint64, value interface{}) { 100 if size <= c.maxSize { 101 c.mu.Lock() 102 defer c.mu.Unlock() 103 104 if _, ok := c.entry(key); ok { 105 // this value is already in the cache; just return 106 return 107 } 108 109 newEl := c.lru.PushBack(key) 110 ce := sizeCacheEntry{size: size, lruEntry: newEl, value: value} 111 c.cache[key] = ce 112 c.totalSize += ce.size 113 for el := c.lru.Front(); el != nil && c.totalSize > c.maxSize; { 114 key1 := el.Value 115 ce, ok := c.cache[key1] 116 if !ok { 117 d.Panic("SizeCache is missing expected value") 118 } 119 next := el.Next() 120 delete(c.cache, key1) 121 c.totalSize -= ce.size 122 c.lru.Remove(el) 123 if c.expireCb != nil { 124 c.expireCb(key1) 125 } 126 el = next 127 } 128 } 129 } 130 131 // Drop will remove the element associated with the given key from the cache. 132 func (c *SizeCache) Drop(key interface{}) { 133 c.mu.Lock() 134 defer c.mu.Unlock() 135 136 if entry, ok := c.entry(key); ok { 137 c.totalSize -= entry.size 138 c.lru.Remove(entry.lruEntry) 139 delete(c.cache, key) 140 } 141 } 142 143 func (c *SizeCache) Purge() { 144 c.mu.Lock() 145 defer c.mu.Unlock() 146 147 for key := range c.cache { 148 delete(c.cache, key) 149 } 150 c.totalSize = 0 151 c.lru = list.List{} 152 } 153 154 func (c *SizeCache) Size() uint64 { 155 return c.maxSize 156 }