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
	}

	// 分批处理：**不能**先把整场的图全读进内存。
	// 一场 100 张 × 1.2MB ≈ 120MB；而 SpiderStart 对同一品牌的多个场次是并发抓的，
	// 峰值会翻好几倍——线上实测 spider 进程一度涨到 1.3GB，把小机器（3.7G）压到只剩 141MB 可用。
	// 按批（默认 8 张）下载→上传→立即释放，峰值稳定在 ~10MB。
	const batchSize = 8
	for start := 0; start < len(urls); start += batchSize {
		end := start + batchSize
		if end > len(urls) {
			end = len(urls)
		}
		u.uploadBatch(ctx, urls, results, start, end)
	}
	return results
}

// uploadBatch 处理 urls[start:end] 这一段：下载 → 生成 key → 预签名 → PUT。
// 返回时本批的图片字节已可回收（不再被引用）。
func (u *Uploader) uploadBatch(ctx context.Context, urls []string, results []Result, start, end int) {
	n := end - start
	data := make([][]byte, n)
	errs := make([]error, n)

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

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

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

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

	// ④ 并发 PUT 本批
	u.parallel(n, func(i int) {
		res := results[start+i]
		if res.Key == "" {
			return
		}
		putURL, ok := presigned[res.Key]
		if !ok {
			results[start+i] = Result{Err: fmt.Errorf("未拿到 %s 的上传地址", res.Key)}
			return
		}
		if perr := u.put(ctx, putURL, urls[start+i], data[i]); perr != nil {
			results[start+i] = Result{Err: perr}
		}
	})

	log.Printf("[media] 本批 %d 张上传完成（%d/%d）", n, end, len(urls))
}

// 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"
	}
}