tver/internal/worker/runner.go

package worker

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"gitea.kazu634.com/kazu634/tver/internal/media"
	"gitea.kazu634.com/kazu634/tver/internal/store"
	"github.com/jackc/pgx/v5"
)

var (
	errCanceled      = errors.New("job canceled")
	errShutdown      = errors.New("worker shutting down")
	errInvalidOutput = errors.New("invalid output")
	headerSecretRE   = regexp.MustCompile(`(?im)(authorization|cookie)([=:][ \t]*)([^\r\n]+)`)
	paramSecretRE    = regexp.MustCompile(`(?i)(token|platform_uid|platform_token)([=:][[:space:]]*)([^&[:space:]]+)`)
)

type JobStore interface {
	ClaimNext(context.Context) (store.Job, error)
	UpdateProgress(context.Context, int64, store.Progress) error
	Heartbeat(context.Context, int64) error
	GetJob(context.Context, int64) (store.Job, error)
	Complete(context.Context, int64, string, int64) error
	Fail(context.Context, int64, string, string) error
	RequeueRunning(context.Context, int64) error
}

type Config struct {
	VideoDir          string
	YTDLPPath         string
	FFprobePath       string
	Concurrency       int
	PollInterval      time.Duration
	ProgressInterval  time.Duration
	HeartbeatInterval time.Duration
	CancelInterval    time.Duration
}

type Runner struct {
	store  JobStore
	config Config
	logger *log.Logger
}

type cappedBuffer struct {
	buffer bytes.Buffer
	limit  int
}

func (b *cappedBuffer) Write(p []byte) (int, error) {
	written := len(p)
	remaining := b.limit - b.buffer.Len()
	if remaining > 0 {
		if remaining < len(p) {
			p = p[:remaining]
		}
		_, _ = b.buffer.Write(p)
	}
	return written, nil
}

func (b *cappedBuffer) String() string { return b.buffer.String() }

func scanLines(reader io.Reader, capture *cappedBuffer) <-chan string {
	lines := make(chan string)
	go func() {
		defer close(lines)
		scanner := bufio.NewScanner(reader)
		scanner.Buffer(make([]byte, 4096), 1<<20)
		for scanner.Scan() {
			line := scanner.Text()
			if capture != nil {
				_, _ = capture.Write([]byte(line + "\n"))
			}
			lines <- line
		}
	}()
	return lines
}

func New(jobs JobStore, config Config, logger *log.Logger) (*Runner, error) {
	if config.VideoDir == "" {
		config.VideoDir = "/videos"
	}
	if config.YTDLPPath == "" {
		config.YTDLPPath = "yt-dlp"
	}
	if config.FFprobePath == "" {
		config.FFprobePath = "ffprobe"
	}
	if config.Concurrency == 0 {
		config.Concurrency = 1
	}
	if config.Concurrency < 1 || config.Concurrency > 2 {
		return nil, errors.New("worker concurrency must be 1 or 2")
	}
	if config.PollInterval <= 0 {
		config.PollInterval = time.Second
	}
	if config.ProgressInterval <= 0 {
		config.ProgressInterval = time.Second
	}
	if config.HeartbeatInterval <= 0 {
		config.HeartbeatInterval = 5 * time.Second
	}
	if config.CancelInterval <= 0 {
		config.CancelInterval = 500 * time.Millisecond
	}
	if logger == nil {
		logger = log.Default()
	}
	return &Runner{store: jobs, config: config, logger: logger}, nil
}

func (r *Runner) Run(ctx context.Context) error {
	sem := make(chan struct{}, r.config.Concurrency)
	var wg sync.WaitGroup
	ticker := time.NewTicker(r.config.PollInterval)
	defer ticker.Stop()
	for {
		for len(sem) < cap(sem) && ctx.Err() == nil {
			job, err := r.store.ClaimNext(ctx)
			if errors.Is(err, pgx.ErrNoRows) {
				break
			}
			if err != nil {
				r.logger.Printf("ジョブ取得に失敗: %v", err)
				break
			}
			sem <- struct{}{}
			wg.Add(1)
			go func() {
				defer wg.Done()
				defer func() { <-sem }()
				r.ProcessJob(ctx, job)
			}()
		}
		select {
		case <-ctx.Done():
			wg.Wait()
			return nil
		case <-ticker.C:
		}
	}
}

