heap.go 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. package cache
  2. import (
  3. "container/heap"
  4. )
  5. type heapEntry struct {
  6. key string
  7. exp uint64
  8. bytes uint
  9. idx int
  10. }
  11. // indexedHeap is a regular min-heap that allows finding
  12. // elements in constant time. It does so by handing out special indices
  13. // and tracking entry movement.
  14. //
  15. // indexdedHeap is used for quickly finding entries with the lowest
  16. // expiration timestamp and deleting arbitrary entries.
  17. type indexedHeap struct {
  18. // Slice the heap is built on
  19. entries []heapEntry
  20. // Mapping "index" to position in heap slice
  21. indices []int
  22. // Max index handed out
  23. maxidx int
  24. }
  25. func (h indexedHeap) Len() int {
  26. return len(h.entries)
  27. }
  28. func (h indexedHeap) Less(i, j int) bool {
  29. return h.entries[i].exp < h.entries[j].exp
  30. }
  31. func (h indexedHeap) Swap(i, j int) {
  32. h.entries[i], h.entries[j] = h.entries[j], h.entries[i]
  33. h.indices[h.entries[i].idx] = i
  34. h.indices[h.entries[j].idx] = j
  35. }
  36. func (h *indexedHeap) Push(x interface{}) {
  37. h.pushInternal(x.(heapEntry)) //nolint:forcetypeassert // Forced type assertion required to implement the heap.Interface interface
  38. }
  39. func (h *indexedHeap) Pop() interface{} {
  40. n := len(h.entries)
  41. h.entries = h.entries[0 : n-1]
  42. return h.entries[0:n][n-1]
  43. }
  44. func (h *indexedHeap) pushInternal(entry heapEntry) {
  45. h.indices[entry.idx] = len(h.entries)
  46. h.entries = append(h.entries, entry)
  47. }
  48. // Returns index to track entry
  49. func (h *indexedHeap) put(key string, exp uint64, bytes uint) int {
  50. idx := 0
  51. if len(h.entries) < h.maxidx {
  52. // Steal index from previously removed entry
  53. // capacity > size is guaranteed
  54. n := len(h.entries)
  55. idx = h.entries[:n+1][n].idx
  56. } else {
  57. idx = h.maxidx
  58. h.maxidx++
  59. h.indices = append(h.indices, idx)
  60. }
  61. // Push manually to avoid allocation
  62. h.pushInternal(heapEntry{
  63. key: key, exp: exp, idx: idx, bytes: bytes,
  64. })
  65. heap.Fix(h, h.Len()-1)
  66. return idx
  67. }
  68. func (h *indexedHeap) removeInternal(realIdx int) (string, uint) {
  69. x := heap.Remove(h, realIdx).(heapEntry) //nolint:forcetypeassert,errcheck // Forced type assertion required to implement the heap.Interface interface
  70. return x.key, x.bytes
  71. }
  72. // Remove entry by index
  73. func (h *indexedHeap) remove(idx int) (string, uint) {
  74. return h.removeInternal(h.indices[idx])
  75. }
  76. // Remove entry with lowest expiration time
  77. func (h *indexedHeap) removeFirst() (string, uint) {
  78. return h.removeInternal(0)
  79. }