Close streams/encodes before deleting file

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())

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)

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
 }

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()
+		}
+	}
+}