-
Notifications
You must be signed in to change notification settings - Fork 0
/
cache.go
88 lines (77 loc) · 2.13 KB
/
cache.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
package goclash
import (
"strconv"
"sync"
"time"
"github.com/go-resty/resty/v2"
cmap "github.com/orcaman/concurrent-map/v2"
)
// Cache is a simple but performant in-memory cache.
type Cache struct {
enabled bool
store cmap.ConcurrentMap[string, *cachedValue]
cacheTime time.Duration
mu sync.Mutex
}
type cachedValue struct {
data []byte // data is the cached value
timer *time.Timer // timer schedules removal of cached value
}
func newCache() *Cache {
return &Cache{
store: cmap.New[*cachedValue](),
enabled: true,
}
}
// UseCache sets whether to use cache.
//
// Cache is capable of storing large amounts of data in memory, across different shards. Using it together with SetCacheTime may replace the need for a store like Redis, depending on your needs.
func (h *Client) UseCache(v bool) {
h.cache.mu.Lock()
defer h.cache.mu.Unlock()
h.cache.enabled = v
}
// SetCacheTime sets a fixed cache time, ignoring CacheControl headers. Disable by passing 0 as argument.
func (h *Client) SetCacheTime(d time.Duration) {
h.cache.mu.Lock()
defer h.cache.mu.Unlock()
h.cache.cacheTime = d
}
// Get gets a value from the cache, and a boolean indicating whether the value was found.
func (c *Cache) Get(key string) ([]byte, bool) {
if !c.enabled {
return nil, false
}
value, ok := c.store.Get(key)
if !ok {
return nil, false
}
return value.data, ok
}
// Set sets a value in the cache, with a duration after it gets removed.
func (c *Cache) Set(key string, data []byte, duration time.Duration) {
if !c.enabled {
return
}
if value, ok := c.store.Get(key); ok {
value.timer.Stop()
}
c.store.Set(key, &cachedValue{
data: data,
timer: time.AfterFunc(duration, func() {
c.store.Remove(key)
}),
})
}
// CacheResponse caches the response body of a resty.Response, using the Cache-Control header to determine the cache time.
func (c *Cache) CacheResponse(url string, res *resty.Response) {
if c.cacheTime > 0 {
c.Set(url, res.Body(), c.cacheTime)
return
}
seconds, err := strconv.Atoi(res.Header().Get("Cache-Control")[8:])
if err != nil {
return
}
c.Set(url, res.Body(), time.Duration(seconds)*time.Second)
}