func (r *Runner) ProcessJob(ctx context.Context, job store.Job) {
	workDir := filepath.Join(r.config.VideoDir, ".incoming", strconv.FormatInt(job.ID, 10))
	if err := os.RemoveAll(workDir); err != nil {
		r.fail(job.ID, "unknown", fmt.Sprintf("作業ディレクトリを初期化できません: %v", err))
		return
	}
	if err := os.MkdirAll(workDir, 0o755); err != nil {
		r.fail(job.ID, "unknown", fmt.Sprintf("作業ディレクトリを作成できません: %v", err))
		return
	}
	defer os.RemoveAll(workDir)
	relative, err := media.RelativeVideoPath(job.SeriesTitle, job.Title, job.EpisodeID)
	if err != nil {
		r.failIfRunning(job.ID, "invalid_output", redact(err.Error()))
		return
	}

	downloaded, stderr, err := r.download(ctx, job, workDir)
	if errors.Is(err, errShutdown) {
		r.requeue(job.ID)
		return
	}
	if errors.Is(err, errCanceled) {
		return
	}
	if err != nil {
		message := redact(stderr)
		if message == "" {
			message = redact(err.Error())
		}
		kind := classifyFailure(message)
		if errors.Is(err, errInvalidOutput) {
			kind = "invalid_output"
		}
		r.failIfRunning(job.ID, kind, message)
		return
	}
	if r.canceled(job.ID) {
		return
	}
	if err := r.probe(ctx, downloaded); err != nil {
		if ctx.Err() != nil {
			r.requeue(job.ID)
			return
		}
		r.failIfRunning(job.ID, "invalid_output", redact(err.Error()))
		return
	}
	if r.canceled(job.ID) {
		return
	}
	if err := r.store.UpdateProgress(ctx, job.ID, store.Progress{Phase: "moving"}); err != nil {
		if r.canceled(job.ID) {
			return
		}
		r.failIfRunning(job.ID, "unknown", "移動状態を記録できません: "+err.Error())
		return
	}
	destination := filepath.Join(r.config.VideoDir, relative)
	if err := os.MkdirAll(filepath.Dir(destination), 0o755); err != nil {
		r.failIfRunning(job.ID, "unknown", "保存先を作成できません: "+err.Error())
		return
	}
	if err := os.Rename(downloaded, destination); err != nil {
		r.failIfRunning(job.ID, "unknown", "動画を保存先へ移動できません: "+err.Error())
		return
	}
	info, err := os.Stat(destination)
	if err != nil {
		r.failIfRunning(job.ID, "unknown", "保存後の動画を確認できません: "+err.Error())
		return
	}
	if err := r.store.Complete(ctx, job.ID, filepath.ToSlash(relative), info.Size()); err != nil {
		if r.canceled(job.ID) {
			if removeErr := os.Remove(destination); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) {
				r.logger.Printf("取り消されたジョブ %d の保存済みファイル削除に失敗: %v", job.ID, removeErr)
			}
			return
		}
		r.logger.Printf("ジョブ %d の完了記録に失敗: %v", job.ID, err)
	}
}

