1package myks
  2
  3import (
  4	"errors"
  5	"iter"
  6	"maps"
  7	"sync"
  8	"time"
  9)
 10
 11var (
 12	ErrNotFound = errors.New("entry was not found")
 13	ErrExpired  = errors.New("entry is expired")
 14)
 15
 16// KeyStore is the main structure of this module
 17type KeyStore[T any] struct {
 18	mu              sync.RWMutex
 19	data            map[string]entry[T]
 20	stopChan        chan struct{}
 21	cleanupInterval time.Duration
 22}
 23
 24// entry is a keystore value, with an optional expiration
 25type entry[T any] struct {
 26	value      T
 27	expiration *time.Time
 28}
 29
 30// New creates a new keystore with an optional goroutine to automatically clean expired values
 31func New[T any](cleanupInterval time.Duration) *KeyStore[T] {
 32	ks := &KeyStore[T]{ // do not set cleanupInterval here
 33		data:     make(map[string]entry[T]),
 34		stopChan: make(chan struct{}),
 35	}
 36
 37	go ks.StartCleanup(cleanupInterval)
 38
 39	return ks
 40}
 41
 42// Set saves a key-value pair with a given expiration.
 43// If duration <= 0, the entry will never expire.
 44func (ks *KeyStore[T]) Set(key string, value T, duration time.Duration) {
 45	ks.mu.Lock()
 46	defer ks.mu.Unlock()
 47
 48	var expiration *time.Time
 49	if duration > 0 {
 50		exp := time.Now().Add(duration)
 51		expiration = &exp
 52	}
 53
 54	ks.data[key] = entry[T]{
 55		value:      value,
 56		expiration: expiration,
 57	}
 58}
 59
 60// Get returns the value for the given key if it exists and it's not expired.
 61func (ks *KeyStore[T]) Get(key string) (*T, error) {
 62	ks.mu.Lock()
 63	defer ks.mu.Unlock()
 64
 65	ent, exists := ks.data[key]
 66	if !exists {
 67		return nil, ErrNotFound
 68	}
 69
 70	if ent.expiration != nil && ent.expiration.Before(time.Now()) {
 71		delete(ks.data, key)
 72		return nil, ErrExpired
 73	}
 74
 75	return &ent.value, nil
 76}
 77
 78// Delete removes a key-value pair from the KeyStore.
 79func (ks *KeyStore[T]) Delete(key string) {
 80	ks.mu.Lock()
 81	defer ks.mu.Unlock()
 82
 83	delete(ks.data, key)
 84}
 85
 86// Keys returns an iterator for all not-expired keys in the KeyStore.
 87func (ks *KeyStore[T]) Keys() iter.Seq[string] {
 88	ks.Clean()
 89
 90	ks.mu.RLock()
 91	defer ks.mu.RUnlock()
 92	return maps.Keys(ks.data)
 93}
 94
 95// Clean deletes all expired keys from the keystore.
 96func (ks *KeyStore[T]) Clean() {
 97	now := time.Now()
 98	ks.mu.Lock()
 99	defer ks.mu.Unlock()
100
101	for k, ent := range ks.data {
102		if ent.expiration != nil && ent.expiration.Before(now) {
103			delete(ks.data, k)
104		}
105	}
106}
107
108// StartCleanup starts a goroutine that periodically deletes all expired keys from the keystore.
109func (ks *KeyStore[T]) StartCleanup(cleanupInterval time.Duration) {
110	if cleanupInterval == 0 {
111		return
112	}
113
114	ks.StopCleanup()
115	ks.cleanupInterval = cleanupInterval
116
117	ticker := time.NewTicker(ks.cleanupInterval)
118	for {
119		select {
120		case <-ticker.C:
121			ks.Clean()
122		case <-ks.stopChan:
123			ticker.Stop()
124			return
125		}
126	}
127}
128
129// StopCleanup stops the cleanup goroutine.
130func (ks *KeyStore[T]) StopCleanup() {
131	ks.mu.Lock()
132	defer ks.mu.Unlock()
133
134	if ks.cleanupInterval != 0 {
135		close(ks.stopChan)
136		ks.cleanupInterval = 0
137	}
138}