From 2cd71785f403a55aa8db91a8c7b41aace1884873 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Fri, 3 Jun 2022 13:29:26 +0200
Subject: [PATCH] hls muxer: cache init.mp4 content

---
 internal/hls/muxer_variant_fmp4.go          | 51 +++++++++++++++++++--
 internal/hls/muxer_variant_fmp4_playlist.go | 17 -------
 2 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/internal/hls/muxer_variant_fmp4.go b/internal/hls/muxer_variant_fmp4.go
index 5f6e23bc..54a1b38a 100644
--- a/internal/hls/muxer_variant_fmp4.go
+++ b/internal/hls/muxer_variant_fmp4.go
@@ -1,6 +1,9 @@
 package hls
 
 import (
+	"bytes"
+	"net/http"
+	"sync"
 	"time"
 
 	"github.com/aler9/gortsplib"
@@ -34,8 +37,15 @@ func (s fmp4AudioSample) duration() time.Duration {
 }
 
 type muxerVariantFMP4 struct {
-	playlist  *muxerVariantFMP4Playlist
-	segmenter *muxerVariantFMP4Segmenter
+	playlist   *muxerVariantFMP4Playlist
+	segmenter  *muxerVariantFMP4Segmenter
+	videoTrack *gortsplib.TrackH264
+	audioTrack *gortsplib.TrackAAC
+
+	mutex        sync.Mutex
+	videoLastSPS []byte
+	videoLastPPS []byte
+	initContent  []byte
 }
 
 func newMuxerVariantFMP4(
@@ -47,7 +57,10 @@ func newMuxerVariantFMP4(
 	videoTrack *gortsplib.TrackH264,
 	audioTrack *gortsplib.TrackAAC,
 ) *muxerVariantFMP4 {
-	v := &muxerVariantFMP4{}
+	v := &muxerVariantFMP4{
+		videoTrack: videoTrack,
+		audioTrack: audioTrack,
+	}
 
 	v.playlist = newMuxerVariantFMP4Playlist(
 		lowLatency,
@@ -84,5 +97,37 @@ func (v *muxerVariantFMP4) writeAAC(pts time.Duration, aus [][]byte) error {
 }
 
 func (v *muxerVariantFMP4) file(name string, msn string, part string, skip string) *MuxerFileResponse {
+	if name == "init.mp4" {
+		v.mutex.Lock()
+		defer v.mutex.Unlock()
+
+		var sps []byte
+		var pps []byte
+		if v.videoTrack != nil {
+			sps = v.videoTrack.SPS()
+			pps = v.videoTrack.PPS()
+		}
+
+		if v.initContent == nil ||
+			(v.videoTrack != nil && (!bytes.Equal(v.videoLastSPS, sps) || !bytes.Equal(v.videoLastPPS, pps))) {
+			initContent, err := mp4InitGenerate(v.videoTrack, v.audioTrack)
+			if err != nil {
+				return &MuxerFileResponse{Status: http.StatusInternalServerError}
+			}
+
+			v.videoLastSPS = sps
+			v.videoLastPPS = pps
+			v.initContent = initContent
+		}
+
+		return &MuxerFileResponse{
+			Status: http.StatusOK,
+			Header: map[string]string{
+				"Content-Type": "video/mp4",
+			},
+			Body: bytes.NewReader(v.initContent),
+		}
+	}
+
 	return v.playlist.file(name, msn, part, skip)
 }
diff --git a/internal/hls/muxer_variant_fmp4_playlist.go b/internal/hls/muxer_variant_fmp4_playlist.go
index 8ba1b7b7..fa72edf3 100644
--- a/internal/hls/muxer_variant_fmp4_playlist.go
+++ b/internal/hls/muxer_variant_fmp4_playlist.go
@@ -359,23 +359,6 @@ func (p *muxerVariantFMP4Playlist) fullPlaylist(isDeltaUpdate bool) io.Reader {
 
 func (p *muxerVariantFMP4Playlist) segmentReader(fname string) *MuxerFileResponse {
 	switch {
-	case fname == "init.mp4":
-		p.mutex.Lock()
-		defer p.mutex.Unlock()
-
-		byts, err := mp4InitGenerate(p.videoTrack, p.audioTrack)
-		if err != nil {
-			return &MuxerFileResponse{Status: http.StatusInternalServerError}
-		}
-
-		return &MuxerFileResponse{
-			Status: http.StatusOK,
-			Header: map[string]string{
-				"Content-Type": "video/mp4",
-			},
-			Body: bytes.NewReader(byts),
-		}
-
 	case strings.HasPrefix(fname, "seg"):
 		base := strings.TrimSuffix(fname, ".mp4")