hanayo/vendor/github.com/thehowl/cieca/cieca.go

118 lines
2.5 KiB
Go
Raw Normal View History

2019-02-23 13:29:15 +00:00
// Package cieca is an in-memory Key-Value datastore for Go. The main purpose
// is to easily cache data in the RAM for a certain amount of time.
package cieca
import (
"sync"
"time"
)
type element struct {
data []byte
expire time.Time
cancel chan struct{}
}
// DataStore is the datastore containing all the data.
type DataStore struct {
data map[string]element
mutex *sync.RWMutex
}
// Get retrieves a value in the datastore. If it is not found, nil is returned.
func (d *DataStore) Get(key string) []byte {
v, _ := d.GetWithExist(key)
return v
}
// GetWithExist retrieves the value of a key, and a boolean indicating whether
// the value existed in the store.
func (d *DataStore) GetWithExist(key string) ([]byte, bool) {
d.setup()
d.mutex.RLock()
defer d.mutex.RUnlock()
el, ex := d.data[key]
return el.data, ex
}
// Expiration gets the time of expiration of a key.
func (d *DataStore) Expiration(key string) *time.Time {
d.setup()
d.mutex.RLock()
defer d.mutex.RUnlock()
el, exist := d.data[key]
if !exist || el.cancel == nil {
return nil
}
t := el.expire
return &t
}
// Delete removes an element from the datastore.
func (d *DataStore) Delete(key string) {
d.setup()
d.mutex.Lock()
defer d.mutex.Unlock()
exp := d.data[key].cancel
if exp != nil && time.Now().Sub(d.data[key].expire) <= 0 {
exp <- struct{}{}
close(exp)
}
delete(d.data, key)
}
// Set sets a key in the datastore with no expiration.
func (d *DataStore) Set(key string, value []byte) {
d.SetWithExpiration(key, value, -1)
}
// SetWithExpiration sets a key in the datastore with a certain value.
func (d *DataStore) SetWithExpiration(key string, value []byte, expiration time.Duration) {
d.setup()
var c chan struct{}
if expiration >= 0 {
c = make(chan struct{})
}
el := element{
data: value,
expire: time.Now().Add(expiration),
cancel: c,
}
d.Delete(key)
d.mutex.Lock()
defer d.mutex.Unlock()
d.data[key] = el
if c != nil {
go d.queueDeletion(expiration, c, key)
}
}
// Clean clears the datastore. Not so thread-safe. Use with care.
func (d *DataStore) Clean() {
if d == nil {
return
}
for el := range d.data {
d.Delete(el)
}
}
func (d *DataStore) queueDeletion(dur time.Duration, canc <-chan struct{}, key string) {
select {
case <-time.NewTimer(dur).C:
d.Delete(key)
case <-canc:
// useless, but explicits what we're doing
return
}
}
func (d *DataStore) setup() {
if d.data == nil {
d.data = make(map[string]element)
}
if d.mutex == nil {
d.mutex = &sync.RWMutex{}
}
}