moteva/server/internal/application/asset_cleanup.go

package application

import (
	"context"
	"encoding/json"
	"fmt"
	"net/url"
	"regexp"
	"strings"
	"time"

	"img_infinite_canvas/internal/domain/design"

	"github.com/zeromicro/go-zero/core/logx"
)

const (
	assetCleanupWorkers    = 5
	assetCleanupQueueSize  = 4096
	assetCleanupMaxRetries = 3
	assetCleanupTimeout    = 60 * time.Second
)

var (
	promptImageTokenPattern     = regexp.MustCompile(`\[@image:#\d+:[^:\]]+:(?:\[[^\]]+\]\(([^)]+)\)|([^\]]+))\]`)
	promptImageDirectivePattern = regexp.MustCompile(`(?i)(?:参考图|Reference image)\s*\d?\s*(?:[（(][^）)\n]+[）)])?\s*[:：]\s*(https?://[^\s\]]+)`)
	messageImageURLPattern      = regexp.MustCompile(`(?i)(https?://[^\s"'<>)\]]*\.(?:png|jpe?g|webp|gif|svg)(?:\?[^\s"'<>)\]]*)?|/[^\s"'<>)\]]*\.(?:png|jpe?g|webp|gif|svg)(?:\?[^\s"'<>)\]]*)?)`)
)

type assetCleanupQueue struct {
	storage    design.AssetStorage
	jobs       chan string
	maxRetries int
}

func newAssetCleanupQueue(storage design.AssetStorage) *assetCleanupQueue {
	if storage == nil {
		return nil
	}
	queue := &assetCleanupQueue{
		storage:    storage,
		jobs:       make(chan string, assetCleanupQueueSize),
		maxRetries: assetCleanupMaxRetries,
	}
	for i := 0; i < assetCleanupWorkers; i++ {
		go queue.run()
	}
	return queue
}

func (s *DesignService) enqueueAssetDeletion(publicURLs []string) {
	publicURLs = normalizeAssetURLs(publicURLs)
	if len(publicURLs) == 0 {
		return
	}
	if s.jobs != nil {
		if err := s.jobs.Enqueue(context.Background(), design.Job{
			Kind:       design.JobAssetCleanup,
			PublicURLs: publicURLs,
		}); err == nil {
			return
		} else {
			logx.Errorf("enqueue asset cleanup failed: %v", err)
		}
	}
	if s.assetCleanup == nil {
		return
	}
	s.assetCleanup.enqueue(publicURLs)
}

func (q *assetCleanupQueue) enqueue(publicURLs []string) {
	for _, publicURL := range normalizeAssetURLs(publicURLs) {
		select {
		case q.jobs <- publicURL:
		default:
			go q.enqueueSlow(publicURL)
		}
	}
}

func (q *assetCleanupQueue) enqueueSlow(publicURL string) {
	select {
	case q.jobs <- publicURL:
	case <-time.After(2 * time.Second):
		logx.Errorf("asset cleanup queue is full, dropping delete job for %s", publicURL)
	}
}

func (q *assetCleanupQueue) run() {
	for publicURL := range q.jobs {
		if err := q.deleteWithRetry(publicURL); err != nil {
			logx.Errorf("asset cleanup failed for %s: %v", publicURL, err)
		}
	}
}

func (q *assetCleanupQueue) deleteWithRetry(publicURL string) error {
	var err error
	for attempt := 1; attempt <= q.maxRetries; attempt++ {
		ctx, cancel := context.WithTimeout(context.Background(), assetCleanupTimeout)
		err = q.storage.DeleteObject(ctx, publicURL)
		cancel()
		if err == nil {
			return nil
		}
		time.Sleep(time.Duration(attempt*attempt) * 200 * time.Millisecond)
	}
	return err
}

func (s *DesignService) deleteAssetURLsWithRetry(publicURLs []string) error {
	if s.assets == nil {
		return nil
	}
	var failures []string
	for _, publicURL := range normalizeAssetURLs(publicURLs) {
		queue := &assetCleanupQueue{
			storage:    s.assets,
			maxRetries: assetCleanupMaxRetries,
		}
		if err := queue.deleteWithRetry(publicURL); err != nil {
			failures = append(failures, fmt.Sprintf("%s: %v", publicURL, err))
		}
	}
	if len(failures) > 0 {
		return fmt.Errorf("asset cleanup failed: %s", strings.Join(failures, "; "))
	}
	return nil
}

func (s *DesignService) assetReferencedByProjects(ctx context.Context, publicURL string) (bool, error) {
	if strings.TrimSpace(publicURL) == "" {
		return false, nil
	}
	summaries, err := s.repo.List(ctx)
	if err != nil {
		return false, err
	}
	for _, summary := range summaries {
		project, err := s.repo.Get(ctx, summary.ID)
		if err != nil {
			continue
		}
		if projectReferencesAssetURL(project, publicURL) {
			return true, nil
		}
	}
	return false, nil
}

func normalizeAssetURLs(publicURLs []string) []string {
	seen := make(map[string]struct{})
	urls := make([]string, 0, len(publicURLs))
	for _, publicURL := range publicURLs {
		publicURL = strings.TrimSpace(publicURL)
		if publicURL == "" {
			continue
		}
		if _, ok := seen[publicURL]; ok {
			continue
		}
		seen[publicURL] = struct{}{}
		urls = append(urls, publicURL)
	}
	return urls
}

func projectReferencesAssetURL(project design.Project, publicURL string) bool {
	protectedURLs := canvasAssetURLSet(project.Nodes)
	for _, url := range canvasSnapshotAssetURLs(project.Canvas) {
		addAssetURLKeys(protectedURLs, url)
	}
	for _, message := range project.Messages {
		for _, url := range messageAssetURLs(message) {
			addAssetURLKeys(protectedURLs, url)
		}
	}
	return assetURLSetContains(protectedURLs, publicURL)
}

