kra-new/internal/integration/cache/cache.go

package cache

import (
	"context"
	"kra/internal/biz/system"
	"strconv"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

type RedisProvider interface {
	RedisClient() redis.UniversalClient
}

type memoryEntry struct {
	value     string
	expiresAt time.Time
}

type Store struct {
	provider RedisProvider
	mu       sync.Mutex
	memory   map[string]memoryEntry
}

const maxMemoryEntries = 10000

func New(provider RedisProvider) system.Cache {
	return &Store{provider: provider, memory: make(map[string]memoryEntry)}
}

func (s *Store) client() redis.UniversalClient {
	if s == nil || s.provider == nil {
		return nil
	}
	return s.provider.RedisClient()
}

func (s *Store) Get(ctx context.Context, key string) (string, bool, error) {
	if client := s.client(); client != nil {
		value, err := client.Get(ctx, key).Result()
		if err == nil {
			return value, true, nil
		}
		if err != redis.Nil {
			return "", false, err
		}
		return "", false, nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	entry, ok := s.memory[key]
	if !ok {
		return "", false, nil
	}
	if !entry.expiresAt.IsZero() && time.Now().After(entry.expiresAt) {
		delete(s.memory, key)
		return "", false, nil
	}
	return entry.value, true, nil
}

func (s *Store) Set(ctx context.Context, key, value string, expiration time.Duration) error {
	if client := s.client(); client != nil {
		return client.Set(ctx, key, value, expiration).Err()
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.makeRoom(key)
	entry := memoryEntry{value: value}
	if expiration > 0 {
		entry.expiresAt = time.Now().Add(expiration)
	}
	s.memory[key] = entry
	return nil
}

func (s *Store) Delete(ctx context.Context, key string) error {
	if client := s.client(); client != nil {
		return client.Del(ctx, key).Err()
	}
	s.mu.Lock()
	delete(s.memory, key)
	s.mu.Unlock()
	return nil
}

func (s *Store) Increment(ctx context.Context, key string, expiration time.Duration) (int64, error) {
	if client := s.client(); client != nil {
		if expiration > 0 {
			expirationMS := expiration.Milliseconds()
			if expirationMS < 1 {
				expirationMS = 1
			}
			const incrementWithTTL = `
local value = redis.call('INCR', KEYS[1])
if redis.call('PTTL', KEYS[1]) < 0 then
  redis.call('PEXPIRE', KEYS[1], ARGV[1])
end
return value`
			return client.Eval(ctx, incrementWithTTL, []string{key}, expirationMS).Int64()
		}
		value, err := client.Incr(ctx, key).Result()
		if err != nil {
			return 0, err
		}
		return value, nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.makeRoom(key)
	entry, ok := s.memory[key]
	if ok && !entry.expiresAt.IsZero() && time.Now().After(entry.expiresAt) {
		ok = false
		entry = memoryEntry{}
	}
	value := int64(0)
	if ok {
		value, _ = strconv.ParseInt(entry.value, 10, 64)
	}
	value++
	entry.value = strconv.FormatInt(value, 10)
	if !ok {
		if expiration > 0 {
			entry.expiresAt = time.Now().Add(expiration)
		} else {
			entry.expiresAt = time.Time{}
		}
	} else if entry.expiresAt.IsZero() && expiration > 0 {
		entry.expiresAt = time.Now().Add(expiration)
	}
	s.memory[key] = entry
	return value, nil
}

func (s *Store) makeRoom(incoming string) {
	if len(s.memory) < maxMemoryEntries {
		return
	}
	now := time.Now()
	for key, entry := range s.memory {
		if !entry.expiresAt.IsZero() && now.After(entry.expiresAt) {
			delete(s.memory, key)
		}
	}
	if len(s.memory) < maxMemoryEntries {
		return
	}
	if _, exists := s.memory[incoming]; exists {
		return
	}
	for key := range s.memory {
		delete(s.memory, key)
		break
	}
}