75 lines
1.2 KiB
Go
75 lines
1.2 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"time"
|
||
|
|
"sync"
|
||
|
|
"github.com/miekg/dns"
|
||
|
|
)
|
||
|
|
|
||
|
|
type MemoryCache struct {
|
||
|
|
Cache
|
||
|
|
|
||
|
|
Backend map[string]Mesg
|
||
|
|
Expire time.Duration
|
||
|
|
Maxcount int
|
||
|
|
mu sync.RWMutex
|
||
|
|
}
|
||
|
|
|
||
|
|
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
|
||
|
|
}
|