lrucache.go 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294
  1. package lru
  2. // Modified by https://github.com/die-net/lrucache
  3. import (
  4. "sync"
  5. "time"
  6. list "github.com/bahlo/generic-list-go"
  7. "github.com/samber/lo"
  8. )
  9. // Option is part of Functional Options Pattern
  10. type Option[K comparable, V any] func(*LruCache[K, V])
  11. // EvictCallback is used to get a callback when a cache entry is evicted
  12. type EvictCallback[K comparable, V any] func(key K, value V)
  13. // WithEvict set the evict callback
  14. func WithEvict[K comparable, V any](cb EvictCallback[K, V]) Option[K, V] {
  15. return func(l *LruCache[K, V]) {
  16. l.onEvict = cb
  17. }
  18. }
  19. // WithUpdateAgeOnGet update expires when Get element
  20. func WithUpdateAgeOnGet[K comparable, V any]() Option[K, V] {
  21. return func(l *LruCache[K, V]) {
  22. l.updateAgeOnGet = true
  23. }
  24. }
  25. // WithAge defined element max age (second)
  26. func WithAge[K comparable, V any](maxAge int64) Option[K, V] {
  27. return func(l *LruCache[K, V]) {
  28. l.maxAge = maxAge
  29. }
  30. }
  31. // WithSize defined max length of LruCache
  32. func WithSize[K comparable, V any](maxSize int) Option[K, V] {
  33. return func(l *LruCache[K, V]) {
  34. l.maxSize = maxSize
  35. }
  36. }
  37. // WithStale decide whether Stale return is enabled.
  38. // If this feature is enabled, element will not get Evicted according to `WithAge`.
  39. func WithStale[K comparable, V any](stale bool) Option[K, V] {
  40. return func(l *LruCache[K, V]) {
  41. l.staleReturn = stale
  42. }
  43. }
  44. // LruCache is a thread-safe, in-memory lru-cache that evicts the
  45. // least recently used entries from memory when (if set) the entries are
  46. // older than maxAge (in seconds). Use the New constructor to create one.
  47. type LruCache[K comparable, V any] struct {
  48. maxAge int64
  49. maxSize int
  50. mu sync.Mutex
  51. cache map[K]*list.Element[*entry[K, V]]
  52. lru *list.List[*entry[K, V]] // Front is least-recent
  53. updateAgeOnGet bool
  54. staleReturn bool
  55. onEvict EvictCallback[K, V]
  56. }
  57. // New creates an LruCache
  58. func New[K comparable, V any](options ...Option[K, V]) *LruCache[K, V] {
  59. lc := &LruCache[K, V]{
  60. lru: list.New[*entry[K, V]](),
  61. cache: make(map[K]*list.Element[*entry[K, V]]),
  62. }
  63. for _, option := range options {
  64. option(lc)
  65. }
  66. return lc
  67. }
  68. // Get returns any representation of a cached response and a bool
  69. // set to true if the key was found.
  70. func (c *LruCache[K, V]) Get(key K) (V, bool) {
  71. c.mu.Lock()
  72. defer c.mu.Unlock()
  73. el := c.get(key)
  74. if el == nil {
  75. return lo.Empty[V](), false
  76. }
  77. value := el.value
  78. return value, true
  79. }
  80. func (c *LruCache[K, V]) GetOrStore(key K, constructor func() V) (V, bool) {
  81. c.mu.Lock()
  82. defer c.mu.Unlock()
  83. el := c.get(key)
  84. if el == nil {
  85. value := constructor()
  86. c.set(key, value)
  87. return value, false
  88. }
  89. value := el.value
  90. return value, true
  91. }
  92. // GetWithExpire returns any representation of a cached response,
  93. // a time.Time Give expected expires,
  94. // and a bool set to true if the key was found.
  95. // This method will NOT check the maxAge of element and will NOT update the expires.
  96. func (c *LruCache[K, V]) GetWithExpire(key K) (V, time.Time, bool) {
  97. c.mu.Lock()
  98. defer c.mu.Unlock()
  99. el := c.get(key)
  100. if el == nil {
  101. return lo.Empty[V](), time.Time{}, false
  102. }
  103. return el.value, time.Unix(el.expires, 0), true
  104. }
  105. // Exist returns if key exist in cache but not put item to the head of linked list
  106. func (c *LruCache[K, V]) Exist(key K) bool {
  107. c.mu.Lock()
  108. defer c.mu.Unlock()
  109. _, ok := c.cache[key]
  110. return ok
  111. }
  112. // Set stores any representation of a response for a given key.
  113. func (c *LruCache[K, V]) Set(key K, value V) {
  114. c.mu.Lock()
  115. defer c.mu.Unlock()
  116. c.set(key, value)
  117. }
  118. func (c *LruCache[K, V]) set(key K, value V) {
  119. expires := int64(0)
  120. if c.maxAge > 0 {
  121. expires = time.Now().Unix() + c.maxAge
  122. }
  123. c.setWithExpire(key, value, time.Unix(expires, 0))
  124. }
  125. // SetWithExpire stores any representation of a response for a given key and given expires.
  126. // The expires time will round to second.
  127. func (c *LruCache[K, V]) SetWithExpire(key K, value V, expires time.Time) {
  128. c.mu.Lock()
  129. defer c.mu.Unlock()
  130. c.setWithExpire(key, value, expires)
  131. }
  132. func (c *LruCache[K, V]) setWithExpire(key K, value V, expires time.Time) {
  133. if le, ok := c.cache[key]; ok {
  134. c.lru.MoveToBack(le)
  135. e := le.Value
  136. e.value = value
  137. e.expires = expires.Unix()
  138. } else {
  139. e := &entry[K, V]{key: key, value: value, expires: expires.Unix()}
  140. c.cache[key] = c.lru.PushBack(e)
  141. if c.maxSize > 0 {
  142. if elLen := c.lru.Len(); elLen > c.maxSize {
  143. c.deleteElement(c.lru.Front())
  144. }
  145. }
  146. }
  147. c.maybeDeleteOldest()
  148. }
  149. // CloneTo clone and overwrite elements to another LruCache
  150. func (c *LruCache[K, V]) CloneTo(n *LruCache[K, V]) {
  151. c.mu.Lock()
  152. defer c.mu.Unlock()
  153. n.mu.Lock()
  154. defer n.mu.Unlock()
  155. n.lru = list.New[*entry[K, V]]()
  156. n.cache = make(map[K]*list.Element[*entry[K, V]])
  157. for e := c.lru.Front(); e != nil; e = e.Next() {
  158. elm := e.Value
  159. n.cache[elm.key] = n.lru.PushBack(elm)
  160. }
  161. }
  162. func (c *LruCache[K, V]) get(key K) *entry[K, V] {
  163. le, ok := c.cache[key]
  164. if !ok {
  165. return nil
  166. }
  167. if !c.staleReturn && c.maxAge > 0 && le.Value.expires <= time.Now().Unix() {
  168. c.deleteElement(le)
  169. c.maybeDeleteOldest()
  170. return nil
  171. }
  172. c.lru.MoveToBack(le)
  173. el := le.Value
  174. if c.maxAge > 0 && c.updateAgeOnGet {
  175. el.expires = time.Now().Unix() + c.maxAge
  176. }
  177. return el
  178. }
  179. // Delete removes the value associated with a key.
  180. func (c *LruCache[K, V]) Delete(key K) {
  181. c.mu.Lock()
  182. defer c.mu.Unlock()
  183. c.delete(key)
  184. }
  185. func (c *LruCache[K, V]) delete(key K) {
  186. if le, ok := c.cache[key]; ok {
  187. c.deleteElement(le)
  188. }
  189. }
  190. func (c *LruCache[K, V]) maybeDeleteOldest() {
  191. if !c.staleReturn && c.maxAge > 0 {
  192. now := time.Now().Unix()
  193. for le := c.lru.Front(); le != nil && le.Value.expires <= now; le = c.lru.Front() {
  194. c.deleteElement(le)
  195. }
  196. }
  197. }
  198. func (c *LruCache[K, V]) deleteElement(le *list.Element[*entry[K, V]]) {
  199. c.lru.Remove(le)
  200. e := le.Value
  201. delete(c.cache, e.key)
  202. if c.onEvict != nil {
  203. c.onEvict(e.key, e.value)
  204. }
  205. }
  206. func (c *LruCache[K, V]) Clear() error {
  207. c.mu.Lock()
  208. defer c.mu.Unlock()
  209. c.cache = make(map[K]*list.Element[*entry[K, V]])
  210. return nil
  211. }
  212. // Compute either sets the computed new value for the key or deletes
  213. // the value for the key. When the delete result of the valueFn function
  214. // is set to true, the value will be deleted, if it exists. When delete
  215. // is set to false, the value is updated to the newValue.
  216. // The ok result indicates whether value was computed and stored, thus, is
  217. // present in the map. The actual result contains the new value in cases where
  218. // the value was computed and stored.
  219. func (c *LruCache[K, V]) Compute(
  220. key K,
  221. valueFn func(oldValue V, loaded bool) (newValue V, delete bool),
  222. ) (actual V, ok bool) {
  223. c.mu.Lock()
  224. defer c.mu.Unlock()
  225. if el := c.get(key); el != nil {
  226. actual, ok = el.value, true
  227. }
  228. if newValue, del := valueFn(actual, ok); del {
  229. if ok { // data not in cache, so needn't delete
  230. c.delete(key)
  231. }
  232. return lo.Empty[V](), false
  233. } else {
  234. c.set(key, newValue)
  235. return newValue, true
  236. }
  237. }
  238. type entry[K comparable, V any] struct {
  239. key K
  240. value V
  241. expires int64
  242. }