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

	// sem 全进程共享的信号量，避免「每场秀各开 N 路」叠乘。
	//
	// 早先每个 UploadAll 调用各自 make 一个信号量：秀级并发 2 × 每场 6 = 12 路
	// 同时打 Vogue，被对方限流，单张从 9.8s 掉到 ~90s，还伴随 50% 失败。
	once sync.Once
	sem  chan struct{}
}

// semaphore 惰性初始化共享信号量（并发数非法时兜底为 1，否则无缓冲 channel 会死锁）。
func (u *Uploader) semaphore() chan struct{} {
	u.once.Do(func() {
		c := u.Concurrency
		if c <= 0 {
			c = 1
		}
		u.sem = make(chan struct{}, c)
	})
	return u.sem
}

// 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
	})
	// 失败时把**第一条**错误打出来：失败通常是同一个原因（限流 / 连接重置），
	// 全打会把日志刷爆，而一条就足以定位。此前只统计个数，导致 50% 失败却查不出原因。
	if ok := countOK(data); ok < n {
		for i := 0; i < n; i++ {
			if errs[i] != nil {
				log.Printf("[media][warn] 下载失败 %s: %v", clipURL(urls[start+i]), errs[i])
				break
			}
		}
		log.Printf("[media] 下载完成 %d 张（失败 %d）", ok, n-ok)
	} else {
		log.Printf("[media] 下载完成 %d 张（失败 0）", ok)
	}

	// ② 生成本批 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
	}
	// 共享信号量：并发上限是**整个进程**的，不是每次调用各自一份。
	// 早先每次 UploadAll 各建一个：秀级并发 2 × 每场 6 = 12 路同时打 Vogue，
	// 被对方限流后单张从 9.8s 掉到 ~90s，并伴随大量失败。
	sem := u.semaphore()
	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
}

// clipURL 截短 URL 只用于日志（完整 URL 太长，会淹没真正有用的错误信息）。
func clipURL(u string) string {
	if len(u) <= 90 {
		return u
	}
	return u[:90] + "..."
}

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