From ce42c53a03a7d130e3ec0ca658fa7a49db6987ed Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sun, 10 Apr 2022 17:38:58 +0200
Subject: [PATCH] hls, rtmp: fix video/audio sync

---
 internal/core/hls_muxer.go | 13 ++++++++++++-
 internal/core/rtmp_conn.go | 20 +++++++++++++++-----
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/internal/core/hls_muxer.go b/internal/core/hls_muxer.go
index 4a6c5ead..1d774a41 100644
--- a/internal/core/hls_muxer.go
+++ b/internal/core/hls_muxer.go
@@ -327,6 +327,8 @@ func (m *hlsMuxer) runInner(innerCtx context.Context, innerReady chan struct{})
 	writerDone := make(chan error)
 	go func() {
 		writerDone <- func() error {
+			var videoInitialPTS *time.Duration
+
 			for {
 				item, ok := m.ringBuffer.Pull()
 				if !ok {
@@ -339,7 +341,16 @@ func (m *hlsMuxer) runInner(innerCtx context.Context, innerReady chan struct{})
 						continue
 					}
 
-					err = m.muxer.WriteH264(data.h264PTS, data.h264NALUs)
+					// video is decoded in another routine,
+					// while audio is decoded in this routine:
+					// we have to sync their PTS.
+					if videoInitialPTS == nil {
+						v := data.h264PTS
+						videoInitialPTS = &v
+					}
+					pts := data.h264PTS - *videoInitialPTS
+
+					err = m.muxer.WriteH264(pts, data.h264NALUs)
 					if err != nil {
 						m.log(logger.Warn, "unable to write segment: %v", err)
 						continue
diff --git a/internal/core/rtmp_conn.go b/internal/core/rtmp_conn.go
index 29efc975..ce290e59 100644
--- a/internal/core/rtmp_conn.go
+++ b/internal/core/rtmp_conn.go
@@ -321,9 +321,10 @@ func (c *rtmpConn) runRead(ctx context.Context) error {
 	// disable read deadline
 	c.conn.SetReadDeadline(time.Time{})
 
-	var videoStartPTS time.Duration
-	var videoDTSEst *h264.DTSEstimator
+	var videoInitialPTS *time.Duration
 	videoFirstIDRFound := false
+	var videoFirstIDRPTS time.Duration
+	var videoDTSEst *h264.DTSEstimator
 
 	for {
 		item, ok := c.ringBuffer.Pull()
@@ -337,6 +338,15 @@ func (c *rtmpConn) runRead(ctx context.Context) error {
 				continue
 			}
 
+			// video is decoded in another routine,
+			// while audio is decoded in this routine:
+			// we have to sync their PTS.
+			if videoInitialPTS == nil {
+				v := data.h264PTS
+				videoInitialPTS = &v
+			}
+			pts := data.h264PTS - *videoInitialPTS
+
 			// wait until we receive an IDR
 			if !videoFirstIDRFound {
 				if !h264.IDRPresent(data.h264NALUs) {
@@ -344,7 +354,7 @@ func (c *rtmpConn) runRead(ctx context.Context) error {
 				}
 
 				videoFirstIDRFound = true
-				videoStartPTS = data.h264PTS
+				videoFirstIDRPTS = pts
 				videoDTSEst = h264.NewDTSEstimator()
 			}
 
@@ -376,7 +386,7 @@ func (c *rtmpConn) runRead(ctx context.Context) error {
 				return err
 			}
 
-			pts := data.h264PTS - videoStartPTS
+			pts -= videoFirstIDRPTS
 			dts := videoDTSEst.Feed(pts)
 
 			c.conn.SetWriteDeadline(time.Now().Add(time.Duration(c.writeTimeout)))
@@ -402,7 +412,7 @@ func (c *rtmpConn) runRead(ctx context.Context) error {
 				continue
 			}
 
-			pts -= videoStartPTS
+			pts -= videoFirstIDRPTS
 			if pts < 0 {
 				continue
 			}