nex/backend/internal/service/routing_cache.go

package service

import (
	"strings"
	"sync"

	"nex/backend/internal/domain"
	"nex/backend/internal/repository"

	"go.uber.org/zap"

	pkglogger "nex/backend/pkg/logger"
)

type RoutingCache struct {
	providers sync.Map
	models    sync.Map

	modelRepo    repository.ModelRepository
	providerRepo repository.ProviderRepository
	logger       *zap.Logger
}

func NewRoutingCache(
	modelRepo repository.ModelRepository,
	providerRepo repository.ProviderRepository,
	logger *zap.Logger,
) *RoutingCache {
	return &RoutingCache{
		modelRepo:    modelRepo,
		providerRepo: providerRepo,
		logger:       pkglogger.WithModule(logger, "service.routing_cache"),
	}
}

func (c *RoutingCache) GetProvider(id string) (*domain.Provider, error) {
	if v, ok := c.providers.Load(id); ok {
		if provider, ok := v.(*domain.Provider); ok {
			return provider, nil
		}
	}

	provider, err := c.providerRepo.GetByID(id)
	if err != nil {
		return nil, err
	}

	if v, ok := c.providers.Load(id); ok {
		if provider, ok := v.(*domain.Provider); ok {
			return provider, nil
		}
	}

	c.providers.Store(id, provider)
	return provider, nil
}

func (c *RoutingCache) GetModel(providerID, modelName string) (*domain.Model, error) {
	key := providerID + "/" + modelName

	if v, ok := c.models.Load(key); ok {
		if model, ok := v.(*domain.Model); ok {
			return model, nil
		}
	}

	model, err := c.modelRepo.FindByProviderAndModelName(providerID, modelName)
	if err != nil {
		return nil, err
	}

	if v, ok := c.models.Load(key); ok {
		if model, ok := v.(*domain.Model); ok {
			return model, nil
		}
	}

	c.models.Store(key, model)
	return model, nil
}

func (c *RoutingCache) SetProvider(provider *domain.Provider) {
	c.providers.Store(provider.ID, provider)
}

func (c *RoutingCache) SetModel(model *domain.Model) {
	key := model.ProviderID + "/" + model.ModelName
	c.models.Store(key, model)
}

func (c *RoutingCache) InvalidateProvider(id string) {
	c.providers.Delete(id)
	c.invalidateModelsByProvider(id)
	c.logger.Debug("Provider 缓存失效", zap.String("provider_id", id))
}

func (c *RoutingCache) InvalidateModel(providerID, modelName string) {
	key := providerID + "/" + modelName
	c.models.Delete(key)
	c.logger.Debug("Model 缓存失效",
		zap.String("provider_id", providerID),
		zap.String("model_name", modelName))
}

func (c *RoutingCache) invalidateModelsByProvider(providerID string) {
	prefix := providerID + "/"
	count := 0
	c.models.Range(func(key, value interface{}) bool {
		keyStr, ok := key.(string)
		if !ok {
			return true
		}

		if strings.HasPrefix(keyStr, prefix) {
			c.models.Delete(key)
			count++
		}
		return true
	})
	if count > 0 {
		c.logger.Debug("清除 Provider 相关 Model 缓存",
			zap.String("provider_id", providerID),
			zap.Int("count", count))
	}
}

func (c *RoutingCache) Preload() error {
	providers, err := c.providerRepo.List()
	if err != nil {
		return err
	}
	for i := range providers {
		c.providers.Store(providers[i].ID, &providers[i])
	}

	models, err := c.modelRepo.List("")
	if err != nil {
		return err
	}
	for i := range models {
		key := models[i].ProviderID + "/" + models[i].ModelName
		c.models.Store(key, &models[i])
	}

	c.logger.Info("缓存预热完成",
		zap.Int("providers", len(providers)),
		zap.Int("models", len(models)))
	return nil
}