diff --git a/cache.go b/cache.go index 50824a7..26844ee 100644 --- a/cache.go +++ b/cache.go @@ -15,11 +15,15 @@ import ( type Item struct { Object interface{} Expiration int64 +} + +type Node struct { + Expiration int64 Key string } -func (item Item) Less(than llrb.Item) bool { - return item.Expiration < than.(Item).Expiration +func (node Node) Less(than llrb.Item) bool { + return node.Expiration < than.(Node).Expiration } @@ -66,23 +70,41 @@ func (c *cache) Set(k string, x interface{}, d time.Duration) { } func (c *cache) set(k string, x interface{}, d time.Duration) { - item := Item{ - Object: x, - Key : k, - } + item := Item{Object: x,} if d == DefaultExpiration { d = c.defaultExpiration } if d > 0 { item.Expiration = time.Now().Add(d).UnixNano() - _, found := c.items[k] - if !found { - c.sortedItems.InsertNoReplace(item) + old, found := c.items[k] + if found && old.Expiration != item.Expiration{ + c.deleteFromBst(Node{Expiration: old.Expiration, Key: k}) + c.sortedItems.InsertNoReplace(Node{Expiration: item.Expiration, Key: k}) + } else if !found { + c.sortedItems.InsertNoReplace(Node{Expiration: item.Expiration, Key: k}) } } c.items[k] = item } +func (c *cache) deleteFromBst (node Node) { + //delete nodes from the tree with the same expiration + //until the required one is found + var toReinsert []Node + for del := c.sortedItems.Delete(node); del != nil; { + delNode := del.(Node) + if delNode.Key == node.Key { + break + } else { + toReinsert = append (toReinsert, delNode) + } + } + //reinsert the nodes in the tree, with modified expiration + for _, delNode := range toReinsert { + c.sortedItems.InsertNoReplace(delNode) + } +} + // Add an item to the cache only if an item doesn't already exist for the given // key, or if the existing item has expired. Returns an error otherwise. @@ -880,34 +902,49 @@ func (c *cache) Delete(k string) { } func (c *cache) delete(k string) (interface{}, bool) { - if c.onEvicted != nil { - if v, found := c.items[k]; found { - delete(c.items, k) + if v, found := c.items[k]; found { + delete(c.items, k) + c.deleteFromBst(Node{Expiration: v.Expiration, Key: k}) + if c.onEvicted != nil { return v.Object, true + } else { + return nil, false } } - delete(c.items, k) return nil, false } + +type keyAndValue struct { + key string + value interface{} +} + // Delete all expired items from the cache. func (c *cache) DeleteExpired() { - var evictedItems []Item + var evictedNodes []Node + var evictedItems []keyAndValue c.mu.Lock() - c.sortedItems.DescendLessOrEqual(Item{Expiration: time.Now().UnixNano()}, func(i llrb.Item) bool { - v := i.(Item) - c.delete(v.Key) - evictedItems = append(evictedItems, v) + c.sortedItems.DescendLessOrEqual(Node{Expiration: time.Now().UnixNano()}, func(i llrb.Item) bool { + v := i.(Node) + k := v.Key + item, found := c.items[k] + if !found { + panic("Item in tree but not in map!!") + } + delete(c.items, k) + evictedItems = append(evictedItems, keyAndValue{k, item.Object}) + evictedNodes = append(evictedNodes, v) return true }) - for _, v := range evictedItems { + for _, v := range evictedNodes { c.sortedItems.Delete(v) } c.mu.Unlock() - for _, v := range evictedItems { - if c.onEvicted != nil { - c.onEvicted(v.Key, v.Object) - } + if c.onEvicted != nil { + for _, n := range evictedItems { + c.onEvicted(n.key, n.value) + } } } @@ -976,10 +1013,10 @@ func (c *cache) Load(r io.Reader) error { ov, found := c.items[k] if !found || ov.Expired() { c.items[k] = v - if found { - c.sortedItems.Delete(ov) + if !found { + c.sortedItems.InsertNoReplace(Node{Expiration: v.Expiration, Key: k}) } - c.sortedItems.InsertNoReplace(v) + } } @@ -1073,9 +1110,8 @@ func newCache(de time.Duration, m map[string]Item) *cache { items: m, } c.sortedItems = llrb.New() - //we can probably do bulk insertion here to speed it up - for _, item := range m { - c.sortedItems.InsertNoReplace(item) + for k, item := range m { + c.sortedItems.InsertNoReplace(Node{Key: k, Expiration: item.Expiration}) } return c } diff --git a/cache_test.go b/cache_test.go index f9b7a96..67de7e7 100644 --- a/cache_test.go +++ b/cache_test.go @@ -1680,5 +1680,13 @@ func benchmarkLargeCache(b *testing.B, nano int) { for i := 0; i < b.N; i++ { tc.Set(strconv.Itoa(i), "bar", DefaultExpiration) } + b.StopTimer() + tc.DeleteExpired() + now := time.Now().UnixNano() + for _, item := range tc.Items() { + if item.Expiration < now { + b.Fatalf("some items have not been correctly evicted") + } + } }