managed_storage.go

package payloadstorage

import (
	"context"
	"sync"
	"time"

	"decentralized-api/logging"

	"github.com/productscience/inference/x/inference/types"
)

const (
	defaultManagedCacheSize = 1000
	maxPruneLookback        = 10
)

type cachedEntry struct {
	promptPayload   []byte
	responsePayload []byte
	expiresAt       time.Time
}

// ManagedStorage wraps PayloadStorage with read caching and automatic epoch pruning.
// - Caches Retrieve results to reduce disk I/O during validation bursts
// - Automatically prunes old epochs in background (only last 10 epochs, older data requires manual prune)
type ManagedStorage struct {
	storage      PayloadStorage
	retainCount  uint64
	cacheTTL     time.Duration
	maxCacheSize int

	mu        sync.RWMutex
	cache     map[string]*cachedEntry
	maxEpoch  uint64
	minPruned uint64
}

func NewManagedStorage(storage PayloadStorage, retainCount uint64, cacheTTL time.Duration) *ManagedStorage {
	return NewManagedStorageWithSize(storage, retainCount, cacheTTL, defaultManagedCacheSize)
}

func NewManagedStorageWithSize(storage PayloadStorage, retainCount uint64, cacheTTL time.Duration, maxCacheSize int) *ManagedStorage {
	m := &ManagedStorage{
		storage:      storage,
		retainCount:  retainCount,
		cacheTTL:     cacheTTL,
		maxCacheSize: maxCacheSize,
		cache:        make(map[string]*cachedEntry),
	}
	go m.cleanupLoop()
	return m
}

func (m *ManagedStorage) Store(ctx context.Context, inferenceId string, epochId uint64, promptPayload, responsePayload []byte) error {
	if err := m.storage.Store(ctx, inferenceId, epochId, promptPayload, responsePayload); err != nil {
		return err
	}
	m.mu.Lock()
	if epochId > m.maxEpoch {
		m.maxEpoch = epochId
	}
	m.mu.Unlock()
	return nil
}

func (m *ManagedStorage) Retrieve(ctx context.Context, inferenceId string, epochId uint64) ([]byte, []byte, error) {
	m.mu.RLock()
	if c, ok := m.cache[inferenceId]; ok && time.Now().Before(c.expiresAt) {
		m.mu.RUnlock()
		return c.promptPayload, c.responsePayload, nil
	}
	m.mu.RUnlock()

	prompt, response, err := m.storage.Retrieve(ctx, inferenceId, epochId)
	if err != nil {
		return nil, nil, err
	}

	m.mu.Lock()
	m.cache[inferenceId] = &cachedEntry{
		promptPayload:   prompt,
		responsePayload: response,
		expiresAt:       time.Now().Add(m.cacheTTL),
	}
	m.mu.Unlock()

	return prompt, response, nil
}

func (m *ManagedStorage) PruneEpoch(ctx context.Context, epochId uint64) error {
	return m.storage.PruneEpoch(ctx, epochId)
}

func (m *ManagedStorage) cleanupLoop() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		m.cleanup()
	}
}

func (m *ManagedStorage) cleanup() {
	m.mu.Lock()

	now := time.Now()
	for id, c := range m.cache {
		if now.After(c.expiresAt) {
			delete(m.cache, id)
		}
	}

	for len(m.cache) > m.maxCacheSize {
		for key := range m.cache {
			delete(m.cache, key)
			break
		}
	}

	var start, threshold uint64
	if m.maxEpoch > m.retainCount {
		threshold = m.maxEpoch - m.retainCount

		if m.minPruned+maxPruneLookback < threshold {
			m.minPruned = threshold - maxPruneLookback
		}
		start = m.minPruned
	}
	m.mu.Unlock()

	m.pruneEpochs(start, threshold)
}

// pruneEpochs prunes epochs [start, threshold) in order, advancing minPruned
// only past epochs that pruned successfully so a failed epoch is retried on
// the next cleanup tick instead of being skipped forever (#850). Pruning runs
// outside m.mu so storage I/O does not block cache reads; cleanupLoop is the
// only caller, so prunes never overlap.
func (m *ManagedStorage) pruneEpochs(start, threshold uint64) {
	for epoch := start; epoch < threshold; epoch++ {
		if err := m.storage.PruneEpoch(context.Background(), epoch); err != nil {
			logging.Warn("Auto-prune failed, will retry on next cleanup", types.PayloadStorage, "epochId", epoch, "error", err)
			return
		}
		logging.Info("Auto-pruned epoch", types.PayloadStorage, "epochId", epoch)

		m.mu.Lock()
		m.minPruned = epoch + 1
		m.mu.Unlock()
	}
}

var _ PayloadStorage = (*ManagedStorage)(nil)