From 22577918fb89fca3e59ba654d6783001163e6a06 Mon Sep 17 00:00:00 2001
From: WithoutPants <53250216+WithoutPants@users.noreply.github.com>
Date: Thu, 25 Jul 2019 08:17:22 +1000
Subject: [PATCH] Add basic live transcoding to webm

---
 pkg/api/routes_scene.go         | 45 ++++++++++++++++++++++++++++++++-
 pkg/ffmpeg/encoder.go           | 20 ++++++++++++++-
 pkg/ffmpeg/encoder_transcode.go | 20 +++++++++++++++
 3 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/pkg/api/routes_scene.go b/pkg/api/routes_scene.go
index cc7572985..965a6cb1a 100644
--- a/pkg/api/routes_scene.go
+++ b/pkg/api/routes_scene.go
@@ -1,12 +1,14 @@
 package api
 
 import (
+	"io"
 	"context"
 	"github.com/go-chi/chi"
 	"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"
@@ -39,8 +41,49 @@ 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)
-	http.ServeFile(w, r, filepath)
+
+	videoCodec := scene.VideoCodec.String
+	hasTranscode, _ := manager.HasTranscode(scene)
+	if ffmpeg.IsValidCodec(videoCodec) || hasTranscode {
+		http.ServeFile(w, r, filepath)
+		return
+	}
+
+	// needs to be transcoded
+	videoFile, err := ffmpeg.NewVideoFile(manager.GetInstance().FFProbePath, scene.Path)
+	if err != nil {
+		logger.Errorf("[stream] error reading video file: %s", err.Error())
+		return
+	}
+	
+	encoder := ffmpeg.NewEncoder(manager.GetInstance().FFMPEGPath)
+
+	stream, process, err := encoder.StreamTranscode(*videoFile)
+	if err != nil {
+		logger.Errorf("[stream] error transcoding video file: %s", err.Error())
+		return
+	}
+
+	w.WriteHeader(http.StatusOK)
+	w.Header().Set("Content-Type", "video/webm")
+
+	logger.Info("[stream] transcoding video file")
+
+	// handle if client closes the connection
+	notify := r.Context().Done()
+	go func() {
+		<-notify
+		logger.Info("[stream] client closed the connection. Killing stream process.")
+		process.Kill()
+	}()
+
+	_, err = io.Copy(w, stream)
+	if err != nil {
+		logger.Errorf("[stream] error serving transcoded video file: %s", err.Error())
+	}
 }
 
 func (rs sceneRoutes) Screenshot(w http.ResponseWriter, r *http.Request) {
diff --git a/pkg/ffmpeg/encoder.go b/pkg/ffmpeg/encoder.go
index da9d4d7a5..caf6413b6 100644
--- a/pkg/ffmpeg/encoder.go
+++ b/pkg/ffmpeg/encoder.go
@@ -1,10 +1,13 @@
 package ffmpeg
 
 import (
-	"github.com/stashapp/stash/pkg/logger"
+	"io"
 	"io/ioutil"
+	"os"
 	"os/exec"
 	"strings"
+
+	"github.com/stashapp/stash/pkg/logger"
 )
 
 type Encoder struct {
@@ -60,3 +63,18 @@ func (e *Encoder) run(probeResult VideoFile, args []string) (string, error) {
 
 	return stdoutString, nil
 }
+
+func (e *Encoder) stream(probeResult VideoFile, args []string) (io.ReadCloser, *os.Process, error) {
+	cmd := exec.Command(e.Path, args...)
+
+	stdout, err := cmd.StdoutPipe()
+	if nil != err {
+		logger.Error("FFMPEG stdout not available: " + err.Error())
+	}
+
+	if err = cmd.Start(); err != nil {
+		return nil, nil, err
+	}
+
+	return stdout, cmd.Process, nil
+}
diff --git a/pkg/ffmpeg/encoder_transcode.go b/pkg/ffmpeg/encoder_transcode.go
index 79b4e8954..d8942b36b 100644
--- a/pkg/ffmpeg/encoder_transcode.go
+++ b/pkg/ffmpeg/encoder_transcode.go
@@ -1,5 +1,10 @@
 package ffmpeg
 
+import (
+	"io"
+	"os"
+)
+
 type TranscodeOptions struct {
 	OutputPath string
 }
@@ -20,3 +25,18 @@ func (e *Encoder) Transcode(probeResult VideoFile, options TranscodeOptions) {
 	}
 	_, _ = e.run(probeResult, args)
 }
+
+func (e *Encoder) StreamTranscode(probeResult VideoFile) (io.ReadCloser, *os.Process, error) {
+	args := []string{
+		"-i", probeResult.Path,
+		"-c:v", "libvpx-vp9",
+		"-vf", "scale=iw:-2",
+		"-deadline", "realtime",
+		"-cpu-used", "5",
+		"-crf", "30",
+		"-b:v", "0",
+		"-f", "webm",
+		"pipe:",
+	}
+	return e.stream(probeResult, args)
+}