spider/internal/embed/worker.go

// Package embed 是「向量回填」的执行端。
//
// 为什么需要它：媒体拆分后 spider 直接传图到对象存储，后端拿到的是 key 而不是图片字节，
// 于是入库时算不出向量（image_embeddings 缺行）。这个 worker 补上这一段：
//
//	向后端要「缺向量的图片」 → 从存储下载 → 送本地 CLIP 编码 → 把向量回传给后端
//
// 它跑在**有 GPU 的机器**上，所以放在 spider 仓库里（和爬虫同一份部署物），
// 但逻辑上只在用 GPU，与抓取完全独立：不抓 Vogue、不写业务库。
package embed

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"sync"
	"time"

	"my-spiders/internal/ingest"
)

// Worker 向量回填执行端。
type Worker struct {
	Client      *ingest.Client // 取任务 + 回传（HMAC）
	Vision      string         // 视觉服务地址，如 http://127.0.0.1:8900
	Concurrency int            // 下载并发
	Batch       int            // 单次送编码的张数
	HTTP        *http.Client
}

// Stats 一轮处理的统计。
type Stats struct {
	Fetched   int // 后端派发的张数
	Saved     int // 后端确认写入的张数
	Failed    int // 失败张数（已计入失败计数）
	Remaining int // 后端统计的剩余量（-1 = 未知）
}

// New 创建 worker。并发/批量传 0 时取默认值（4 / 8）。
func New(client *ingest.Client, vision string, concurrency, batch int) *Worker {
	if concurrency <= 0 {
		concurrency = 4
	}
	if batch <= 0 {
		batch = 8
	}
	return &Worker{
		Client:      client,
		Vision:      vision,
		Concurrency: concurrency,
		Batch:       batch,
		// 图片可能几百 KB~几 MB，60s 足够；向量 POST 走同一个 client。
		HTTP: &http.Client{Timeout: 60 * time.Second},
	}
}

// RunOnce 处理一轮：拉取 → 下载 → 编码 → 回传。
//
// 不会因为单张失败而中断：下载失败/编码跳过的图进 failures，由后端累加计数，
// 连续失败到上限后自然退出队列（坏图不会无限占位）。
func (w *Worker) RunOnce(ctx context.Context, kind string, limit int) (Stats, error) {
	var st Stats
	pending, remaining, err := w.Client.FetchPendingEmbeddings(ctx, kind, limit)
	if err != nil {
		return st, err
	}
	st.Fetched = len(pending)
	st.Remaining = remaining
	if len(pending) == 0 {
		log.Printf("[embed] %s 无待回填图片（取时剩余 %d）", kind, remaining)
		return st, nil
	}

	// ① 并发下载：后端给的是可直接访问的 url，worker 不关心存储实现。
	data := make([][]byte, len(pending))
	errs := make([]error, len(pending))
	w.parallel(len(pending), func(i int) {
		body, derr := w.download(ctx, pending[i].URL)
		if derr != nil {
			errs[i] = derr
			return
		}
		data[i] = body
	})

	// ② 编码：只把下载成功的送进去，按 Batch 分批（视觉服务对单次张数有上限）。
	type job struct {
		pos int
		img []byte
	}
	jobs := make([]job, 0, len(pending))
	for i := range pending {
		if errs[i] == nil {
			jobs = append(jobs, job{pos: i, img: data[i]})
		}
	}

	items := make([]ingest.EmbeddingVector, 0, len(jobs))
	failures := make([]ingest.EmbeddingFailure, 0, len(pending)-len(jobs))

	for start := 0; start < len(jobs); start += w.Batch {
		end := start + w.Batch
		if end > len(jobs) {
			end = len(jobs)
		}
		batch := jobs[start:end]
		imgs := make([][]byte, len(batch))
		for i := range batch {
			imgs[i] = batch[i].img
		}
		vecs, idxs, skipped, eerr := w.embed(ctx, imgs)
		if eerr != nil {
			for _, b := range batch {
				p := pending[b.pos]
				failures = append(failures, ingest.EmbeddingFailure{Kind: kind, ImageID: p.ImageID, Error: eerr.Error()})
			}
			continue
		}
		// indexes 缺省时按位置对应；存在时严格按它映射（视觉服务可能跳过坏图）。
		for j := range vecs {
			posInBatch := j
			if len(idxs) == len(vecs) {
				posInBatch = idxs[j]
			}
			if posInBatch < 0 || posInBatch >= len(batch) {
				continue
			}
			p := pending[batch[posInBatch].pos]
			items = append(items, ingest.EmbeddingVector{Kind: kind, ImageID: p.ImageID, Vector: vecs[j]})
		}
		for _, s := range skipped {
			if s < 0 || s >= len(batch) {
				continue
			}
			p := pending[batch[s].pos]
			failures = append(failures, ingest.EmbeddingFailure{Kind: kind, ImageID: p.ImageID, Error: "视觉服务跳过（图片无法解码）"})
		}
	}

	// ③ 下载失败的也一并上报，让后端计数
	for i := range pending {
		if errs[i] != nil {
			failures = append(failures, ingest.EmbeddingFailure{Kind: kind, ImageID: pending[i].ImageID, Error: errs[i].Error()})
		}
	}

	if len(items) == 0 && len(failures) == 0 {
		return st, nil
	}
	saved, failed, serr := w.Client.SubmitEmbeddings(ctx, items, failures)
	if serr != nil {
		return st, serr
	}
	st.Saved = saved
	st.Failed = failed
	log.Printf("[embed] %s 本轮：取到 %d 张 → 写入 %d，失败 %d（剩余 %d）",
		kind, st.Fetched, st.Saved, st.Failed, st.Remaining)
	return st, nil
}

// embed 调本地视觉服务编码一批图片。
// 返回的 idxs 是「vectors 对应入参下标」；skipped 是无法解码的下标。
func (w *Worker) embed(ctx context.Context, imgs [][]byte) ([][]float32, []int, []int, error) {
	payload := struct {
		Images []string `json:"images"`
	}{Images: make([]string, len(imgs))}
	for i, b := range imgs {
		payload.Images[i] = base64.StdEncoding.EncodeToString(b)
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return nil, nil, nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, w.Vision+"/embed", bytes.NewReader(body))
	if err != nil {
		return nil, nil, nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := w.HTTP.Do(req)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("视觉服务不可用: %w", err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusOK {
		return nil, nil, nil, fmt.Errorf("视觉服务返回 %d: %s", resp.StatusCode, truncate(string(raw), 120))
	}
	var out struct {
		Dim     int         `json:"dim"`
		Vectors [][]float32 `json:"vectors"`
		Indexes []int       `json:"indexes"`
		Skipped []int       `json:"skipped"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, nil, nil, fmt.Errorf("视觉服务响应解析失败: %w", err)
	}
	return out.Vectors, out.Indexes, out.Skipped, nil
}

func (w *Worker) 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 := w.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
}

// parallel 以 Concurrency 路并发执行 fn(0..n-1)。
func (w *Worker) parallel(n int, fn func(i int)) {
	if n == 0 {
		return
	}
	if w.Concurrency <= 0 {
		w.Concurrency = 1
	}
	sem := make(chan struct{}, w.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 truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}