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]")) }