forked from aws/aws-dax-go
/
lru.go
156 lines (131 loc) · 3.14 KB
/
lru.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
/*
Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License").
You may not use this file except in compliance with the License.
A copy of the License is located at
http://www.apache.org/licenses/LICENSE-2.0
or in the "license" file accompanying this file. This file is distributed
on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
express or implied. See the License for the specific language governing
permissions and limitations under the License.
*/
package lru
import (
"github.com/aws/aws-sdk-go/aws"
"sync"
)
// Lru is a cache which is safe for concurrent access.
type Lru struct {
// MaxEntries is the maximum number of cache entries
// before an item is evicted. Zero means no limit.
MaxEntries int
// LoadFunc specifies the function that loads a value
// for a specific key when not found in the cache.
LoadFunc func(ctx aws.Context, key Key) (interface{}, error)
loadGroup loadGroup
// Optional KeyMarshaller. Caller should provide one when using
// Key type which is not comparable. eg. slice
KeyMarshaller func(key Key) Key
mu sync.RWMutex
cache map[Key]*entry
head, tail *entry
}
type Key interface{}
type entry struct {
key Key
value interface{}
prev, next *entry
}
func (c *Lru) contains(key Key) bool {
c.mu.RLock()
defer c.mu.RUnlock()
_, ok := c.cache[key]
return ok
}
func (c *Lru) lookup(key Key) (*entry, bool) {
c.mu.RLock()
defer c.mu.RUnlock()
v, ok := c.cache[key]
return v, ok
}
func (c *Lru) GetWithContext(ctx aws.Context, okey Key) (interface{}, error) {
ikey := okey
if c.KeyMarshaller != nil {
ikey = c.KeyMarshaller(okey)
}
if en, ok := c.lookup(ikey); ok {
return en.value, nil
}
v, err := c.loadGroup.do(ikey, func() (interface{}, error) {
if en, ok := c.lookup(ikey); ok {
return en.value, nil
}
val, err := c.LoadFunc(ctx, okey)
if err != nil {
return nil, err
}
c.mu.Lock()
defer c.mu.Unlock()
en := &entry{key: ikey, value: val}
if c.tail == nil {
c.head = en
c.tail = en
} else {
en.prev = c.tail
c.tail.next = en
c.tail = en
}
if c.cache == nil {
c.cache = make(map[Key]*entry)
}
c.cache[ikey] = en
// Evict oldest entry if over the max.
if c.MaxEntries > 0 && len(c.cache) > c.MaxEntries {
evict := c.head
if evict != nil {
delete(c.cache, evict.key)
c.head = evict.next
if c.head != nil {
c.head.prev = nil
}
evict.next = nil
}
}
return val, nil
})
return v, err
}
type loader struct {
wg sync.WaitGroup
value interface{}
err error
}
type loadGroup struct {
mu sync.Mutex
m map[Key]*loader
}
func (g *loadGroup) do(key Key, loadFn func() (interface{}, error)) (interface{}, error) {
g.mu.Lock()
if g.m == nil {
g.m = make(map[Key]*loader)
}
if l, ok := g.m[key]; ok {
g.mu.Unlock()
l.wg.Wait()
return l.value, l.err
}
v := &loader{}
v.wg.Add(1)
g.m[key] = v
g.mu.Unlock()
v.value, v.err = loadFn()
v.wg.Done()
// cleanup loader
g.mu.Lock()
defer g.mu.Unlock()
delete(g.m, key)
if len(g.m) <= 0 {
g.m = nil
}
return v.value, v.err
}