stash/internal/manager/task_generate.go

package manager

import (
	"context"
	"fmt"
	"time"

	"github.com/remeh/sizedwaitgroup"
	"github.com/stashapp/stash/internal/manager/config"
	"github.com/stashapp/stash/pkg/job"
	"github.com/stashapp/stash/pkg/logger"
	"github.com/stashapp/stash/pkg/models"
	"github.com/stashapp/stash/pkg/scene"
	"github.com/stashapp/stash/pkg/scene/generate"
	"github.com/stashapp/stash/pkg/sliceutil/stringslice"
	"github.com/stashapp/stash/pkg/utils"
)

type GenerateMetadataInput struct {
	Sprites             *bool                        `json:"sprites"`
	Previews            *bool                        `json:"previews"`
	ImagePreviews       *bool                        `json:"imagePreviews"`
	PreviewOptions      *GeneratePreviewOptionsInput `json:"previewOptions"`
	Markers             *bool                        `json:"markers"`
	MarkerImagePreviews *bool                        `json:"markerImagePreviews"`
	MarkerScreenshots   *bool                        `json:"markerScreenshots"`
	Transcodes          *bool                        `json:"transcodes"`
	// Generate transcodes even if not required
	ForceTranscodes           *bool `json:"forceTranscodes"`
	Phashes                   *bool `json:"phashes"`
	InteractiveHeatmapsSpeeds *bool `json:"interactiveHeatmapsSpeeds"`
	// scene ids to generate for
	SceneIDs []string `json:"sceneIDs"`
	// marker ids to generate for
	MarkerIDs []string `json:"markerIDs"`
	// overwrite existing media
	Overwrite *bool `json:"overwrite"`
}

type GeneratePreviewOptionsInput struct {
	// Number of segments in a preview file
	PreviewSegments *int `json:"previewSegments"`
	// Preview segment duration, in seconds
	PreviewSegmentDuration *float64 `json:"previewSegmentDuration"`
	// Duration of start of video to exclude when generating previews
	PreviewExcludeStart *string `json:"previewExcludeStart"`
	// Duration of end of video to exclude when generating previews
	PreviewExcludeEnd *string `json:"previewExcludeEnd"`
	// Preset when generating preview
	PreviewPreset *models.PreviewPreset `json:"previewPreset"`
}

const generateQueueSize = 200000

type GenerateJob struct {
	txnManager Repository
	input      GenerateMetadataInput

	overwrite      bool
	fileNamingAlgo models.HashAlgorithm
}

type totalsGenerate struct {
	sprites                  int64
	previews                 int64
	imagePreviews            int64
	markers                  int64
	transcodes               int64
	phashes                  int64
	interactiveHeatmapSpeeds int64

	tasks int
}

