stash/pkg/manager/task_generate_screenshot.go
SmallCoccinelle 4a0c4c4847
Reorder waitgroup completion (#1748)
Rather than passing a pointer to a waitgroup into task.Start(..)
functions, handle the waitgroup.Done() at the callsite.

This makes waitgroup handling local to its definition rather than it
being spread out over multiple files. Tasks now simply execute, and
the policy of waiting on them is handled by the caller.
2021-09-22 13:22:59 +10:00

88 lines
2.1 KiB
Go

package manager
import (
"context"
"fmt"
"io/ioutil"
"os"
"time"
"github.com/stashapp/stash/pkg/ffmpeg"
"github.com/stashapp/stash/pkg/logger"
"github.com/stashapp/stash/pkg/models"
)
type GenerateScreenshotTask struct {
Scene models.Scene
ScreenshotAt *float64
fileNamingAlgorithm models.HashAlgorithm
txnManager models.TransactionManager
}
func (t *GenerateScreenshotTask) Start() {
scenePath := t.Scene.Path
probeResult, err := ffmpeg.NewVideoFile(instance.FFProbePath, scenePath, false)
if err != nil {
logger.Error(err.Error())
return
}
var at float64
if t.ScreenshotAt == nil {
at = float64(probeResult.Duration) * 0.2
} else {
at = *t.ScreenshotAt
}
checksum := t.Scene.GetHash(t.fileNamingAlgorithm)
normalPath := instance.Paths.Scene.GetScreenshotPath(checksum)
// we'll generate the screenshot, grab the generated data and set it
// in the database. We'll use SetSceneScreenshot to set the data
// which also generates the thumbnail
logger.Debugf("Creating screenshot for %s", scenePath)
makeScreenshot(*probeResult, normalPath, 2, probeResult.Width, at)
f, err := os.Open(normalPath)
if err != nil {
logger.Errorf("Error reading screenshot: %s", err.Error())
return
}
defer f.Close()
coverImageData, err := ioutil.ReadAll(f)
if err != nil {
logger.Errorf("Error reading screenshot: %s", err.Error())
return
}
if err := t.txnManager.WithTxn(context.TODO(), func(r models.Repository) error {
qb := r.Scene()
updatedTime := time.Now()
updatedScene := models.ScenePartial{
ID: t.Scene.ID,
UpdatedAt: &models.SQLiteTimestamp{Timestamp: updatedTime},
}
if err := SetSceneScreenshot(checksum, coverImageData); err != nil {
return fmt.Errorf("error writing screenshot: %v", err)
}
// update the scene cover table
if err := qb.UpdateCover(t.Scene.ID, coverImageData); err != nil {
return fmt.Errorf("error setting screenshot: %v", err)
}
// update the scene with the update date
_, err = qb.Update(updatedScene)
if err != nil {
return fmt.Errorf("error updating scene: %v", err)
}
return nil
}); err != nil {
logger.Error(err.Error())
}
}