On layer switch don't send packets until Keyframe

A better implementation would send old video until new is available, this keeps it simple
Glimesh · Sep 18, 2024 · 98abef0 · 98abef0
1 parent f2fffcd
commit 98abef0
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 9 deletions.
diff --git a/internal/webrtc/keyframe_detector.go b/internal/webrtc/keyframe_detector.go
@@ -22,5 +22,5 @@ func isKeyframe(pkt *rtp.Packet, codec videoTrackCodec, depacketizer rtp.Depacke
 		firstNaluType := nalu[4] & naluTypeBitmask
 		return firstNaluType == idrNALUType || firstNaluType == spsNALUType || firstNaluType == ppsNALUType
 	}
-	return false
+	return true
 }
diff --git a/internal/webrtc/whep.go b/internal/webrtc/whep.go
@@ -15,11 +15,12 @@ import (
 
 type (
 	whepSession struct {
-		videoTrack     *trackMultiCodec
-		currentLayer   atomic.Value
-		sequenceNumber uint16
-		timestamp      uint32
-		packetsWritten uint64
+		videoTrack         *trackMultiCodec
+		currentLayer       atomic.Value
+		waitingForKeyframe atomic.Bool
+		sequenceNumber     uint16
+		timestamp          uint32
+		packetsWritten     uint64
 	}
 
 	simulcastLayerResponse struct {
@@ -64,6 +65,7 @@ func WHEPChangeLayer(whepSessionId, layer string) error {
 
 		if _, ok := streamMap[streamKey].whepSessions[whepSessionId]; ok {
 			streamMap[streamKey].whepSessions[whepSessionId].currentLayer.Store(layer)
+			streamMap[streamKey].whepSessions[whepSessionId].waitingForKeyframe.Store(true)
 			streamMap[streamKey].pliChan <- true
 		}
 	}
@@ -151,14 +153,22 @@ func WHEP(offer, streamKey string) (string, string, error) {
 		timestamp:  50000,
 	}
 	stream.whepSessions[whepSessionId].currentLayer.Store("")
+	stream.whepSessions[whepSessionId].waitingForKeyframe.Store(false)
+
 	return appendOffer(peerConnection.LocalDescription().SDP), whepSessionId, nil
 }
 
-func (w *whepSession) sendVideoPacket(rtpPkt *rtp.Packet, layer string, timeDiff int64, sequenceDiff int, codec videoTrackCodec) {
+func (w *whepSession) sendVideoPacket(rtpPkt *rtp.Packet, layer string, timeDiff int64, sequenceDiff int, codec videoTrackCodec, isKeyframe bool) {
 	if w.currentLayer.Load() == "" {
 		w.currentLayer.Store(layer)
 	} else if layer != w.currentLayer.Load() {
 		return
+	} else if w.waitingForKeyframe.Load() {
+		if !isKeyframe {
+			return
+		}
+
+		w.waitingForKeyframe.Store(false)
 	}
 
 	w.packetsWritten += 1

diff --git a/internal/webrtc/whip.go b/internal/webrtc/whip.go
@@ -100,7 +100,9 @@ func videoWriter(remoteTrack *webrtc.TrackRemote, stream *stream, peerConnection
 
 		videoTrack.packetsReceived.Add(1)
 
-		if isKeyframe(rtpPkt, codec, depacketizer) {
+		// Keyframe detection has only been implemented for H264
+		isKeyframe := isKeyframe(rtpPkt, codec, depacketizer)
+		if isKeyframe && codec == videoTrackCodecH264 {
 			videoTrack.lastKeyFrameSeen.Store(time.Now())
 		}
 
@@ -130,7 +132,7 @@ func videoWriter(remoteTrack *webrtc.TrackRemote, stream *stream, peerConnection
 
 		s.whepSessionsLock.RLock()
 		for i := range s.whepSessions {
-			s.whepSessions[i].sendVideoPacket(rtpPkt, id, timeDiff, sequenceDiff, codec)
+			s.whepSessions[i].sendVideoPacket(rtpPkt, id, timeDiff, sequenceDiff, codec, isKeyframe)
 		}
 		s.whepSessionsLock.RUnlock()