func (r *Runner) download(ctx context.Context, job store.Job, workDir string) (string, string, error) {
	output := filepath.Join(workDir, "%(id)s.%(ext)s")
	args := []string{"--no-playlist", "--no-simulate", "--quiet", "--progress", "--newline",
		"--continue", "--retries", "0", "--fragment-retries", "0", "--extractor-retries", "0",
		"-f", "bv*[vcodec^=avc1]+ba/b", "--merge-output-format", "mp4",
		"--progress-template", `download:{"event":"download","format_id":%(info.format_id)j,"progress":%(progress)j}`,
		"--progress-template", `postprocess:{"event":"postprocess","progress":%(progress)j}`,
		"--print", `after_move:{"event":"after_move","filepath":%(filepath)j}`, "-o", output,
		"https://tver.jp/episodes/" + job.EpisodeID}
	cmd := exec.Command(r.config.YTDLPPath, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return "", "", err
	}
	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return "", "", err
	}
	if err := cmd.Start(); err != nil {
		return "", "", err
	}
	stderr := cappedBuffer{limit: 64 << 10}
	stdoutLines := scanLines(stdout, nil)
	stderrLines := scanLines(stderrPipe, &stderr)
	waited := make(chan error, 1)
	go func() { waited <- cmd.Wait() }()
	heartbeat := time.NewTicker(r.config.HeartbeatInterval)
	cancelCheck := time.NewTicker(r.config.CancelInterval)
	defer heartbeat.Stop()
	defer cancelCheck.Stop()
	_ = r.store.Heartbeat(ctx, job.ID)
	var afterMove string
	var pending *store.Progress
	nextProgress := time.Time{}
	lastPhase := ""
	recordProgress := func(progress store.Progress) {
		pending = &progress
		now := time.Now()
		if progress.Phase != lastPhase || now.After(nextProgress) {
			_ = r.store.UpdateProgress(ctx, job.ID, progress)
			nextProgress = now.Add(r.config.ProgressInterval)
			lastPhase = progress.Phase
			pending = nil
		}
	}
	stopping := ""
	var forceKill <-chan time.Time
	ctxDone := ctx.Done()
	waitComplete := false
	var waitErr error
	for {
		if waitComplete && stdoutLines == nil && stderrLines == nil {
			if pending != nil && stopping == "" {
				updateCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				_ = r.store.UpdateProgress(updateCtx, job.ID, *pending)
				cancel()
			}
			if stopping == "canceled" {
				return "", stderr.String(), errCanceled
			}
			if stopping == "shutdown" {
				return "", stderr.String(), errShutdown
			}
			if waitErr != nil {
				r.logger.Printf("yt-dlp ジョブ %d 失敗: %s", job.ID, redact(stderr.String()))
				return "", stderr.String(), waitErr
			}
			path, err := pathInside(workDir, afterMove)
			if err != nil {
				return "", stderr.String(), err
			}
			return path, stderr.String(), nil
		}
		select {
		case line, ok := <-stdoutLines:
			if !ok {
				stdoutLines = nil
				continue
			}
			if path, ok := parseAfterMove(line); ok {
				afterMove = path
				continue
			}
			if progress, ok := parseProgress(line); ok {
				recordProgress(progress)
			}
		case line, ok := <-stderrLines:
			if !ok {
				stderrLines = nil
				continue
			}
			if progress, ok := parseProgress(line); ok {
				recordProgress(progress)
			}
		case <-heartbeat.C:
			_ = r.store.Heartbeat(ctx, job.ID)
		case <-cancelCheck.C:
			if r.canceled(job.ID) && stopping == "" {
				stopping = "canceled"
				if !waitComplete {
					stopProcessGroup(cmd)
					forceKill = time.After(2 * time.Second)
				}
			}
		case <-ctxDone:
			if stopping == "" {
				stopping = "shutdown"
				if !waitComplete {
					stopProcessGroup(cmd)
					forceKill = time.After(2 * time.Second)
				}
			}
			ctxDone = nil
		case <-forceKill:
			_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
			forceKill = nil
		case waitErr = <-waited:
			waitComplete = true
			waited = nil
			forceKill = nil
		}
	}
}

func stopProcessGroup(cmd *exec.Cmd) {
	if cmd.Process == nil {
		return
	}
	_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM)
}

