golib/gcache/memory.go

//
// memory.go
// Copyright (C) 2023 tiglog <me@tiglog.com>
//
// Distributed under terms of the MIT license.
//

package gcache

import (
	"runtime"
	"sync"
	"time"
)

func NewMemCache(opts ...ICacheOption) ICache {
	conf := NewConfig()
	for _, opt := range opts {
		opt(conf)
	}

	c := &memCache{
		shards:    make([]*memCacheShard, conf.shards),
		closed:    make(chan struct{}),
		shardMask: uint64(conf.shards - 1),
		config:    conf,
		hash:      conf.hash,
	}
	for i := 0; i < len(c.shards); i++ {
		c.shards[i] = newMemCacheShard(conf)
	}
	if conf.clearInterval > 0 {
		go func() {
			ticker := time.NewTicker(conf.clearInterval)
			defer ticker.Stop()
			for {
				select {
				case <-ticker.C:
					for _, shard := range c.shards {
						shard.checkExpire()
					}
				case <-c.closed:
					return
				}
			}
		}()
	}
	cache := &MemCache{c}
	// Associated finalizer function with obj.
	// When the obj is unreachable, close the obj.
	runtime.SetFinalizer(cache, func(cache *MemCache) { close(cache.closed) })
	return cache
}

type MemCache struct {
	*memCache
}

type memCache struct {
	shards    []*memCacheShard
	hash      IHash
	shardMask uint64
	config    *Config
	closed    chan struct{}
}

func (c *memCache) Set(k string, v interface{}, opts ...SetIOption) bool {
	item := Item{v: v}
	for _, opt := range opts {
		if pass := opt(c, k, &item); !pass {
			return false
		}
	}
	hashedKey := c.hash.Sum64(k)
	shard := c.getShard(hashedKey)
	shard.set(k, &item)
	return true
}

func (c *memCache) Get(k string) (interface{}, bool) {
	hashedKey := c.hash.Sum64(k)
	shard := c.getShard(hashedKey)
	return shard.get(k)
}

func (c *memCache) GetSet(k string, v interface{}, opts ...SetIOption) (interface{}, bool) {
	defer c.Set(k, v, opts...)
	return c.Get(k)
}

func (c *memCache) GetDel(k string) (interface{}, bool) {
	defer c.Del(k)
	return c.Get(k)
}

func (c *memCache) Del(ks ...string) int {
	var count int
	for _, k := range ks {
		hashedKey := c.hash.Sum64(k)
		shard := c.getShard(hashedKey)
		count += shard.del(k)
	}
	return count
}

// DelExpired Only delete when key expires
func (c *memCache) DelExpired(k string) bool {
	hashedKey := c.hash.Sum64(k)
	shard := c.getShard(hashedKey)
	return shard.delExpired(k)
}

func (c *memCache) Exists(ks ...string) bool {
	for _, k := range ks {
		if _, found := c.Get(k); !found {
			return false
		}
	}
	return true
}

func (c *memCache) Expire(k string, d time.Duration) bool {
	v, found := c.Get(k)
	if !found {
		return false
	}
	return c.Set(k, v, WithEx(d))
}

func (c *memCache) ExpireAt(k string, t time.Time) bool {
	v, found := c.Get(k)
	if !found {
		return false
	}
	return c.Set(k, v, WithExAt(t))
}

func (c *memCache) Persist(k string) bool {
	v, found := c.Get(k)
	if !found {
		return false
	}
	return c.Set(k, v)
}

func (c *memCache) Ttl(k string) (time.Duration, bool) {
	hashedKey := c.hash.Sum64(k)
	shard := c.getShard(hashedKey)
	return shard.ttl(k)
}

func (c *memCache) getShard(hashedKey uint64) (shard *memCacheShard) {
	return c.shards[hashedKey&c.shardMask]
}

//////////////
//   shard  //
//////////////

type memCacheShard struct {
	hashmap         map[string]Item
	lock            sync.RWMutex
	expiredCallback ExpiredCallback
}

func newMemCacheShard(conf *Config) *memCacheShard {
	return &memCacheShard{expiredCallback: conf.expiredCallback, hashmap: map[string]Item{}}
}

func (c *memCacheShard) set(k string, item *Item) {
	c.lock.Lock()
	c.hashmap[k] = *item
	c.lock.Unlock()
	return
}

func (c *memCacheShard) get(k string) (interface{}, bool) {
	c.lock.RLock()
	item, exist := c.hashmap[k]
	c.lock.RUnlock()
	if !exist {
		return nil, false
	}
	if !item.Expired() {
		return item.v, true
	}
	if c.delExpired(k) {
		return nil, false
	}
	return c.get(k)
}

func (c *memCacheShard) getSet(k string, item *Item) (interface{}, bool) {
	defer c.set(k, item)
	return c.get(k)
}

func (c *memCacheShard) getDel(k string) (interface{}, bool) {
	defer c.del(k)
	return c.get(k)
}

func (c *memCacheShard) del(k string) int {
	var count int
	c.lock.Lock()
	v, found := c.hashmap[k]
	if found {
		delete(c.hashmap, k)
		if !v.Expired() {
			count++
		}
	}
	c.lock.Unlock()
	return count
}

// delExpired Only delete when key expires
func (c *memCacheShard) delExpired(k string) bool {
	c.lock.Lock()
	item, found := c.hashmap[k]
	if !found || !item.Expired() {
		c.lock.Unlock()
		return false
	}
	delete(c.hashmap, k)
	c.lock.Unlock()
	if c.expiredCallback != nil {
		_ = c.expiredCallback(k, item.v)
	}
	return true
}

func (c *memCacheShard) ttl(k string) (time.Duration, bool) {
	c.lock.RLock()
	v, found := c.hashmap[k]
	c.lock.RUnlock()
	if !found || !v.CanExpire() || v.Expired() {
		return 0, false
	}
	return v.expire.Sub(time.Now()), true
}

func (c *memCacheShard) checkExpire() {
	var expiredKeys []string
	c.lock.RLock()
	for k, item := range c.hashmap {
		if item.Expired() {
			expiredKeys = append(expiredKeys, k)
		}
	}
	c.lock.RUnlock()
	for _, k := range expiredKeys {
		c.delExpired(k)
	}
}