nettools/data_structures/cache.go

package dsUtils

import (
	"context"
	"sync"
	"time"
)

type KeyValue struct {
	Key   string
	Value []byte
}

type MaybeExpirableData struct {
	Data           []byte
	ExpirationTime time.Time
	CanExpire      bool
}
type ToBeInvalidatedData struct {
	Key            string
	ExpirationTime time.Time
}

type ICache interface {
	Get(key string) (item []byte, isSet bool)
	Set(data map[string][]byte, canExpire bool, updateExpirationTime bool)
	Remove(key string)
	SetIfNotExists(key string, value []byte, canExpire bool) (isSet bool)
	GetAndRemove(key string) (item []byte, isSet bool)

	Run(ctx context.Context) error
}

// Cache provides simple cache functionality
// All objects in the same cache instance are getting same TTL
type Cache struct {
	ICache
	newExpirableItem *sync.Cond
	lock             *sync.Mutex
	data             map[string]MaybeExpirableData
	invalidateOrder  []ToBeInvalidatedData
	ttl              time.Duration
}

// NewCache Creates new Cache instance
func NewCache(ttl time.Duration) *Cache {
	return &Cache{
		ttl:              ttl,
		newExpirableItem: &sync.Cond{L: &sync.Mutex{}},
		lock:             &sync.Mutex{},
		data:             make(map[string]MaybeExpirableData),
		invalidateOrder:  make([]ToBeInvalidatedData, 0),
	}
}

func (c *Cache) get(key string) (item MaybeExpirableData, isSet bool) {
	item, isSet = c.data[key]
	if item.ExpirationTime.Before(time.Now()) {
		delete(c.data, key)
		return MaybeExpirableData{}, false
	}
	return item, isSet
}

func (c *Cache) Get(key string) (data []byte, isSet bool) {
	c.lock.Lock()
	defer c.lock.Unlock()
	item, isSet := c.get(key)
	return item.Data, isSet
}

func (c *Cache) GetAndRemove(key string) (data []byte, isSet bool) {
	c.lock.Lock()
	item, isSet := c.get(key)
	delete(c.data, key)
	c.lock.Unlock()
	return item.Data, isSet
}

// Internal function
func (c *Cache) set(key string, value []byte, canExpire bool, expirationTime time.Time) {
	c.data[key] = MaybeExpirableData{
		Data:           value,
		ExpirationTime: expirationTime,
		CanExpire:      canExpire,
	}
	if canExpire {
		c.invalidateOrder = append(c.invalidateOrder, ToBeInvalidatedData{
			Key:            key,
			ExpirationTime: expirationTime,
		})
		c.newExpirableItem.Signal()
	}
}

func (c *Cache) Set(data map[string][]byte, canExpire bool, updateExpirationTime bool) {
	c.lock.Lock()
	for key, val := range data {
		if !updateExpirationTime {
			if item, exists := c.data[key]; exists {
				c.set(key, val, item.CanExpire, item.ExpirationTime)
				return
			}
		}
		c.set(key, val, canExpire, time.Now().Add(c.ttl))
	}
	c.lock.Unlock()
}

func (c *Cache) Remove(key string) {
	c.lock.Lock()
	delete(c.data, key)
	c.lock.Unlock()
}

func (c *Cache) SetIfNotExists(key string, value []byte, canExpire bool) (isSet bool) {
	c.lock.Lock()
	defer c.lock.Unlock()
	_, exists := c.data[key]
	if exists {
		return false
	}
	c.set(key, value, canExpire, time.Now().Add(c.ttl))
	return true
}

func (c *Cache) timeInvalidate() {
	c.lock.Lock()
	if len(c.invalidateOrder) == 0 {
		c.lock.Unlock()
		c.newExpirableItem.Wait()
		c.lock.Lock()
	}
	toBeInvalidated := c.invalidateOrder[0]
	c.lock.Unlock()
	time.After(time.Now().Sub(toBeInvalidated.ExpirationTime))
	c.lock.Lock()
	item, exists := c.data[toBeInvalidated.Key]
	if exists && item.CanExpire {
		if item.ExpirationTime.After(time.Now()) {
			delete(c.data, toBeInvalidated.Key)
		}
	}
	c.invalidateOrder = c.invalidateOrder[1:]
	c.lock.Unlock()
}

func (c *Cache) Run(ctx context.Context) error {
	// Invalidator loop
	go func() {
		for {
			c.timeInvalidate()
		}
	}()
	<-ctx.Done()
	return ctx.Err()
}