func pathInside(root, candidate string) (string, error) {
	if candidate == "" {
		return "", errors.New("yt-dlp did not report after_move filepath")
	}
	rootAbs, err := filepath.Abs(root)
	if err != nil {
		return "", err
	}
	candidateAbs, err := filepath.Abs(candidate)
	if err != nil {
		return "", err
	}
	rel, err := filepath.Rel(rootAbs, candidateAbs)
	if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("%w: yt-dlp filepath leaves working directory", errInvalidOutput)
	}
	info, err := os.Lstat(candidateAbs)
	if err != nil {
		return "", fmt.Errorf("%w: yt-dlp output file is missing: %v", errInvalidOutput, err)
	}
	if !info.Mode().IsRegular() {
		return "", fmt.Errorf("%w: yt-dlp output is not a regular file", errInvalidOutput)
	}
	rootReal, err := filepath.EvalSymlinks(rootAbs)
	if err != nil {
		return "", fmt.Errorf("%w: resolve working directory: %v", errInvalidOutput, err)
	}
	candidateReal, err := filepath.EvalSymlinks(candidateAbs)
	if err != nil {
		return "", fmt.Errorf("%w: resolve yt-dlp output: %v", errInvalidOutput, err)
	}
	realRelative, err := filepath.Rel(rootReal, candidateReal)
	if err != nil || realRelative == ".." || strings.HasPrefix(realRelative, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("%w: yt-dlp output resolves outside working directory", errInvalidOutput)
	}
	return candidateAbs, nil
}

func (r *Runner) probe(ctx context.Context, path string) error {
	cmd := exec.CommandContext(ctx, r.config.FFprobePath, "-v", "error", "-show_entries",
		"format=format_name", "-show_entries", "stream=codec_type,codec_name", "-of", "json", path)
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("ffprobe failed: %w", err)
	}
	var result struct {
		Format struct {
			FormatName string `json:"format_name"`
		} `json:"format"`
		Streams []struct {
			CodecType string `json:"codec_type"`
			CodecName string `json:"codec_name"`
		} `json:"streams"`
	}
	if err := json.Unmarshal(output, &result); err != nil {
		return fmt.Errorf("ffprobe output is not JSON: %w", err)
	}
	video, audio := false, false
	for _, stream := range result.Streams {
		video = video || stream.CodecType == "video" && stream.CodecName == "h264"
		audio = audio || stream.CodecType == "audio" && stream.CodecName == "aac"
	}
	if !strings.Contains(result.Format.FormatName, "mp4") || !video || !audio {
		return fmt.Errorf("expected MP4 with H.264 and AAC, got format=%s video=%t audio=%t", result.Format.FormatName, video, audio)
	}
	return nil
}

func (r *Runner) canceled(id int64) bool {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	job, err := r.store.GetJob(ctx, id)
	return err == nil && job.Status == "canceled"
}

func (r *Runner) failIfRunning(id int64, kind, message string) {
	if r.canceled(id) {
		return
	}
	r.fail(id, kind, message)
}

func (r *Runner) fail(id int64, kind, message string) {
	if len(message) > 8192 {
		message = message[:8192]
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := r.store.Fail(ctx, id, kind, message); err != nil {
		r.logger.Printf("ジョブ %d の失敗記録に失敗: %v", id, err)
	}
}

func (r *Runner) requeue(id int64) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := r.store.RequeueRunning(ctx, id); err != nil && !errors.Is(err, store.ErrInvalidTransition) {
		r.logger.Printf("ジョブ %d の待機中への復帰に失敗: %v", id, err)
	}
}

func classifyFailure(message string) string {
	lower := strings.ToLower(message)
	for _, test := range []struct {
		kind     string
		patterns []string
	}{
		{"drm", []string{"drm", "digital rights management", "protected by"}},
		{"geo", []string{"geo-restricted", "not available in your country", "not available from your location", "geo restriction", "region restriction", "geographic restriction"}},
		{"network", []string{"network", "connection", "timed out", "temporary failure", "unable to download", "proxy error", "http error 5"}},
		{"unavailable", []string{"media_not_found", "video unavailable", "not available", "no longer available", "http error 404", "episode does not exist"}},
		{"extract", []string{"extractor error", "unable to extract", "unsupported url", "no video formats", "after_move", "output file is missing"}},
	} {
		for _, pattern := range test.patterns {
			if strings.Contains(lower, pattern) {
				return test.kind
			}
		}
	}
	return "unknown"
}

func redact(message string) string {
	message = headerSecretRE.ReplaceAllString(message, "$1$2[REDACTED]")
	return strings.TrimSpace(paramSecretRE.ReplaceAllString(message, "$1$2[REDACTED]"))
}