127 lines
2.6 KiB
Go
127 lines
2.6 KiB
Go
package store
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// memoryStoreItem holds the value and expiration timestamp for a key.
|
|
type memoryStoreItem struct {
|
|
value []byte
|
|
expiresAt int64 // Unix-nano timestamp. 0 for no expiry.
|
|
}
|
|
|
|
// MemoryStore is an in-memory key-value store that is safe for concurrent use.
|
|
type MemoryStore struct {
|
|
mu sync.RWMutex
|
|
data map[string]memoryStoreItem
|
|
stopCh chan struct{} // Channel to stop the cleanup goroutine
|
|
}
|
|
|
|
// NewMemoryStore creates and returns a new MemoryStore instance.
|
|
// It also starts a background goroutine to periodically clean up expired keys.
|
|
func NewMemoryStore() *MemoryStore {
|
|
s := &MemoryStore{
|
|
data: make(map[string]memoryStoreItem),
|
|
stopCh: make(chan struct{}),
|
|
}
|
|
go s.cleanupLoop(1 * time.Minute)
|
|
return s
|
|
}
|
|
|
|
// Close stops the background cleanup goroutine.
|
|
func (s *MemoryStore) Close() error {
|
|
close(s.stopCh)
|
|
return nil
|
|
}
|
|
|
|
// cleanupLoop periodically iterates through the store and removes expired keys.
|
|
func (s *MemoryStore) cleanupLoop(interval time.Duration) {
|
|
ticker := time.NewTicker(interval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
s.mu.Lock()
|
|
now := time.Now().UnixNano()
|
|
for key, item := range s.data {
|
|
if item.expiresAt > 0 && now > item.expiresAt {
|
|
delete(s.data, key)
|
|
}
|
|
}
|
|
s.mu.Unlock()
|
|
case <-s.stopCh:
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// Set stores a key-value pair.
|
|
func (s *MemoryStore) Set(key string, value []byte, ttl time.Duration) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
var expiresAt int64
|
|
if ttl > 0 {
|
|
expiresAt = time.Now().UnixNano() + ttl.Nanoseconds()
|
|
}
|
|
|
|
s.data[key] = memoryStoreItem{
|
|
value: value,
|
|
expiresAt: expiresAt,
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Get retrieves a value by its key.
|
|
func (s *MemoryStore) Get(key string) ([]byte, error) {
|
|
s.mu.RLock()
|
|
item, exists := s.data[key]
|
|
s.mu.RUnlock()
|
|
|
|
if !exists {
|
|
return nil, ErrNotFound
|
|
}
|
|
|
|
// Check for expiration
|
|
if item.expiresAt > 0 && time.Now().UnixNano() > item.expiresAt {
|
|
// Lazy deletion
|
|
s.mu.Lock()
|
|
delete(s.data, key)
|
|
s.mu.Unlock()
|
|
return nil, ErrNotFound
|
|
}
|
|
|
|
return item.value, nil
|
|
}
|
|
|
|
// Delete removes a value by its key.
|
|
func (s *MemoryStore) Delete(key string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
delete(s.data, key)
|
|
return nil
|
|
}
|
|
|
|
// Exists checks if a key exists.
|
|
func (s *MemoryStore) Exists(key string) (bool, error) {
|
|
s.mu.RLock()
|
|
item, exists := s.data[key]
|
|
s.mu.RUnlock()
|
|
|
|
if !exists {
|
|
return false, nil
|
|
}
|
|
|
|
if item.expiresAt > 0 && time.Now().UnixNano() > item.expiresAt {
|
|
// Lazy deletion
|
|
s.mu.Lock()
|
|
delete(s.data, key)
|
|
s.mu.Unlock()
|
|
return false, nil
|
|
}
|
|
|
|
return true, nil
|
|
}
|