func deletedCanvasAssetURLs(previous []design.Node, next []design.Node, messages []design.Message) []string {
	protectedURLs := canvasAssetURLSet(next)
	for _, message := range messages {
		for _, url := range messageAssetURLs(message) {
			addAssetURLKeys(protectedURLs, url)
		}
	}
	seen := make(map[string]struct{})
	var deleted []string
	for _, node := range previous {
		for _, url := range nodeAssetURLs(node) {
			if assetURLSetContains(protectedURLs, url) {
				continue
			}
			if _, ok := seen[url]; ok {
				continue
			}
			seen[url] = struct{}{}
			deleted = append(deleted, url)
		}
	}
	return deleted
}

func projectAssetURLs(project design.Project) []string {
	seen := make(map[string]struct{})
	var urls []string
	add := func(url string) {
		url = strings.TrimSpace(url)
		if url == "" {
			return
		}
		if _, ok := seen[url]; ok {
			return
		}
		seen[url] = struct{}{}
		urls = append(urls, url)
	}
	for _, node := range project.Nodes {
		for _, url := range nodeAssetURLs(node) {
			add(url)
		}
	}
	for _, url := range canvasSnapshotAssetURLs(project.Canvas) {
		add(url)
	}
	for _, message := range project.Messages {
		for _, url := range messageAssetURLs(message) {
			add(url)
		}
	}
	return urls
}

func messageAssetURLs(message design.Message) []string {
	var urls []string
	content := strings.TrimSpace(message.Content)
	if content == "" {
		return nil
	}
	var events []struct {
		EventData struct {
			Artifact []struct {
				ImageURL string `json:"image_url"`
			} `json:"artifact"`
		} `json:"event_data"`
	}
	if err := json.Unmarshal([]byte(content), &events); err == nil {
		for _, event := range events {
			for _, artifact := range event.EventData.Artifact {
				if strings.TrimSpace(artifact.ImageURL) != "" {
					urls = append(urls, artifact.ImageURL)
				}
			}
		}
	}
	urls = append(urls, promptMessageAssetURLs(content)...)
	urls = append(urls, looseMessageImageURLs(content)...)
	return urls
}

func canvasSnapshotAssetURLs(canvas string) []string {
	snapshot, err := design.DecodeCanvasSnapshot(canvas)
	if err != nil {
		return nil
	}
	urls := make([]string, 0, len(snapshot.Nodes))
	for _, node := range snapshot.Nodes {
		urls = append(urls, nodeAssetURLs(node)...)
	}
	return urls
}

func canvasAssetURLSet(nodes []design.Node) map[string]struct{} {
	items := make(map[string]struct{})
	for _, node := range nodes {
		for _, url := range nodeAssetURLs(node) {
			addAssetURLKeys(items, url)
		}
	}
	return items
}

func nodeAssetURLs(node design.Node) []string {
	if node.Type != design.NodeTypeImage && node.Type != design.NodeTypeFrame {
		return nil
	}
	urls := make([]string, 0, 2)
	for _, content := range []string{node.Content, node.RenderContent} {
		content = strings.TrimSpace(content)
		if content == "" || strings.HasPrefix(content, "data:image/") || !isGeneratedImageContent(content) {
			continue
		}
		urls = append(urls, content)
	}
	return normalizeAssetURLs(urls)
}

func promptMessageAssetURLs(content string) []string {
	var urls []string
	for _, match := range promptImageTokenPattern.FindAllStringSubmatch(content, -1) {
		if len(match) > 1 && strings.TrimSpace(match[1]) != "" {
			urls = append(urls, strings.TrimSpace(match[1]))
			continue
		}
		if len(match) > 2 && strings.TrimSpace(match[2]) != "" {
			urls = append(urls, strings.TrimSpace(match[2]))
		}
	}
	for _, match := range promptImageDirectivePattern.FindAllStringSubmatch(content, -1) {
		if len(match) > 1 && strings.TrimSpace(match[1]) != "" {
			urls = append(urls, strings.TrimSpace(match[1]))
		}
	}
	return urls
}

func looseMessageImageURLs(content string) []string {
	matches := messageImageURLPattern.FindAllString(content, -1)
	urls := make([]string, 0, len(matches))
	for _, match := range matches {
		url := strings.TrimRight(strings.TrimSpace(match), ".,，。;；")
		if url != "" {
			urls = append(urls, url)
		}
	}
	return urls
}

func assetURLSetContains(items map[string]struct{}, publicURL string) bool {
	for key := range assetURLKeys(publicURL) {
		if _, ok := items[key]; ok {
			return true
		}
	}
	return false
}

func addAssetURLKeys(items map[string]struct{}, publicURL string) {
	for key := range assetURLKeys(publicURL) {
		items[key] = struct{}{}
	}
}

func assetURLKeys(publicURL string) map[string]struct{} {
	publicURL = strings.TrimSpace(publicURL)
	keys := make(map[string]struct{})
	if publicURL == "" || strings.HasPrefix(publicURL, "data:image/") || strings.HasPrefix(publicURL, "blob:") {
		return keys
	}
	keys[publicURL] = struct{}{}
	if queryIndex := strings.Index(publicURL, "?"); queryIndex > 0 {
		keys[publicURL[:queryIndex]] = struct{}{}
	}
	parsed, err := url.Parse(publicURL)
	if err == nil && parsed.Path != "" {
		keys[parsed.Path] = struct{}{}
		pathWithoutQuery := parsed.Path
		if parsed.RawQuery != "" {
			keys[pathWithoutQuery+"?"+parsed.RawQuery] = struct{}{}
		}
	}
	return keys
}