From 82e4efa5ebf25ebe70e9c355837d5abf070c8a01 Mon Sep 17 00:00:00 2001
From: Daniel Lawrence <lawrence.danielt@gmail.com>
Date: Tue, 22 Oct 2024 11:46:06 -0400
Subject: [PATCH] Support multi-variant HLS streams

---
 pkg/media/pipeline.go    | 21 ++++++++++++++++++---
 pkg/media/webrtc_sink.go | 16 +++++++++++++++-
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/pkg/media/pipeline.go b/pkg/media/pipeline.go
index 3119842e..bc94e0f5 100644
--- a/pkg/media/pipeline.go
+++ b/pkg/media/pipeline.go
@@ -48,8 +48,10 @@ type Pipeline struct {
 	sink     *WebRTCSink
 	input    *Input
 
-	closed core.Fuse
-	cancel atomic.Pointer[context.CancelFunc]
+	audioSinkCreated atomic.Bool
+	videoSinkCreated atomic.Bool
+	closed           core.Fuse
+	cancel           atomic.Pointer[context.CancelFunc]
 
 	pipelineErr chan error
 }
@@ -119,6 +121,19 @@ func (p *Pipeline) onOutputReady(pad *gst.Pad, kind types.StreamKind) {
 }
 
 func (p *Pipeline) onParamsReady(kind types.StreamKind, gPad *gst.GhostPad, param *glib.ParamSpec) {
+	// Keep track of whether we've already created audio/video outputs, and skip
+	// adding more if so
+	switch kind {
+	case types.Audio:
+		if !p.audioSinkCreated.CompareAndSwap(false, true) {
+			return
+		}
+	case types.Video:
+		if !p.videoSinkCreated.CompareAndSwap(false, true) {
+			return
+		}
+	}
+
 	var err error
 
 	// TODO fix go-gst to not create non nil gst.Caps for a NULL native caps pointer?
@@ -139,7 +154,7 @@ func (p *Pipeline) onParamsReady(kind types.StreamKind, gPad *gst.GhostPad, para
 		p.SendStateUpdate(context.Background())
 	}()
 
-	bin, err := p.sink.AddTrack(kind, caps.(*gst.Caps))
+	bin, err := p.sink.AddTrack(kind, caps.(*gst.Caps), p.Params)
 	if err != nil {
 		return
 	}
diff --git a/pkg/media/webrtc_sink.go b/pkg/media/webrtc_sink.go
index 30e79e4d..d14d51da 100644
--- a/pkg/media/webrtc_sink.go
+++ b/pkg/media/webrtc_sink.go
@@ -198,7 +198,7 @@ func (s *WebRTCSink) addVideoTrack(w, h int) ([]*Output, error) {
 	return outputs, nil
 }
 
-func (s *WebRTCSink) AddTrack(kind types.StreamKind, caps *gst.Caps) (*gst.Bin, error) {
+func (s *WebRTCSink) AddTrack(kind types.StreamKind, caps *gst.Caps, p *params.Params) (*gst.Bin, error) {
 	var bin *gst.Bin
 
 	switch kind {
@@ -218,6 +218,20 @@ func (s *WebRTCSink) AddTrack(kind types.StreamKind, caps *gst.Caps) (*gst.Bin,
 		}
 
 		logger.Infow("source resolution parsed", "width", w, "height", h)
+		// In cases like multi-variant HLS, the initial source resolution may be very low
+		// (e.g. 320x180). Since the input capsfilter element in the output bin will maintain this
+		// resolution, upscale to the highest layer's dimensions to prevent downscaling if we
+		// get a higher resolution variant later.
+		if len(p.VideoEncodingOptions.GetLayers()) > 0 {
+			layerHigh := p.VideoEncodingOptions.GetLayers()[0]
+			lw := int(layerHigh.GetWidth())
+			lh := int(layerHigh.GetHeight())
+			if lw > w || lh > h {
+				w = lw
+				h = lh
+				logger.Infow("max layer resolution greater than source, sizing up", "width", w, "height", h)
+			}
+		}
 
 		outputs, err := s.addVideoTrack(w, h)
 		if err != nil {