118 lines
2.5 KiB
Go
118 lines
2.5 KiB
Go
|
// 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{}
|
||
|
|
}
|
||
|
|
}
|