worker を追加(yt-dlp によるダウンロード、段階別進捗、ffprobe 検査、実行中の取り消し)
This commit is contained in:
@@ -0,0 +1,500 @@
|
||||
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)
|
||||
|
||||
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
|
||||
}
|
||||
relative := media.RelativeVideoPath(job.SeriesTitle, job.Title, job.EpisodeID)
|
||||
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]"))
|
||||
}
|
||||
Reference in New Issue
Block a user