昔念
699db8406b
feat(base): 添加邮箱注册码功能及用户注册接口 - 在 `sessionManager` 中新增邮件注册码缓存管理实例和相关方法 - 实现生成、保存、验证、删除邮件注册码的逻辑 - 新增 `/reg` 和 `/email` 接口用于用户注册和发送验证码 - 引入 `golang-lru` 依赖以支持限流缓存功能 - 调整包导入顺序，优化代码结构 ```
346 lines
8.7 KiB
Go
346 lines
8.7 KiB
Go
// Copyright IBM Corp. 2014, 2025
|
|
// SPDX-License-Identifier: MPL-2.0
|
|
|
|
package expirable
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/hashicorp/golang-lru/v2/internal"
|
|
)
|
|
|
|
// EvictCallback is used to get a callback when a cache entry is evicted
|
|
type EvictCallback[K comparable, V any] func(key K, value V)
|
|
|
|
// LRU implements a thread-safe LRU with expirable entries.
|
|
type LRU[K comparable, V any] struct {
|
|
size int
|
|
evictList *internal.LruList[K, V]
|
|
items map[K]*internal.Entry[K, V]
|
|
onEvict EvictCallback[K, V]
|
|
|
|
// expirable options
|
|
mu sync.Mutex
|
|
ttl time.Duration
|
|
done chan struct{}
|
|
|
|
// buckets for expiration
|
|
buckets []bucket[K, V]
|
|
// uint8 because it's number between 0 and numBuckets
|
|
nextCleanupBucket uint8
|
|
}
|
|
|
|
// bucket is a container for holding entries to be expired
|
|
type bucket[K comparable, V any] struct {
|
|
entries map[K]*internal.Entry[K, V]
|
|
newestEntry time.Time
|
|
}
|
|
|
|
// noEvictionTTL - very long ttl to prevent eviction
|
|
const noEvictionTTL = time.Hour * 24 * 365 * 10
|
|
|
|
// because of uint8 usage for nextCleanupBucket, should not exceed 256.
|
|
// casting it as uint8 explicitly requires type conversions in multiple places
|
|
const numBuckets = 100
|
|
|
|
// NewLRU returns a new thread-safe cache with expirable entries.
|
|
//
|
|
// Size parameter set to 0 makes cache of unlimited size, e.g. turns LRU mechanism off.
|
|
//
|
|
// Providing 0 TTL turns expiring off.
|
|
//
|
|
// Delete expired entries every 1/100th of ttl value. Goroutine which deletes expired entries runs indefinitely.
|
|
func NewLRU[K comparable, V any](size int, onEvict EvictCallback[K, V], ttl time.Duration) *LRU[K, V] {
|
|
if size < 0 {
|
|
size = 0
|
|
}
|
|
if ttl <= 0 {
|
|
ttl = noEvictionTTL
|
|
}
|
|
|
|
res := LRU[K, V]{
|
|
ttl: ttl,
|
|
size: size,
|
|
evictList: internal.NewList[K, V](),
|
|
items: make(map[K]*internal.Entry[K, V]),
|
|
onEvict: onEvict,
|
|
done: make(chan struct{}),
|
|
}
|
|
|
|
// initialize the buckets
|
|
res.buckets = make([]bucket[K, V], numBuckets)
|
|
for i := 0; i < numBuckets; i++ {
|
|
res.buckets[i] = bucket[K, V]{entries: make(map[K]*internal.Entry[K, V])}
|
|
}
|
|
|
|
// enable deleteExpired() running in separate goroutine for cache with non-zero TTL
|
|
//
|
|
// Important: done channel is never closed, so deleteExpired() goroutine will never exit,
|
|
// it's decided to add functionality to close it in the version later than v2.
|
|
if res.ttl != noEvictionTTL {
|
|
go func(done <-chan struct{}) {
|
|
ticker := time.NewTicker(res.ttl / numBuckets)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-done:
|
|
return
|
|
case <-ticker.C:
|
|
res.deleteExpired()
|
|
}
|
|
}
|
|
}(res.done)
|
|
}
|
|
return &res
|
|
}
|
|
|
|
// Purge clears the cache completely.
|
|
// onEvict is called for each evicted key.
|
|
func (c *LRU[K, V]) Purge() {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
for k, v := range c.items {
|
|
if c.onEvict != nil {
|
|
c.onEvict(k, v.Value)
|
|
}
|
|
delete(c.items, k)
|
|
}
|
|
for _, b := range c.buckets {
|
|
for _, ent := range b.entries {
|
|
delete(b.entries, ent.Key)
|
|
}
|
|
}
|
|
c.evictList.Init()
|
|
}
|
|
|
|
// Add adds a value to the cache. Returns true if an eviction occurred.
|
|
// Returns false if there was no eviction: the item was already in the cache,
|
|
// or the size was not exceeded.
|
|
func (c *LRU[K, V]) Add(key K, value V) (evicted bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
now := time.Now()
|
|
|
|
// Check for existing item
|
|
if ent, ok := c.items[key]; ok {
|
|
c.evictList.MoveToFront(ent)
|
|
c.removeFromBucket(ent) // remove the entry from its current bucket as expiresAt is renewed
|
|
ent.Value = value
|
|
ent.ExpiresAt = now.Add(c.ttl)
|
|
c.addToBucket(ent)
|
|
return false
|
|
}
|
|
|
|
// Add new item
|
|
ent := c.evictList.PushFrontExpirable(key, value, now.Add(c.ttl))
|
|
c.items[key] = ent
|
|
c.addToBucket(ent) // adds the entry to the appropriate bucket and sets entry.expireBucket
|
|
|
|
evict := c.size > 0 && c.evictList.Length() > c.size
|
|
// Verify size not exceeded
|
|
if evict {
|
|
c.removeOldest()
|
|
}
|
|
return evict
|
|
}
|
|
|
|
// Get looks up a key's value from the cache.
|
|
func (c *LRU[K, V]) Get(key K) (value V, ok bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
var ent *internal.Entry[K, V]
|
|
if ent, ok = c.items[key]; ok {
|
|
// Expired item check
|
|
if time.Now().After(ent.ExpiresAt) {
|
|
return value, false
|
|
}
|
|
c.evictList.MoveToFront(ent)
|
|
return ent.Value, true
|
|
}
|
|
return
|
|
}
|
|
|
|
// Contains checks if a key is in the cache, without updating the recent-ness
|
|
// or deleting it for being stale.
|
|
func (c *LRU[K, V]) Contains(key K) (ok bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
_, ok = c.items[key]
|
|
return ok
|
|
}
|
|
|
|
// Peek returns the key value (or undefined if not found) without updating
|
|
// the "recently used"-ness of the key.
|
|
func (c *LRU[K, V]) Peek(key K) (value V, ok bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
var ent *internal.Entry[K, V]
|
|
if ent, ok = c.items[key]; ok {
|
|
// Expired item check
|
|
if time.Now().After(ent.ExpiresAt) {
|
|
return value, false
|
|
}
|
|
return ent.Value, true
|
|
}
|
|
return
|
|
}
|
|
|
|
// Remove removes the provided key from the cache, returning if the
|
|
// key was contained.
|
|
func (c *LRU[K, V]) Remove(key K) bool {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if ent, ok := c.items[key]; ok {
|
|
c.removeElement(ent)
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// RemoveOldest removes the oldest item from the cache.
|
|
func (c *LRU[K, V]) RemoveOldest() (key K, value V, ok bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if ent := c.evictList.Back(); ent != nil {
|
|
c.removeElement(ent)
|
|
return ent.Key, ent.Value, true
|
|
}
|
|
return
|
|
}
|
|
|
|
// GetOldest returns the oldest entry
|
|
func (c *LRU[K, V]) GetOldest() (key K, value V, ok bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if ent := c.evictList.Back(); ent != nil {
|
|
return ent.Key, ent.Value, true
|
|
}
|
|
return
|
|
}
|
|
|
|
// Keys returns a slice of the keys in the cache, from oldest to newest.
|
|
// Expired entries are filtered out.
|
|
func (c *LRU[K, V]) Keys() []K {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
keys := make([]K, 0, len(c.items))
|
|
now := time.Now()
|
|
for ent := c.evictList.Back(); ent != nil; ent = ent.PrevEntry() {
|
|
if now.After(ent.ExpiresAt) {
|
|
continue
|
|
}
|
|
keys = append(keys, ent.Key)
|
|
}
|
|
return keys
|
|
}
|
|
|
|
// Values returns a slice of the values in the cache, from oldest to newest.
|
|
// Expired entries are filtered out.
|
|
func (c *LRU[K, V]) Values() []V {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
values := make([]V, 0, len(c.items))
|
|
now := time.Now()
|
|
for ent := c.evictList.Back(); ent != nil; ent = ent.PrevEntry() {
|
|
if now.After(ent.ExpiresAt) {
|
|
continue
|
|
}
|
|
values = append(values, ent.Value)
|
|
}
|
|
return values
|
|
}
|
|
|
|
// Len returns the number of items in the cache.
|
|
func (c *LRU[K, V]) Len() int {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
return c.evictList.Length()
|
|
}
|
|
|
|
// Resize changes the cache size. Size of 0 means unlimited.
|
|
func (c *LRU[K, V]) Resize(size int) (evicted int) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if size <= 0 {
|
|
c.size = 0
|
|
return 0
|
|
}
|
|
diff := c.evictList.Length() - size
|
|
if diff < 0 {
|
|
diff = 0
|
|
}
|
|
for i := 0; i < diff; i++ {
|
|
c.removeOldest()
|
|
}
|
|
c.size = size
|
|
return diff
|
|
}
|
|
|
|
// Close destroys cleanup goroutine. To clean up the cache, run Purge() before Close().
|
|
// func (c *LRU[K, V]) Close() {
|
|
// c.mu.Lock()
|
|
// defer c.mu.Unlock()
|
|
// select {
|
|
// case <-c.done:
|
|
// return
|
|
// default:
|
|
// }
|
|
// close(c.done)
|
|
// }
|
|
|
|
// removeOldest removes the oldest item from the cache. Has to be called with lock!
|
|
func (c *LRU[K, V]) removeOldest() {
|
|
if ent := c.evictList.Back(); ent != nil {
|
|
c.removeElement(ent)
|
|
}
|
|
}
|
|
|
|
// removeElement is used to remove a given list element from the cache. Has to be called with lock!
|
|
func (c *LRU[K, V]) removeElement(e *internal.Entry[K, V]) {
|
|
c.evictList.Remove(e)
|
|
delete(c.items, e.Key)
|
|
c.removeFromBucket(e)
|
|
if c.onEvict != nil {
|
|
c.onEvict(e.Key, e.Value)
|
|
}
|
|
}
|
|
|
|
// deleteExpired deletes expired records from the oldest bucket, waiting for the newest entry
|
|
// in it to expire first.
|
|
func (c *LRU[K, V]) deleteExpired() {
|
|
c.mu.Lock()
|
|
bucketIdx := c.nextCleanupBucket
|
|
timeToExpire := time.Until(c.buckets[bucketIdx].newestEntry)
|
|
// wait for newest entry to expire before cleanup without holding lock
|
|
if timeToExpire > 0 {
|
|
c.mu.Unlock()
|
|
time.Sleep(timeToExpire)
|
|
c.mu.Lock()
|
|
}
|
|
for _, ent := range c.buckets[bucketIdx].entries {
|
|
c.removeElement(ent)
|
|
}
|
|
c.nextCleanupBucket = (c.nextCleanupBucket + 1) % numBuckets
|
|
c.mu.Unlock()
|
|
}
|
|
|
|
// addToBucket adds entry to expire bucket so that it will be cleaned up when the time comes. Has to be called with lock!
|
|
func (c *LRU[K, V]) addToBucket(e *internal.Entry[K, V]) {
|
|
bucketID := (numBuckets + c.nextCleanupBucket - 1) % numBuckets
|
|
e.ExpireBucket = bucketID
|
|
c.buckets[bucketID].entries[e.Key] = e
|
|
if c.buckets[bucketID].newestEntry.Before(e.ExpiresAt) {
|
|
c.buckets[bucketID].newestEntry = e.ExpiresAt
|
|
}
|
|
}
|
|
|
|
// removeFromBucket removes the entry from its corresponding bucket. Has to be called with lock!
|
|
func (c *LRU[K, V]) removeFromBucket(e *internal.Entry[K, V]) {
|
|
delete(c.buckets[e.ExpireBucket].entries, e.Key)
|
|
}
|
|
|
|
// Cap returns the capacity of the cache
|
|
func (c *LRU[K, V]) Cap() int {
|
|
return c.size
|
|
} |