spider/internal/media/upload.go

// Package media 负责「媒体机直传」：把图片下载下来、算内容 sha1、
// 向后端要预签名 PUT 地址、直接传到对象存储，最后只把 key 交给后端。
//
// 为什么由 spider 做这件事：
//   - spider 部署在带宽好的机器上，下载 Vogue 图 + 上传存储是带宽密集的活；
//   - 后端机器只需写库，不再被出口带宽卡住；
//   - 爬虫**不持有**存储密钥：预签名地址由后端签发（15 分钟有效），
//     密钥只留在后端，减少泄露面。
//
// key 采用内容寻址（runway/<sha1>.jpg），同一张图重复上传是幂等的，
// 重爬不会产生孤儿文件。
package media

import (
	"bytes"
	"context"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"my-spiders/internal/ingest"
)

// Uploader 媒体直传器。
type Uploader struct {
	Client      *ingest.Client // 复用它来签发预签名 URL（HMAC）
	Prefix      string         // 存储前缀：runway / street
	Concurrency int            // 下载与上传的并发数
	HTTP        *http.Client
}

// Result 单张图的直传结果；Err 非空表示这张失败（调用方应回退为上报 URL）。
type Result struct {
	Key  string
	Size int64
	Err  error
}

// New 创建直传器。concurrency <= 0 时取 6。
func New(client *ingest.Client, prefix string, concurrency int) *Uploader {
	if concurrency <= 0 {
		concurrency = 6
	}
	return &Uploader{
		Client:      client,
		Prefix:      prefix,
		Concurrency: concurrency,
		HTTP:        &http.Client{Timeout: 60 * time.Second},
	}
}

// UploadAll 并发完成「下载 → 生成 key → 预签名 → PUT」，返回与 urls 下标一一对应的结果。
// 单张失败不影响其他张：该下标 Err 非空，调用方据此对这张回退成 URL 上报。
func (u *Uploader) UploadAll(ctx context.Context, urls []string) []Result {
	results := make([]Result, len(urls))
	if len(urls) == 0 {
		return results
	}

	// ① 并发下载
	data := make([][]byte, len(urls))
	errs := make([]error, len(urls))
	u.parallel(len(urls), func(i int) {
		body, err := u.download(ctx, urls[i])
		if err != nil {
			errs[i] = err
			return
		}
		data[i] = body
	})
	log.Printf("[media] 下载完成 %d 张（失败 %d）", countOK(data), len(urls)-countOK(data))

	// ② 生成本批 key（内容 sha1 + 原扩展名）
	for i := range urls {
		if errs[i] != nil {
			results[i] = Result{Err: errs[i]}
			continue
		}
		results[i] = Result{
			Key:  u.Prefix + "/" + sha1Hex(data[i]) + extOf(urls[i]),
			Size: int64(len(data[i])),
		}
	}

	// ③ 批量申请预签名地址（同一张图可能在同一场里出现多次，按 key 去重）
	want := map[string]ingest.UploadItem{}
	for i := range urls {
		if results[i].Key == "" {
			continue
		}
		if _, ok := want[results[i].Key]; ok {
			continue
		}
		want[results[i].Key] = ingest.UploadItem{Key: results[i].Key, ContentType: contentTypeOf(urls[i])}
	}
	items := make([]ingest.UploadItem, 0, len(want))
	for _, it := range want {
		items = append(items, it)
	}

	presigned := make(map[string]string, len(items))
	const chunk = 100 // 后端单次上限 200，这里留一半余量
	for start := 0; start < len(items); start += chunk {
		end := start + chunk
		if end > len(items) {
			end = len(items)
		}
		got, err := u.Client.RequestUploadURLs(ctx, items[start:end])
		if err != nil {
			// 申请失败：整批回退成 URL 上报（后端自己下载），不中断整场抓取。
			for i := range results {
				if results[i].Key != "" {
					results[i] = Result{Err: fmt.Errorf("申请上传地址失败: %w", err)}
				}
			}
			return results
		}
		for _, g := range got {
			presigned[g.Key] = g.URL
		}
	}

	log.Printf("[media] 已拿到 %d 个预签名地址，开始上传", len(presigned))

	// ④ 并发 PUT（内容已在内存里，直接传字节）
	u.parallel(len(urls), func(i int) {
		if results[i].Key == "" {
			return
		}
		putURL, ok := presigned[results[i].Key]
		if !ok {
			results[i] = Result{Err: fmt.Errorf("未拿到 %s 的上传地址", results[i].Key)}
			return
		}
		if err := u.put(ctx, putURL, urls[i], data[i]); err != nil {
			results[i] = Result{Err: err}
		}
	})
	return results
}

// parallel 以 Concurrency 路并发执行 fn(0..n-1)。
func (u *Uploader) parallel(n int, fn func(i int)) {
	if n == 0 {
		return
	}
	if u.Concurrency <= 0 {
		// 兜底：并发为 0 会造成无缓冲 channel 死锁（曾把整条抓取卡住）。
		u.Concurrency = 1
	}
	sem := make(chan struct{}, u.Concurrency)
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()
			fn(i)
		}(i)
	}
	wg.Wait()
}

func (u *Uploader) download(ctx context.Context, url string) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	resp, err := u.HTTP.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("下载失败 status=%d", resp.StatusCode)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if len(body) == 0 {
		return nil, fmt.Errorf("下载到空文件")
	}
	return body, nil
}

func (u *Uploader) put(ctx context.Context, putURL, originalURL string, data []byte) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodPut, putURL, bytes.NewReader(data))
	if err != nil {
		return err
	}
	// Content-Type 必须与申请预签名时一致，否则签名校验失败（403）。
	req.Header.Set("Content-Type", contentTypeOf(originalURL))
	resp, err := u.HTTP.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent {
		raw, _ := io.ReadAll(io.LimitReader(resp.Body, 300))
		return fmt.Errorf("PUT 失败 status=%d %s", resp.StatusCode, string(raw))
	}
	return nil
}

// countOK 数一下成功下载的张数（用于日志定位卡在哪一步）。
func countOK(data [][]byte) int {
	n := 0
	for i := range data {
		if len(data[i]) > 0 {
			n++
		}
	}
	return n
}

func sha1Hex(b []byte) string {
	sum := sha1.Sum(b)
	return hex.EncodeToString(sum[:])
}

// extOf 从图片 URL 推断扩展名（取路径最后一段，忽略查询串）。
func extOf(rawURL string) string {
	p := rawURL
	if i := strings.IndexAny(p, "?#"); i >= 0 {
		p = p[:i]
	}
	if i := strings.LastIndex(p, "."); i >= 0 && len(p)-i <= 6 {
		ext := strings.ToLower(p[i:])
		switch ext {
		case ".jpg", ".jpeg", ".png", ".gif", ".webp", ".avif":
			if ext == ".jpeg" {
				return ".jpg"
			}
			return ext
		}
	}
	return ".jpg"
}

// contentTypeOf 与 extOf 保持一致，保证预签名与 PUT 的 Content-Type 相同。
func contentTypeOf(rawURL string) string {
	switch extOf(rawURL) {
	case ".png":
		return "image/png"
	case ".gif":
		return "image/gif"
	case ".webp":
		return "image/webp"
	case ".avif":
		return "image/avif"
	default:
		return "image/jpeg"
	}
}