hls, rtmp: set DTS = PTS when a IDR frame is received

internal/core/rtmp_conn.go

@@ -329,13 +329,26 @@ func (c *rtmpConn) runRead(ctx context.Context) error {
 					return err
 				}
 
-				dts := videoDTSEst.Feed(pts + rtmpConnPTSOffset)
+				idrPresent := func() bool {
+					for _, nalu := range nalus {
+						typ := h264.NALUType(nalu[0] & 0x1F)
+						if typ == h264.NALUTypeIDR {
+							return true
+						}
+					}
+					return false
+				}()
+
+				pts += rtmpConnPTSOffset
+
+				dts := videoDTSEst.Feed(idrPresent, pts)
+
 				c.conn.NetConn().SetWriteDeadline(time.Now().Add(c.writeTimeout))
 				err = c.conn.WritePacket(av.Packet{
 					Type:  av.H264,
 					Data:  data,
 					Time:  dts,
-					CTime: pts + rtmpConnPTSOffset - dts,
+					CTime: pts - dts,
 				})
 				if err != nil {
 					return err

internal/h264/dtsestimator.go

@@ -20,7 +20,7 @@ func NewDTSEstimator() *DTSEstimator {
 }
 
 // Feed provides PTS to the estimator, and returns the estimated DTS.
-func (d *DTSEstimator) Feed(pts time.Duration) time.Duration {
+func (d *DTSEstimator) Feed(idrPresent bool, pts time.Duration) time.Duration {
 	if d.initializing > 0 {
 		d.initializing--
 		dts := d.prevDTS + time.Millisecond
@@ -31,6 +31,12 @@ func (d *DTSEstimator) Feed(pts time.Duration) time.Duration {
 	}
 
 	dts := func() time.Duration {
+		// IDR
+		if idrPresent {
+			// DTS is always PTS
+			return pts
+		}
+
 		// P or I frame
 		if pts > d.prevPTS {
 			// previous frame was B

internal/h264/dtsestimator_test.go

@@ -9,9 +9,19 @@ import (
 
 func TestDTSEstimator(t *testing.T) {
 	est := NewDTSEstimator()
-	est.Feed(2 * time.Second)
-	est.Feed(2*time.Second - 200*time.Millisecond)
-	est.Feed(2*time.Second - 400*time.Millisecond)
-	dts := est.Feed(2*time.Second + 200*time.Millisecond)
+
+	dts := est.Feed(false, 2*time.Second)
+	require.Equal(t, time.Millisecond, dts)
+
+	dts = est.Feed(false, 2*time.Second-200*time.Millisecond)
+	require.Equal(t, 2*time.Millisecond, dts)
+
+	dts = est.Feed(false, 2*time.Second-400*time.Millisecond)
+	require.Equal(t, 3*time.Millisecond, dts)
+
+	dts = est.Feed(false, 2*time.Second+200*time.Millisecond)
 	require.Equal(t, 2*time.Second-400*time.Millisecond, dts)
+
+	dts = est.Feed(true, 2*time.Second+300*time.Millisecond)
+	require.Equal(t, 2*time.Second+300*time.Millisecond, dts)
 }

internal/hls/muxer.go

@@ -115,7 +115,7 @@ func (m *Muxer) WriteH264(pts time.Duration, nalus [][]byte) error {
 	pts = pts + ptsOffset - m.startPTS
 
 	err := m.currentSegment.writeH264(
-		m.videoDTSEst.Feed(pts),
+		m.videoDTSEst.Feed(idrPresent, pts),
 		pts,
 		idrPresent,
 		nalus)