136 lines
2.8 KiB
Go
136 lines
2.8 KiB
Go
package cache
|
|
|
|
import (
|
|
"container/list"
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type MemoryConfig struct {
|
|
MaxEntries int `json:"maxEntries"`
|
|
MaxBytes int `json:"maxBytes"` // Sum of key and value bytes, excluding Go object overhead.
|
|
}
|
|
|
|
type entry struct {
|
|
key string
|
|
value []byte
|
|
expires time.Time
|
|
}
|
|
|
|
// Memory is an LRU cache with lazy expiration and no background goroutines.
|
|
type Memory struct {
|
|
mu sync.Mutex
|
|
config MemoryConfig
|
|
items map[string]*list.Element
|
|
lru *list.List
|
|
bytes int
|
|
closed bool
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewMemory(cfg MemoryConfig) (*Memory, error) {
|
|
if cfg.MaxEntries < 0 || cfg.MaxBytes < 0 {
|
|
return nil, errors.New("cache: negative capacity")
|
|
}
|
|
if cfg.MaxEntries == 0 {
|
|
cfg.MaxEntries = 4096
|
|
}
|
|
if cfg.MaxBytes == 0 {
|
|
cfg.MaxBytes = 32 << 20
|
|
}
|
|
return &Memory{config: cfg, items: make(map[string]*list.Element), lru: list.New(), now: time.Now}, nil
|
|
}
|
|
|
|
func (m *Memory) Get(ctx context.Context, key string) ([]byte, error) {
|
|
if err := Check(ctx, key, 0); err != nil {
|
|
return nil, err
|
|
}
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.closed {
|
|
return nil, ErrClosed
|
|
}
|
|
e := m.items[key]
|
|
if e == nil {
|
|
return nil, ErrMiss
|
|
}
|
|
item := e.Value.(*entry)
|
|
if !item.expires.IsZero() && !m.now().Before(item.expires) {
|
|
m.remove(e)
|
|
return nil, ErrMiss
|
|
}
|
|
m.lru.MoveToFront(e)
|
|
return append([]byte(nil), item.value...), nil
|
|
}
|
|
|
|
func (m *Memory) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error {
|
|
if err := Check(ctx, key, ttl); err != nil {
|
|
return err
|
|
}
|
|
if len(value) > m.config.MaxBytes-len(key) {
|
|
return ErrTooLarge
|
|
}
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.closed {
|
|
return ErrClosed
|
|
}
|
|
if e := m.items[key]; e != nil {
|
|
m.remove(e)
|
|
}
|
|
for len(m.items) >= m.config.MaxEntries || m.bytes+len(key)+len(value) > m.config.MaxBytes {
|
|
m.remove(m.lru.Back())
|
|
}
|
|
item := &entry{key: key, value: append([]byte(nil), value...)}
|
|
if ttl > 0 {
|
|
item.expires = m.now().Add(ttl)
|
|
}
|
|
m.items[key] = m.lru.PushFront(item)
|
|
m.bytes += len(key) + len(value)
|
|
return nil
|
|
}
|
|
|
|
func (m *Memory) Delete(ctx context.Context, keys ...string) error {
|
|
if ctx == nil {
|
|
return errors.New("cache: nil context")
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
for _, key := range keys {
|
|
if err := Check(ctx, key, 0); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.closed {
|
|
return ErrClosed
|
|
}
|
|
for _, key := range keys {
|
|
if e := m.items[key]; e != nil {
|
|
m.remove(e)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *Memory) remove(e *list.Element) {
|
|
item := e.Value.(*entry)
|
|
m.bytes -= len(item.key) + len(item.value)
|
|
delete(m.items, item.key)
|
|
m.lru.Remove(e)
|
|
}
|
|
|
|
func (m *Memory) Close() error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.closed = true
|
|
clear(m.items)
|
|
m.lru.Init()
|
|
m.bytes = 0
|
|
return nil
|
|
}
|