func (j *GenerateJob) Execute(ctx context.Context, progress *job.Progress) {
	var scenes []*models.Scene
	var err error
	var markers []*models.SceneMarker

	if j.input.Overwrite != nil {
		j.overwrite = *j.input.Overwrite
	}
	j.fileNamingAlgo = config.GetInstance().GetVideoFileNamingAlgorithm()

	config := config.GetInstance()
	parallelTasks := config.GetParallelTasksWithAutoDetection()

	logger.Infof("Generate started with %d parallel tasks", parallelTasks)

	queue := make(chan Task, generateQueueSize)
	go func() {
		defer close(queue)

		var totals totalsGenerate
		sceneIDs, err := stringslice.StringSliceToIntSlice(j.input.SceneIDs)
		if err != nil {
			logger.Error(err.Error())
		}
		markerIDs, err := stringslice.StringSliceToIntSlice(j.input.MarkerIDs)
		if err != nil {
			logger.Error(err.Error())
		}

		g := &generate.Generator{
			Encoder:     instance.FFMPEG,
			LockManager: instance.ReadLockManager,
			MarkerPaths: instance.Paths.SceneMarkers,
			ScenePaths:  instance.Paths.Scene,
			Overwrite:   j.overwrite,
		}

		if err := j.txnManager.WithTxn(ctx, func(ctx context.Context) error {
			qb := j.txnManager.Scene
			if len(j.input.SceneIDs) == 0 && len(j.input.MarkerIDs) == 0 {
				totals = j.queueTasks(ctx, g, queue)
			} else {
				if len(j.input.SceneIDs) > 0 {
					scenes, err = qb.FindMany(ctx, sceneIDs)
					for _, s := range scenes {
						if err := s.LoadFiles(ctx, qb); err != nil {
							return err
						}

						j.queueSceneJobs(ctx, g, s, queue, &totals)
					}
				}

				if len(j.input.MarkerIDs) > 0 {
					markers, err = j.txnManager.SceneMarker.FindMany(ctx, markerIDs)
					if err != nil {
						return err
					}
					for _, m := range markers {
						j.queueMarkerJob(g, m, queue, &totals)
					}
				}
			}

			return nil
		}); err != nil {
			logger.Error(err.Error())
			return
		}

		logger.Infof("Generating %d sprites %d previews %d image previews %d markers %d transcodes %d phashes %d heatmaps & speeds", totals.sprites, totals.previews, totals.imagePreviews, totals.markers, totals.transcodes, totals.phashes, totals.interactiveHeatmapSpeeds)

		progress.SetTotal(int(totals.tasks))
	}()

	wg := sizedwaitgroup.New(parallelTasks)

	// Start measuring how long the generate has taken. (consider moving this up)
	start := time.Now()
	if err = instance.Paths.Generated.EnsureTmpDir(); err != nil {
		logger.Warnf("could not create temporary directory: %v", err)
	}

	defer func() {
		if err := instance.Paths.Generated.EmptyTmpDir(); err != nil {
			logger.Warnf("failure emptying temporary directory: %v", err)
		}
	}()

	for f := range queue {
		if job.IsCancelled(ctx) {
			break
		}

		wg.Add()
		// #1879 - need to make a copy of f - otherwise there is a race condition
		// where f is changed when the goroutine runs
		localTask := f
		go progress.ExecuteTask(localTask.GetDescription(), func() {
			localTask.Start(ctx)
			wg.Done()
			progress.Increment()
		})
	}

	wg.Wait()

	if job.IsCancelled(ctx) {
		logger.Info("Stopping due to user request")
		return
	}

	elapsed := time.Since(start)
	logger.Info(fmt.Sprintf("Generate finished (%s)", elapsed))
}

func (j *GenerateJob) queueTasks(ctx context.Context, g *generate.Generator, queue chan<- Task) totalsGenerate {
	var totals totalsGenerate

	const batchSize = 1000

	findFilter := models.BatchFindFilter(batchSize)

	for more := true; more; {
		if job.IsCancelled(ctx) {
			return totals
		}

		scenes, err := scene.Query(ctx, j.txnManager.Scene, nil, findFilter)
		if err != nil {
			logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
			return totals
		}

		for _, ss := range scenes {
			if job.IsCancelled(ctx) {
				return totals
			}

			if err := ss.LoadFiles(ctx, j.txnManager.Scene); err != nil {
				logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
				return totals
			}

			j.queueSceneJobs(ctx, g, ss, queue, &totals)
		}

		if len(scenes) != batchSize {
			more = false
		} else {
			*findFilter.Page++
		}
	}

	return totals
}

func getGeneratePreviewOptions(optionsInput GeneratePreviewOptionsInput) generate.PreviewOptions {
	config := config.GetInstance()

	ret := generate.PreviewOptions{
		Segments:        config.GetPreviewSegments(),
		SegmentDuration: config.GetPreviewSegmentDuration(),
		ExcludeStart:    config.GetPreviewExcludeStart(),
		ExcludeEnd:      config.GetPreviewExcludeEnd(),
		Preset:          config.GetPreviewPreset().String(),
		Audio:           config.GetPreviewAudio(),
	}

	if optionsInput.PreviewSegments != nil {
		ret.Segments = *optionsInput.PreviewSegments
	}

	if optionsInput.PreviewSegmentDuration != nil {
		ret.SegmentDuration = *optionsInput.PreviewSegmentDuration
	}

	if optionsInput.PreviewExcludeStart != nil {
		ret.ExcludeStart = *optionsInput.PreviewExcludeStart
	}

	if optionsInput.PreviewExcludeEnd != nil {
		ret.ExcludeEnd = *optionsInput.PreviewExcludeEnd
	}

	if optionsInput.PreviewPreset != nil {
		ret.Preset = optionsInput.PreviewPreset.String()
	}

	return ret
}

