spider/cmd/serve.go

package cmd

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/spf13/cobra"

	"my-spiders/internal/config"
	"my-spiders/internal/media"
	"my-spiders/internal/spider"
)

var (
	serveWorkers  int
	serveInterval int
	serveMaxCo    int
)

// serveCmd 把爬虫变成「常驻服务」：不再由人决定爬谁，而是从后台领取任务。
//
// 与 vogue 子命令的区别：
//   - vogue：一次性的，参数指定品牌（--brand），跑完退出；
//   - serve：常驻，循环「领任务 → 执行 → 回报」，没有任务就空转等待。
//
// 这样后台（/admin/crawl）就能控制爬哪些品牌，爬虫部署在哪台机器都不影响。
var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "以服务方式运行：从后台领取品牌抓取任务并执行",
	Run: func(cmd *cobra.Command, args []string) {
		if IngestClient == nil {
			log.Fatal("[致命] 未启用入库管线：config.yml 里 ingest.enabled 必须为 true（serve 模式依赖后台接口领取任务）")
		}

		host, _ := os.Hostname()
		if strings.TrimSpace(host) == "" {
			host = "spider"
		}

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		// Ctrl+C / SIGTERM 优雅退出：不再领新任务，等在跑的品牌结束。
		sig := make(chan os.Signal, 1)
		signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
		var wg sync.WaitGroup

		log.Printf("[serve] 启动 %d 个 worker，轮询间隔 %ds，单品牌并发 %d", serveWorkers, serveInterval, serveMaxCo)

		for i := 0; i < serveWorkers; i++ {
			wg.Add(1)
			go func(idx int) {
				defer wg.Done()
				workerID := fmt.Sprintf("%s-%d-%d", host, os.Getpid(), idx)
				runServeWorker(ctx, workerID)
			}(i)
		}

		<-sig
		log.Println("[serve] 收到退出信号，等待进行中的任务结束…")
		cancel()
		wg.Wait()
		log.Println("[serve] 已退出")
	},
}

func runServeWorker(ctx context.Context, workerID string) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		task, err := IngestClient.FetchTask(ctx, workerID)
		if err != nil {
			log.Printf("[serve] %s 领取任务失败: %v", workerID, err)
			sleepOrDone(ctx, serveInterval)
			continue
		}
		if task == nil {
			// 队列空：退避轮询，避免打满后台接口。
			sleepOrDone(ctx, serveInterval)
			continue
		}

		log.Printf("[serve] %s 领取任务 task=%d 品牌=%s(uid=%s) 场次上限=%d",
			workerID, task.ID, task.BrandName, task.BrandUID, task.MaxShows)

		start := time.Now()
		shows, errMsg := executeBrand(ctx, task.BrandUID, task.BrandName, int(task.MaxShows))
		elapsed := time.Since(start).Round(time.Second)

		status := "done"
		if errMsg != "" {
			status = "failed"
		}
		if err := IngestClient.ReportTask(ctx, task.ID, status, shows, errMsg); err != nil {
			log.Printf("[serve] %s 回报任务 %d 失败: %v", workerID, task.ID, err)
			// 回报失败不要紧：租约（CrawlStuckTimeoutSec）到期后后台会把任务重新放回队列。
			sleepOrDone(ctx, serveInterval)
			continue
		}
		log.Printf("[serve] %s 任务 %d 完成：status=%s 提交 %d 场，耗时 %s", workerID, task.ID, status, shows, elapsed)
	}
}

// executeBrand 执行单个品牌，异常（panic）也转成错误返回，避免一个坏品牌把整个 serve 进程带崩。
func executeBrand(ctx context.Context, brandUID, brandName string, maxShows int) (shows int, errMsg string) {
	defer func() {
		if r := recover(); r != nil {
			shows = 0
			errMsg = fmt.Sprintf("panic: %v", r)
			log.Printf("[serve] 执行品牌 %s 时 panic: %v", brandName, r)
		}
	}()

	s := spider.NewVogueSpider(serveMaxCo, 0, maxShows, IngestClient)
	s.Debug = debug
	// 默认跳过已入库的秀：图片 key 是内容 sha1，不下载无法判重，
	// 场次级预检是唯一能在下载前省掉整场带宽的手段。
	s.Force = forceAll
	// 媒体机直传：spider 自己下载并上传图片，后端只写库。
	if config.GlobalConfig.Media.Upload && IngestClient != nil {
		s.Media = media.New(IngestClient, "runway", config.GlobalConfig.Media.Concurrency)
	}
	shows = s.RunBrand(brandUID, brandName, maxShows)
	return shows, ""
}

// sleepOrDone 可被 ctx 取消打断的 sleep。
func sleepOrDone(ctx context.Context, seconds int) {
	if seconds < 1 {
		seconds = 5
	}
	t := time.NewTimer(time.Duration(seconds) * time.Second)
	defer t.Stop()
	select {
	case <-ctx.Done():
	case <-t.C:
	}
}

func init() {
	rootCmd.AddCommand(serveCmd)
	serveCmd.Flags().IntVarP(&serveWorkers, "workers", "w", 1, "并发执行的品牌数（每路一个 goroutine）")
	serveCmd.Flags().IntVarP(&serveInterval, "interval", "i", 5, "队列为空时的轮询间隔（秒）")
	serveCmd.Flags().IntVarP(&serveMaxCo, "max-co", "m", 5, "单个品牌内部的抓取并发数")
	// 与 vogue 子命令共用同一个开关：默认跳过已入库的秀，确需刷新时加 --force。
	serveCmd.Flags().BoolVarP(&forceAll, "force", "f", false, "忽略「已入库」预检，强制重抓全部场次")
}