-
-
Notifications
You must be signed in to change notification settings - Fork 1.6k
/
memory.go
135 lines (116 loc) · 2.55 KB
/
memory.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
// Package memory Is a copy of the storage memory from the external storage packet as a purpose to test the behavior
// in the unittests when using a storages from these packets
package memory
import (
"sync"
"sync/atomic"
"time"
"github.com/gofiber/fiber/v2/utils"
)
// Storage interface that is implemented by storage providers
type Storage struct {
mux sync.RWMutex
db map[string]entry
gcInterval time.Duration
done chan struct{}
}
type entry struct {
data []byte
// max value is 4294967295 -> Sun Feb 07 2106 06:28:15 GMT+0000
expiry uint32
}
// New creates a new memory storage
func New(config ...Config) *Storage {
// Set default config
cfg := configDefault(config...)
// Create storage
store := &Storage{
db: make(map[string]entry),
gcInterval: cfg.GCInterval,
done: make(chan struct{}),
}
// Start garbage collector
utils.StartTimeStampUpdater()
go store.gc()
return store
}
// Get value by key
func (s *Storage) Get(key string) ([]byte, error) {
if len(key) <= 0 {
return nil, nil
}
s.mux.RLock()
v, ok := s.db[key]
s.mux.RUnlock()
if !ok || v.expiry != 0 && v.expiry <= atomic.LoadUint32(&utils.Timestamp) {
return nil, nil
}
return v.data, nil
}
// Set key with value
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
// Ain't Nobody Got Time For That
if len(key) <= 0 || len(val) <= 0 {
return nil
}
var expire uint32
if exp != 0 {
expire = uint32(exp.Seconds()) + atomic.LoadUint32(&utils.Timestamp)
}
s.mux.Lock()
s.db[key] = entry{val, expire}
s.mux.Unlock()
return nil
}
// Delete key by key
func (s *Storage) Delete(key string) error {
// Ain't Nobody Got Time For That
if len(key) <= 0 {
return nil
}
s.mux.Lock()
delete(s.db, key)
s.mux.Unlock()
return nil
}
// Reset all keys
func (s *Storage) Reset() error {
s.mux.Lock()
s.db = make(map[string]entry)
s.mux.Unlock()
return nil
}
// Close the memory storage
func (s *Storage) Close() error {
s.done <- struct{}{}
return nil
}
func (s *Storage) gc() {
ticker := time.NewTicker(s.gcInterval)
defer ticker.Stop()
var expired []string
for {
select {
case <-s.done:
return
case <-ticker.C:
expired = expired[:0]
s.mux.RLock()
for id, v := range s.db {
if v.expiry != 0 && v.expiry < atomic.LoadUint32(&utils.Timestamp) {
expired = append(expired, id)
}
}
s.mux.RUnlock()
s.mux.Lock()
for i := range expired {
delete(s.db, expired[i])
}
s.mux.Unlock()
}
}
}
// Return database client
func (s *Storage) Conn() map[string]entry {
return s.db
}