func (j *GenerateJob) queueSceneJobs(ctx context.Context, g *generate.Generator, scene *models.Scene, queue chan<- Task, totals *totalsGenerate) {
	if utils.IsTrue(j.input.Sprites) {
		task := &GenerateSpriteTask{
			Scene:               *scene,
			Overwrite:           j.overwrite,
			fileNamingAlgorithm: j.fileNamingAlgo,
		}

		if j.overwrite || task.required() {
			totals.sprites++
			totals.tasks++
			queue <- task
		}
	}

	generatePreviewOptions := j.input.PreviewOptions
	if generatePreviewOptions == nil {
		generatePreviewOptions = &GeneratePreviewOptionsInput{}
	}
	options := getGeneratePreviewOptions(*generatePreviewOptions)

	if utils.IsTrue(j.input.Previews) {
		task := &GeneratePreviewTask{
			Scene:               *scene,
			ImagePreview:        utils.IsTrue(j.input.ImagePreviews),
			Options:             options,
			Overwrite:           j.overwrite,
			fileNamingAlgorithm: j.fileNamingAlgo,
			generator:           g,
		}

		sceneHash := scene.GetHash(task.fileNamingAlgorithm)
		addTask := false
		if j.overwrite || !task.doesVideoPreviewExist(sceneHash) {
			totals.previews++
			addTask = true
		}

		if utils.IsTrue(j.input.ImagePreviews) && (j.overwrite || !task.doesImagePreviewExist(sceneHash)) {
			totals.imagePreviews++
			addTask = true
		}

		if addTask {
			totals.tasks++
			queue <- task
		}
	}

	if utils.IsTrue(j.input.Markers) {
		task := &GenerateMarkersTask{
			TxnManager:          j.txnManager,
			Scene:               scene,
			Overwrite:           j.overwrite,
			fileNamingAlgorithm: j.fileNamingAlgo,
			ImagePreview:        utils.IsTrue(j.input.MarkerImagePreviews),
			Screenshot:          utils.IsTrue(j.input.MarkerScreenshots),

			generator: g,
		}

		markers := task.markersNeeded(ctx)
		if markers > 0 {
			totals.markers += int64(markers)
			totals.tasks++

			queue <- task
		}
	}

	if utils.IsTrue(j.input.Transcodes) {
		forceTranscode := utils.IsTrue(j.input.ForceTranscodes)
		task := &GenerateTranscodeTask{
			Scene:               *scene,
			Overwrite:           j.overwrite,
			Force:               forceTranscode,
			fileNamingAlgorithm: j.fileNamingAlgo,
			g:                   g,
		}
		if task.isTranscodeNeeded() {
			totals.transcodes++
			totals.tasks++
			queue <- task
		}
	}

	if utils.IsTrue(j.input.Phashes) {
		// generate for all files in scene
		for _, f := range scene.Files.List() {
			task := &GeneratePhashTask{
				File:                f,
				fileNamingAlgorithm: j.fileNamingAlgo,
				txnManager:          j.txnManager,
				fileUpdater:         j.txnManager.File,
				Overwrite:           j.overwrite,
			}

			if task.shouldGenerate() {
				totals.phashes++
				totals.tasks++
				queue <- task
			}
		}
	}

	if utils.IsTrue(j.input.InteractiveHeatmapsSpeeds) {
		task := &GenerateInteractiveHeatmapSpeedTask{
			Scene:               *scene,
			Overwrite:           j.overwrite,
			fileNamingAlgorithm: j.fileNamingAlgo,
			TxnManager:          j.txnManager,
		}

		if task.shouldGenerate() {
			totals.interactiveHeatmapSpeeds++
			totals.tasks++
			queue <- task
		}
	}
}

func (j *GenerateJob) queueMarkerJob(g *generate.Generator, marker *models.SceneMarker, queue chan<- Task, totals *totalsGenerate) {
	task := &GenerateMarkersTask{
		TxnManager:          j.txnManager,
		Marker:              marker,
		Overwrite:           j.overwrite,
		fileNamingAlgorithm: j.fileNamingAlgo,
		generator:           g,
	}
	totals.markers++
	totals.tasks++
	queue <- task
}