golib/gcache/memory_adapter.go

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

package gcache

import (
	"sync"
	"time"
)

// Option defines options for creating a memory store
type MCOption func(i *InMemoryStore)

// BufferSize configures the store to allow a maximum of n
func BufferSize(n int) MCOption {
	return func(i *InMemoryStore) {
		i.bufferSize = n
	}
}

// KeyFunc allows for dynamic generation of cache keys
func MCCacheKeyGenerator(fn KeyFunc) MCOption {
	return func(i *InMemoryStore) {
		i.keyfn = fn
	}
}

// New returns a configured in memory store
func NewMemoryStore(opts ...MCOption) *InMemoryStore {
	i := &InMemoryStore{}

	for _, opt := range opts {
		opt(i)
	}

	if i.keyfn == nil {
		i.keyfn = DefaultKeyFunc
	}

	if i.data == nil {
		if i.bufferSize == 0 {
			i.bufferSize = 100
		}

		i.data = make(map[string]*Item, i.bufferSize)
	}

	return i
}

// Represents an in-memory store
type InMemoryStore struct {
	lock sync.RWMutex
	data map[string]*Item

	bufferSize int
	keyfn      KeyFunc
}

func (i *InMemoryStore) Set(key string, data []byte, expires time.Duration) error {
	i.lock.Lock()

	i.data[i.keyfn(key)] = &Item{
		ExpiresAt: time.Now().Add(expires),
		Data:      copyData(data),
	}

	i.lock.Unlock()
	return nil
}

func (i *InMemoryStore) Get(key string) ([]byte, error) {
	i.lock.RLock()

	item := i.data[i.keyfn(key)]
	if item == nil {
		i.lock.RUnlock()
		return nil, ErrCacheMiss
	}

	if item.IsExpired() {
		i.lock.RUnlock()
		i.Delete(key)
		return nil, ErrCacheMiss
	}

	i.lock.RUnlock()
	return copyData(item.Data), nil
}

func (i *InMemoryStore) Delete(key string) error {
	i.lock.RLock()

	_, ok := i.data[i.keyfn(key)]
	if !ok {
		i.lock.RUnlock()
		return ErrCacheMiss
	}

	i.lock.RUnlock()

	i.lock.Lock()
	delete(i.data, i.keyfn(key))
	i.lock.Unlock()
	return nil
}

func (i *InMemoryStore) Flush() error {
	i.lock.Lock()

	i.data = make(map[string]*Item, i.bufferSize)
	i.lock.Unlock()
	return nil
}

func (i *InMemoryStore) Has(key string) bool {
	i.lock.RLock()

	_, ok := i.data[i.keyfn(key)]
	i.lock.RUnlock()
	return ok
}

func (i *InMemoryStore) GC() {
	i.lock.Lock()

	for k, item := range i.data {
		if item.IsExpired() {
			//No need to spawn a new goroutine since we
			//still have the lock here
			delete(i.data, k)
		}
	}

	i.lock.Unlock()
}

func (i *InMemoryStore) count() int {
	i.lock.Lock()
	n := len(i.data)
	i.lock.Unlock()

	return n
}

func copyData(data []byte) []byte {
	result := make([]byte, len(data))
	copy(result, data)

	return result
}