Rooba
/
stash
mirror of https://github.com/stashapp/stash.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor generate task (#1871) 

* Refactor generate task
This commit is contained in:
WithoutPants 2021-10-23 10:48:42 +11:00 committed by GitHub
parent f80a5e3222
commit 0ab8d32687
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 341 additions and 328 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

322
pkg/manager/manager_tasks.go
View File

@ -6,9 +6,6 @@ import (
"fmt"
"strconv"
"sync"
"time"
"github.com/remeh/sizedwaitgroup"
"github.com/stashapp/stash/pkg/job"
"github.com/stashapp/stash/pkg/logger"
@ -165,224 +162,10 @@ func (s *singleton) Generate(ctx context.Context, input models.GenerateMetadataI
logger.Warnf("could not generate temporary directory: %v", err)
}
sceneIDs, err := utils.StringSliceToIntSlice(input.SceneIDs)
if err != nil {
logger.Error(err.Error())
j := &GenerateJob{
txnManager: s.TxnManager,
input: input,
}
markerIDs, err := utils.StringSliceToIntSlice(input.MarkerIDs)
if err != nil {
logger.Error(err.Error())
}
// TODO - formalise this
j := job.MakeJobExec(func(ctx context.Context, progress *job.Progress) {
var scenes []*models.Scene
var err error
var markers []*models.SceneMarker
if err := s.TxnManager.WithReadTxn(context.TODO(), func(r models.ReaderRepository) error {
qb := r.Scene()
if len(sceneIDs) > 0 {
scenes, err = qb.FindMany(sceneIDs)
} else {
scenes, err = qb.All()
}
if err != nil {
return err
}
if len(markerIDs) > 0 {
markers, err = r.SceneMarker().FindMany(markerIDs)
if err != nil {
return err
}
}
return nil
}); err != nil {
logger.Error(err.Error())
return
}
config := config.GetInstance()
parallelTasks := config.GetParallelTasksWithAutoDetection()
logger.Infof("Generate started with %d parallel tasks", parallelTasks)
wg := sizedwaitgroup.New(parallelTasks)
lenScenes := len(scenes)
total := lenScenes + len(markers)
progress.SetTotal(total)
if job.IsCancelled(ctx) {
logger.Info("Stopping due to user request")
return
}
// TODO - consider removing this. Even though we're only waiting a maximum of
// 90 seconds for this, it is all for a simple log message, and probably not worth
// waiting for
var totalsNeeded *totalsGenerate
progress.ExecuteTask("Calculating content to generate...", func() {
totalsNeeded = s.neededGenerate(scenes, input)
if totalsNeeded == nil {
logger.Infof("Taking too long to count content. Skipping...")
logger.Infof("Generating content")
} else {
logger.Infof("Generating %d sprites %d previews %d image previews %d markers %d transcodes %d phashes", totalsNeeded.sprites, totalsNeeded.previews, totalsNeeded.imagePreviews, totalsNeeded.markers, totalsNeeded.transcodes, totalsNeeded.phashes)
}
})
fileNamingAlgo := config.GetVideoFileNamingAlgorithm()
overwrite := false
if input.Overwrite != nil {
overwrite = *input.Overwrite
}
generatePreviewOptions := input.PreviewOptions
if generatePreviewOptions == nil {
generatePreviewOptions = &models.GeneratePreviewOptionsInput{}
}
setGeneratePreviewOptionsInput(generatePreviewOptions)
// 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)
}
for _, scene := range scenes {
progress.Increment()
if job.IsCancelled(ctx) {
logger.Info("Stopping due to user request")
wg.Wait()
if err := instance.Paths.Generated.EmptyTmpDir(); err != nil {
logger.Warnf("failure emptying temporary directory: %v", err)
}
return
}
if scene == nil {
logger.Errorf("nil scene, skipping generate")
continue
}
if utils.IsTrue(input.Sprites) {
task := GenerateSpriteTask{
Scene: *scene,
Overwrite: overwrite,
fileNamingAlgorithm: fileNamingAlgo,
}
wg.Add()
go progress.ExecuteTask(fmt.Sprintf("Generating sprites for %s", scene.Path), func() {
task.Start()
wg.Done()
})
}
if utils.IsTrue(input.Previews) {
task := GeneratePreviewTask{
Scene: *scene,
ImagePreview: utils.IsTrue(input.ImagePreviews),
Options: *generatePreviewOptions,
Overwrite: overwrite,
fileNamingAlgorithm: fileNamingAlgo,
}
wg.Add()
go progress.ExecuteTask(fmt.Sprintf("Generating preview for %s", scene.Path), func() {
task.Start()
wg.Done()
})
}
if utils.IsTrue(input.Markers) {
wg.Add()
task := GenerateMarkersTask{
TxnManager: s.TxnManager,
Scene: scene,
Overwrite: overwrite,
fileNamingAlgorithm: fileNamingAlgo,
ImagePreview: utils.IsTrue(input.MarkerImagePreviews),
Screenshot: utils.IsTrue(input.MarkerScreenshots),
}
go progress.ExecuteTask(fmt.Sprintf("Generating markers for %s", scene.Path), func() {
task.Start()
wg.Done()
})
}
if utils.IsTrue(input.Transcodes) {
wg.Add()
task := GenerateTranscodeTask{
Scene: *scene,
Overwrite: overwrite,
fileNamingAlgorithm: fileNamingAlgo,
}
go progress.ExecuteTask(fmt.Sprintf("Generating transcode for %s", scene.Path), func() {
task.Start()
wg.Done()
})
}
if utils.IsTrue(input.Phashes) {
task := GeneratePhashTask{
Scene: *scene,
fileNamingAlgorithm: fileNamingAlgo,
txnManager: s.TxnManager,
Overwrite: overwrite,
}
wg.Add()
go progress.ExecuteTask(fmt.Sprintf("Generating phash for %s", scene.Path), func() {
task.Start()
wg.Done()
})
}
}
wg.Wait()
for _, marker := range markers {
progress.Increment()
if job.IsCancelled(ctx) {
logger.Info("Stopping due to user request")
wg.Wait()
if err := instance.Paths.Generated.EmptyTmpDir(); err != nil {
logger.Warnf("failure emptying temporary directory: %v", err)
}
elapsed := time.Since(start)
logger.Info(fmt.Sprintf("Generate finished (%s)", elapsed))
return
}
if marker == nil {
logger.Errorf("nil marker, skipping generate")
continue
}
wg.Add()
task := GenerateMarkersTask{
TxnManager: s.TxnManager,
Marker: marker,
Overwrite: overwrite,
fileNamingAlgorithm: fileNamingAlgo,
}
go progress.ExecuteTask(fmt.Sprintf("Generating marker preview for marker ID %d", marker.ID), func() {
task.Start()
wg.Done()
})
}
wg.Wait()
if err = instance.Paths.Generated.EmptyTmpDir(); err != nil {
logger.Warnf("failure emptying temporary directory: %v", err)
}
elapsed := time.Since(start)
logger.Info(fmt.Sprintf("Generate finished (%s)", elapsed))
})
return s.JobManager.Add(ctx, "Generating...", j), nil
}
@ -425,7 +208,7 @@ func (s *singleton) generateScreenshot(ctx context.Context, sceneId string, at *
fileNamingAlgorithm: config.GetInstance().GetVideoFileNamingAlgorithm(),
}
task.Start()
task.Start(ctx)
logger.Infof("Generate screenshot finished")
})
@ -500,103 +283,6 @@ func (s *singleton) MigrateHash(ctx context.Context) int {
return s.JobManager.Add(ctx, "Migrating scene hashes...", j)
}
type totalsGenerate struct {
sprites int64
previews int64
imagePreviews int64
markers int64
transcodes int64
phashes int64
}
func (s *singleton) neededGenerate(scenes []*models.Scene, input models.GenerateMetadataInput) *totalsGenerate {
var totals totalsGenerate
const timeout = 90 * time.Second
// Set a deadline.
chTimeout := time.After(timeout)
fileNamingAlgo := config.GetInstance().GetVideoFileNamingAlgorithm()
overwrite := false
if input.Overwrite != nil {
overwrite = *input.Overwrite
}
logger.Infof("Counting content to generate...")
for _, scene := range scenes {
if scene != nil {
if utils.IsTrue(input.Sprites) {
task := GenerateSpriteTask{
Scene: *scene,
fileNamingAlgorithm: fileNamingAlgo,
}
if overwrite || task.required() {
totals.sprites++
}
}
if utils.IsTrue(input.Previews) {
task := GeneratePreviewTask{
Scene: *scene,
ImagePreview: utils.IsTrue(input.ImagePreviews),
fileNamingAlgorithm: fileNamingAlgo,
}
sceneHash := scene.GetHash(task.fileNamingAlgorithm)
if overwrite || !task.doesVideoPreviewExist(sceneHash) {
totals.previews++
}
if utils.IsTrue(input.ImagePreviews) && (overwrite || !task.doesImagePreviewExist(sceneHash)) {
totals.imagePreviews++
}
}
if utils.IsTrue(input.Markers) {
task := GenerateMarkersTask{
TxnManager: s.TxnManager,
Scene: scene,
Overwrite: overwrite,
fileNamingAlgorithm: fileNamingAlgo,
}
totals.markers += int64(task.isMarkerNeeded())
}
if utils.IsTrue(input.Transcodes) {
task := GenerateTranscodeTask{
Scene: *scene,
Overwrite: overwrite,
fileNamingAlgorithm: fileNamingAlgo,
}
if task.isTranscodeNeeded() {
totals.transcodes++
}
}
if utils.IsTrue(input.Phashes) {
task := GeneratePhashTask{
Scene: *scene,
fileNamingAlgorithm: fileNamingAlgo,
}
if task.shouldGenerate() {
totals.phashes++
}
}
}
// check for timeout
select {
case <-chTimeout:
return nil
default:
}
}
return &totals
}
func (s *singleton) StashBoxBatchPerformerTag(ctx context.Context, input models.StashBoxBatchPerformerTagInput) int {
j := job.MakeJobExec(func(ctx context.Context, progress *job.Progress) {
logger.Infof("Initiating stash-box batch performer tag")

288
pkg/manager/task_generate.go Normal file
View File

@ -0,0 +1,288 @@
package manager
import (
"context"
"errors"
"fmt"
"time"
"github.com/remeh/sizedwaitgroup"
"github.com/stashapp/stash/pkg/job"
"github.com/stashapp/stash/pkg/logger"
"github.com/stashapp/stash/pkg/manager/config"
"github.com/stashapp/stash/pkg/models"
"github.com/stashapp/stash/pkg/utils"
)
const generateQueueSize = 200000
type GenerateJob struct {
txnManager models.TransactionManager
input models.GenerateMetadataInput
overwrite bool
fileNamingAlgo models.HashAlgorithm
}
type totalsGenerate struct {
sprites int64
previews int64
imagePreviews int64
markers int64
transcodes int64
phashes 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() {
var totals totalsGenerate
sceneIDs, err := utils.StringSliceToIntSlice(j.input.SceneIDs)
if err != nil {
logger.Error(err.Error())
}
markerIDs, err := utils.StringSliceToIntSlice(j.input.MarkerIDs)
if err != nil {
logger.Error(err.Error())
}
if err := j.txnManager.WithReadTxn(ctx, func(r models.ReaderRepository) error {
qb := r.Scene()
if len(j.input.SceneIDs) == 0 && len(j.input.MarkerIDs) == 0 {
totals = j.queueTasks(ctx, queue)
} else {
if len(j.input.SceneIDs) > 0 {
scenes, err = qb.FindMany(sceneIDs)
for _, s := range scenes {
j.queueSceneJobs(s, queue, &totals)
}
}
if len(j.input.MarkerIDs) > 0 {
markers, err = r.SceneMarker().FindMany(markerIDs)
if err != nil {
return err
}
for _, m := range markers {
j.queueMarkerJob(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", totals.sprites, totals.previews, totals.imagePreviews, totals.markers, totals.transcodes, totals.phashes)
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()
go progress.ExecuteTask(f.GetDescription(), func() {
f.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, queue chan<- Task) totalsGenerate {
defer close(queue)
var totals totalsGenerate
const batchSize = 1000
findFilter := models.BatchFindFilter(batchSize)
if err := j.txnManager.WithReadTxn(ctx, func(r models.ReaderRepository) error {
for more := true; more; {
if job.IsCancelled(ctx) {
return context.Canceled
}
scenes, _, err := r.Scene().Query(nil, findFilter)
if err != nil {
return err
}
for _, ss := range scenes {
if job.IsCancelled(ctx) {
return context.Canceled
}
j.queueSceneJobs(ss, queue, &totals)
}
if len(scenes) != batchSize {
more = false
} else {
*findFilter.Page++
}
}
return nil
}); err != nil {
if !errors.Is(err, context.Canceled) {
logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
}
}
return totals
}
func (j *GenerateJob) queueSceneJobs(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
}
}
if utils.IsTrue(j.input.Previews) {
generatePreviewOptions := j.input.PreviewOptions
if generatePreviewOptions == nil {
generatePreviewOptions = &models.GeneratePreviewOptionsInput{}
}
setGeneratePreviewOptionsInput(generatePreviewOptions)
task := &GeneratePreviewTask{
Scene: *scene,
ImagePreview: utils.IsTrue(j.input.ImagePreviews),
Options: *generatePreviewOptions,
Overwrite: j.overwrite,
fileNamingAlgorithm: j.fileNamingAlgo,
}
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),
}
markers := task.markersNeeded()
if markers > 0 {
totals.markers += int64(markers)
totals.tasks++
queue <- task
}
}
if utils.IsTrue(j.input.Transcodes) {
task := &GenerateTranscodeTask{
Scene: *scene,
Overwrite: j.overwrite,
fileNamingAlgorithm: j.fileNamingAlgo,
}
if task.isTranscodeNeeded() {
totals.transcodes++
totals.tasks++
queue <- task
}
}
if utils.IsTrue(j.input.Phashes) {
task := &GeneratePhashTask{
Scene: *scene,
fileNamingAlgorithm: j.fileNamingAlgo,
txnManager: j.txnManager,
Overwrite: j.overwrite,
}
if task.shouldGenerate() {
totals.phashes++
totals.tasks++
queue <- task
}
}
}
func (j *GenerateJob) queueMarkerJob(marker *models.SceneMarker, queue chan<- Task, totals *totalsGenerate) {
task := &GenerateMarkersTask{
TxnManager: j.txnManager,
Marker: marker,
Overwrite: j.overwrite,
fileNamingAlgorithm: j.fileNamingAlgo,
}
totals.markers++
totals.tasks++
queue <- task
}

15
pkg/manager/task_generate_markers.go
View File

@ -2,6 +2,7 @@ package manager
import (
"context"
"fmt"
"path/filepath"
"strconv"
@ -22,7 +23,17 @@ type GenerateMarkersTask struct {
Screenshot bool
}
func (t *GenerateMarkersTask) Start() {
func (t *GenerateMarkersTask) GetDescription() string {
if t.Scene != nil {
return fmt.Sprintf("Generating markers for %s", t.Scene.Path)
} else if t.Marker != nil {
return fmt.Sprintf("Generating marker preview for marker ID %d", t.Marker.ID)
}
return "Generating markers"
}
func (t *GenerateMarkersTask) Start(ctx context.Context) {
if t.Scene != nil {
t.generateSceneMarkers()
}
@ -155,7 +166,7 @@ func (t *GenerateMarkersTask) generateMarker(videoFile *ffmpeg.VideoFile, scene
}
}
func (t *GenerateMarkersTask) isMarkerNeeded() int {
func (t *GenerateMarkersTask) markersNeeded() int {
markers := 0
var sceneMarkers []*models.SceneMarker
if err := t.TxnManager.WithReadTxn(context.TODO(), func(r models.ReaderRepository) error {

7
pkg/manager/task_generate_phash.go
View File

@ -3,6 +3,7 @@ package manager
import (
"context"
"database/sql"
"fmt"
"github.com/stashapp/stash/pkg/logger"
"github.com/stashapp/stash/pkg/models"
@ -15,7 +16,11 @@ type GeneratePhashTask struct {
txnManager models.TransactionManager
}
func (t *GeneratePhashTask) Start() {
func (t *GeneratePhashTask) GetDescription() string {
return fmt.Sprintf("Generating phash for %s", t.Scene.Path)
}
func (t *GeneratePhashTask) Start(ctx context.Context) {
if !t.shouldGenerate() {
return
}

9
pkg/manager/task_generate_preview.go
View File

@ -1,6 +1,9 @@
package manager
import (
"context"
"fmt"
"github.com/stashapp/stash/pkg/logger"
"github.com/stashapp/stash/pkg/manager/config"
"github.com/stashapp/stash/pkg/models"
@ -17,7 +20,11 @@ type GeneratePreviewTask struct {
fileNamingAlgorithm models.HashAlgorithm
}
func (t *GeneratePreviewTask) Start() {
func (t *GeneratePreviewTask) GetDescription() string {
return fmt.Sprintf("Generating preview for %s", t.Scene.Path)
}
func (t *GeneratePreviewTask) Start(ctx context.Context) {
videoFilename := t.videoFilename()
videoChecksum := t.Scene.GetHash(t.fileNamingAlgorithm)
imageFilename := t.imageFilename()

2
pkg/manager/task_generate_screenshot.go
View File

@ -18,7 +18,7 @@ type GenerateScreenshotTask struct {
txnManager models.TransactionManager
}
func (t *GenerateScreenshotTask) Start() {
func (t *GenerateScreenshotTask) Start(ctx context.Context) {
scenePath := t.Scene.Path
ffprobe := instance.FFProbe
probeResult, err := ffprobe.NewVideoFile(scenePath, false)

9
pkg/manager/task_generate_sprite.go
View File

@ -1,6 +1,9 @@
package manager
import (
"context"
"fmt"
"github.com/stashapp/stash/pkg/logger"
"github.com/stashapp/stash/pkg/models"
"github.com/stashapp/stash/pkg/utils"
@ -12,7 +15,11 @@ type GenerateSpriteTask struct {
fileNamingAlgorithm models.HashAlgorithm
}
func (t *GenerateSpriteTask) Start() {
func (t *GenerateSpriteTask) GetDescription() string {
return fmt.Sprintf("Generating sprites for %s", t.Scene.Path)
}
func (t *GenerateSpriteTask) Start(ctx context.Context) {
if !t.Overwrite && !t.required() {
return
}

6
pkg/manager/task_scan.go
View File

@ -286,7 +286,7 @@ func (t *ScanTask) Start(ctx context.Context) {
Overwrite: false,
fileNamingAlgorithm: t.fileNamingAlgorithm,
}
taskSprite.Start()
taskSprite.Start(ctx)
iwg.Done()
})
}
@ -300,7 +300,7 @@ func (t *ScanTask) Start(ctx context.Context) {
fileNamingAlgorithm: t.fileNamingAlgorithm,
txnManager: t.TxnManager,
}
taskPhash.Start()
taskPhash.Start(ctx)
iwg.Done()
})
}
@ -332,7 +332,7 @@ func (t *ScanTask) Start(ctx context.Context) {
Overwrite: false,
fileNamingAlgorithm: t.fileNamingAlgorithm,
}
taskPreview.Start()
taskPreview.Start(ctx)
iwg.Done()
})
}

9
pkg/manager/task_transcode.go
View File

@ -1,6 +1,9 @@
package manager
import (
"context"
"fmt"
"github.com/stashapp/stash/pkg/ffmpeg"
"github.com/stashapp/stash/pkg/logger"
"github.com/stashapp/stash/pkg/manager/config"
@ -14,7 +17,11 @@ type GenerateTranscodeTask struct {
fileNamingAlgorithm models.HashAlgorithm
}
func (t *GenerateTranscodeTask) Start() {
func (t *GenerateTranscodeTask) GetDescription() string {
return fmt.Sprintf("Generating transcode for %s", t.Scene.Path)
}
func (t *GenerateTranscodeTask) Start(ctc context.Context) {
hasTranscode := HasTranscode(&t.Scene, t.fileNamingAlgorithm)
if !t.Overwrite && hasTranscode {
return

2
ui/v2.5/src/components/Changelog/versions/v0110.md
View File

@ -5,11 +5,13 @@
* Added interface options to disable creating performers/studios/tags from dropdown selectors. ([#1814](https://github.com/stashapp/stash/pull/1814))
### 🎨 Improvements
* Optimised generate process. ([#1871](https://github.com/stashapp/stash/pull/1871))
* Added clear button to query text field. ([#1845](https://github.com/stashapp/stash/pull/1845))
* Moved Performer rating stars from details/edit tabs to heading section of performer page. ([#1844](https://github.com/stashapp/stash/pull/1844))
* Optimised scanning process. ([#1816](https://github.com/stashapp/stash/pull/1816))
### 🐛 Bug fixes
* Fix marker generation task reading video files unnecessarily. ([#1871](https://github.com/stashapp/stash/pull/1871))
* Fix accessing Stash via IPv6 link local address causing security tripwire to be activated. ([#1841](https://github.com/stashapp/stash/pull/1841))
* Fix Twitter value defaulting to freeones in built-in Freeones scraper. ([#1853](https://github.com/stashapp/stash/pull/1853))
* Fix colour codes not outputting correctly when logging to file on Windows. ([#1846](https://github.com/stashapp/stash/pull/1846))