100 lines
1.6 KiB
Go
100 lines
1.6 KiB
Go
package cache
|
|
|
|
import (
|
|
"github.com/miekg/dns"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
func NewMemoryCache(expire time.Duration, maxCount int) *MemoryCache {
|
|
return &MemoryCache{
|
|
backend: make(map[string]Mesg),
|
|
Expire: expire,
|
|
Maxcount: maxCount,
|
|
}
|
|
}
|
|
|
|
type Mesg struct {
|
|
Msg *dns.Msg
|
|
Expire time.Time
|
|
}
|
|
|
|
type MemoryCache struct {
|
|
Cache
|
|
|
|
backend map[string]Mesg
|
|
Expire time.Duration
|
|
Maxcount int
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
func (c *MemoryCache) initialize() {
|
|
c.backend = make(map[string]Mesg)
|
|
}
|
|
|
|
func (c *MemoryCache) Get(key string) (*dns.Msg, error) {
|
|
c.mu.RLock()
|
|
mesg, ok := c.backend[key]
|
|
c.mu.RUnlock()
|
|
if !ok {
|
|
return nil, KeyNotFound{key}
|
|
}
|
|
|
|
if mesg.Expire.Before(time.Now()) {
|
|
c.Remove(key)
|
|
return nil, KeyExpired{key}
|
|
}
|
|
|
|
return mesg.Msg, nil
|
|
|
|
}
|
|
|
|
func (c *MemoryCache) Set(key string, msg *dns.Msg) error {
|
|
if c.Full() && !c.Exists(key) {
|
|
return CacheIsFull{}
|
|
}
|
|
|
|
expire := time.Now().Add(c.Expire)
|
|
mesg := Mesg{msg, expire}
|
|
c.mu.Lock()
|
|
c.backend[key] = mesg
|
|
c.mu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
func (c *MemoryCache) Remove(key string) error {
|
|
c.mu.Lock()
|
|
delete(c.backend, key)
|
|
c.mu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
func (c *MemoryCache) Exists(key string) bool {
|
|
c.mu.RLock()
|
|
_, ok := c.backend[key]
|
|
c.mu.RUnlock()
|
|
return ok
|
|
}
|
|
|
|
func (c *MemoryCache) Length() int {
|
|
c.mu.RLock()
|
|
defer c.mu.RUnlock()
|
|
return len(c.backend)
|
|
}
|
|
|
|
func (c *MemoryCache) Full() bool {
|
|
// if Maxcount is zero. the cache will never be full.
|
|
if c.Maxcount == 0 {
|
|
return false
|
|
}
|
|
return c.Length() >= c.Maxcount
|
|
}
|
|
|
|
func (c *MemoryCache) Purge() error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
c.initialize()
|
|
|
|
return nil
|
|
} |