diff --git a/pkg/api/resolver_mutation_scene.go b/pkg/api/resolver_mutation_scene.go index c6dab2dc3..354f43822 100644 --- a/pkg/api/resolver_mutation_scene.go +++ b/pkg/api/resolver_mutation_scene.go @@ -157,6 +157,9 @@ func (r *mutationResolver) SceneDestroy(ctx context.Context, input models.SceneD // if delete file is true, then delete the file as well // if it fails, just log a message if input.DeleteFile != nil && *input.DeleteFile { + // kill any running encoders + manager.KillRunningStreams(scene.Path) + err = os.Remove(scene.Path) if err != nil { logger.Warnf("Could not delete file %s: %s", scene.Path, err.Error()) @@ -216,6 +219,9 @@ func deleteGeneratedSceneFiles(scene *models.Scene) { transcodePath := manager.GetInstance().Paths.Scene.GetTranscodePath(scene.Checksum) exists, _ = utils.FileExists(transcodePath) if exists { + // kill any running streams + manager.KillRunningStreams(transcodePath) + err := os.Remove(transcodePath) if err != nil { logger.Warnf("Could not delete file %s: %s", transcodePath, err.Error()) diff --git a/pkg/api/routes_scene.go b/pkg/api/routes_scene.go index 965a6cb1a..1a06d79c9 100644 --- a/pkg/api/routes_scene.go +++ b/pkg/api/routes_scene.go @@ -1,17 +1,18 @@ package api import ( - "io" "context" + "io" + "net/http" + "strconv" + "strings" + "github.com/go-chi/chi" + "github.com/stashapp/stash/pkg/ffmpeg" "github.com/stashapp/stash/pkg/logger" "github.com/stashapp/stash/pkg/manager" "github.com/stashapp/stash/pkg/models" "github.com/stashapp/stash/pkg/utils" - "github.com/stashapp/stash/pkg/ffmpeg" - "net/http" - "strconv" - "strings" ) type sceneRoutes struct{} @@ -41,14 +42,16 @@ func (rs sceneRoutes) Routes() chi.Router { func (rs sceneRoutes) Stream(w http.ResponseWriter, r *http.Request) { scene := r.Context().Value(sceneKey).(*models.Scene) - + // detect if not a streamable file and try to transcode it instead filepath := manager.GetInstance().Paths.Scene.GetStreamPath(scene.Path, scene.Checksum) - videoCodec := scene.VideoCodec.String hasTranscode, _ := manager.HasTranscode(scene) if ffmpeg.IsValidCodec(videoCodec) || hasTranscode { + manager.RegisterStream(filepath, &w) http.ServeFile(w, r, filepath) + manager.WaitAndDeregisterStream(filepath, &w, r) + return } @@ -58,7 +61,7 @@ func (rs sceneRoutes) Stream(w http.ResponseWriter, r *http.Request) { logger.Errorf("[stream] error reading video file: %s", err.Error()) return } - + encoder := ffmpeg.NewEncoder(manager.GetInstance().FFMPEGPath) stream, process, err := encoder.StreamTranscode(*videoFile) diff --git a/pkg/ffmpeg/encoder.go b/pkg/ffmpeg/encoder.go index caf6413b6..d3cf909d1 100644 --- a/pkg/ffmpeg/encoder.go +++ b/pkg/ffmpeg/encoder.go @@ -1,11 +1,13 @@ package ffmpeg import ( + "fmt" "io" "io/ioutil" "os" "os/exec" "strings" + "time" "github.com/stashapp/stash/pkg/logger" ) @@ -14,12 +16,63 @@ type Encoder struct { Path string } +var runningEncoders map[string][]*os.Process = make(map[string][]*os.Process) + func NewEncoder(ffmpegPath string) Encoder { return Encoder{ Path: ffmpegPath, } } +func registerRunningEncoder(path string, process *os.Process) { + processes := runningEncoders[path] + + runningEncoders[path] = append(processes, process) +} + +func deregisterRunningEncoder(path string, process *os.Process) { + processes := runningEncoders[path] + + for i, v := range processes { + if v == process { + runningEncoders[path] = append(processes[:i], processes[i+1:]...) + return + } + } +} + +func waitAndDeregister(path string, cmd *exec.Cmd) error { + err := cmd.Wait() + deregisterRunningEncoder(path, cmd.Process) + + return err +} + +func KillRunningEncoders(path string) { + processes := runningEncoders[path] + + for _, process := range processes { + // assume it worked, don't check for error + fmt.Printf("Killing encoder process for file: %s", path) + process.Kill() + + // wait for the process to die before returning + // don't wait more than a few seconds + done := make(chan error) + go func() { + _, err := process.Wait() + done <- err + }() + + select { + case <-done: + return + case <-time.After(5 * time.Second): + return + } + } +} + func (e *Encoder) run(probeResult VideoFile, args []string) (string, error) { cmd := exec.Command(e.Path, args...) @@ -56,7 +109,10 @@ func (e *Encoder) run(probeResult VideoFile, args []string) (string, error) { stdoutData, _ := ioutil.ReadAll(stdout) stdoutString := string(stdoutData) - if err := cmd.Wait(); err != nil { + registerRunningEncoder(probeResult.Path, cmd.Process) + err = waitAndDeregister(probeResult.Path, cmd) + + if err != nil { logger.Errorf("ffmpeg error when running command <%s>", strings.Join(cmd.Args, " ")) return stdoutString, err } @@ -76,5 +132,8 @@ func (e *Encoder) stream(probeResult VideoFile, args []string) (io.ReadCloser, * return nil, nil, err } + registerRunningEncoder(probeResult.Path, cmd.Process) + go waitAndDeregister(probeResult.Path, cmd) + return stdout, cmd.Process, nil } diff --git a/pkg/manager/running_streams.go b/pkg/manager/running_streams.go new file mode 100644 index 000000000..f573d7aa4 --- /dev/null +++ b/pkg/manager/running_streams.go @@ -0,0 +1,58 @@ +package manager + +import ( + "net/http" + + "github.com/stashapp/stash/pkg/ffmpeg" + "github.com/stashapp/stash/pkg/logger" +) + +var streamingFiles = make(map[string][]*http.ResponseWriter) + +func RegisterStream(filepath string, w *http.ResponseWriter) { + streams := streamingFiles[filepath] + + streamingFiles[filepath] = append(streams, w) +} + +func deregisterStream(filepath string, w *http.ResponseWriter) { + streams := streamingFiles[filepath] + + for i, v := range streams { + if v == w { + streamingFiles[filepath] = append(streams[:i], streams[i+1:]...) + return + } + } +} + +func WaitAndDeregisterStream(filepath string, w *http.ResponseWriter, r *http.Request) { + notify := r.Context().Done() + go func() { + <-notify + deregisterStream(filepath, w) + }() +} + +func KillRunningStreams(path string) { + ffmpeg.KillRunningEncoders(path) + + streams := streamingFiles[path] + + for _, w := range streams { + hj, ok := (*w).(http.Hijacker) + if !ok { + // if we can't close the connection can't really do anything else + logger.Warnf("cannot close running stream for: %s", path) + return + } + + // hijack and close the connection + conn, _, err := hj.Hijack() + if err != nil { + logger.Errorf("cannot close running stream for '%s' due to error: %s", path, err.Error()) + } else { + conn.Close() + } + } +}