hls: correctly compute segment duration

Segment duration is now (current PTS - first PTS in segment)
Previously it was (last PTS in segment - first PTS in segment)

internal/hls/muxer_test.go

@@ -79,9 +79,9 @@ func TestMuxer(t *testing.T) {
 	re := regexp.MustCompile(`^#EXTM3U\n` +
 		`#EXT-X-VERSION:3\n` +
 		`#EXT-X-ALLOW-CACHE:NO\n` +
-		`#EXT-X-TARGETDURATION:2\n` +
+		`#EXT-X-TARGETDURATION:4\n` +
 		`#EXT-X-MEDIA-SEQUENCE:0\n` +
-		`#EXTINF:2,\n` +
+		`#EXTINF:4,\n` +
 		`([0-9]+\.ts)\n$`)
 	ma := re.FindStringSubmatch(string(byts))
 	require.NotEqual(t, 0, len(ma))

internal/hls/muxer_ts_generator.go

@@ -1,13 +1,11 @@
 package hls
 
 import (
-	"context"
 	"time"
 
 	"github.com/aler9/gortsplib"
 	"github.com/aler9/gortsplib/pkg/aac"
 	"github.com/aler9/gortsplib/pkg/h264"
-	"github.com/asticode/go-astits"
 )
 
 const (
@@ -26,7 +24,7 @@ type muxerTSGenerator struct {
 	aacConf            *gortsplib.TrackConfigAAC
 	streamPlaylist     *muxerStreamPlaylist
 
-	tsmuxer        *astits.Muxer
+	writer         *muxerTSWriter
 	currentSegment *muxerTSSegment
 	videoDTSEst    *h264.DTSEstimator
 	audioAUCount   int
@@ -48,42 +46,17 @@ func newMuxerTSGenerator(
 		hlsSegmentDuration: hlsSegmentDuration,
 		videoTrack:         videoTrack,
 		audioTrack:         audioTrack,
-		streamPlaylist:     streamPlaylist,
 		h264Conf:           h264Conf,
 		aacConf:            aacConf,
+		streamPlaylist:     streamPlaylist,
+		writer:             newMuxerTSWriter(videoTrack, audioTrack),
 	}
 
-	m.tsmuxer = astits.NewMuxer(context.Background(), m)
+	m.currentSegment = newMuxerTSSegment(m.videoTrack, m.writer)
-
-	if videoTrack != nil {
-		m.tsmuxer.AddElementaryStream(astits.PMTElementaryStream{
-			ElementaryPID: 256,
-			StreamType:    astits.StreamTypeH264Video,
-		})
-	}
-
-	if audioTrack != nil {
-		m.tsmuxer.AddElementaryStream(astits.PMTElementaryStream{
-			ElementaryPID: 257,
-			StreamType:    astits.StreamTypeAACAudio,
-		})
-	}
-
-	if videoTrack != nil {
-		m.tsmuxer.SetPCRPID(256)
-	} else {
-		m.tsmuxer.SetPCRPID(257)
-	}
-
-	m.currentSegment = newMuxerTSSegment(m.videoTrack, m.tsmuxer)
 
 	return m
 }
 
-func (m *muxerTSGenerator) Write(p []byte) (int, error) {
-	return m.currentSegment.write(p)
-}
-
 func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
 	idrPresent := func() bool {
 		for _, nalu := range nalus {
@@ -100,14 +73,7 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
 		return nil
 	}
 
-	// switch segment or initialize the first segment
+	if !m.currentSegment.firstPacketWritten {
-	if m.currentSegment.firstPacketWritten {
-		if idrPresent &&
-			m.currentSegment.duration() >= m.hlsSegmentDuration {
-			m.streamPlaylist.pushSegment(m.currentSegment)
-			m.currentSegment = newMuxerTSSegment(m.videoTrack, m.tsmuxer)
-		}
-	} else {
 		m.startPCR = time.Now()
 		m.startPTS = pts
 		m.videoDTSEst = h264.NewDTSEstimator()
@@ -116,6 +82,16 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
 	dts := m.videoDTSEst.Feed(pts-m.startPTS) + pcrOffset
 	pts = pts - m.startPTS + pcrOffset
 
+	// switch segment or initialize the first segment
+	if m.currentSegment.firstPacketWritten {
+		if idrPresent &&
+			(pts-m.currentSegment.startPTS) >= m.hlsSegmentDuration {
+			m.currentSegment.endPTS = pts
+			m.streamPlaylist.pushSegment(m.currentSegment)
+			m.currentSegment = newMuxerTSSegment(m.videoTrack, m.writer)
+		}
+	}
+
 	// prepend an AUD. This is required by video.js and iOS
 	filteredNALUs := [][]byte{
 		{byte(h264.NALUTypeAccessUnitDelimiter), 240},
@@ -147,18 +123,23 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error {
 }
 
 func (m *muxerTSGenerator) writeAAC(pts time.Duration, aus [][]byte) error {
+	if m.videoTrack == nil && !m.currentSegment.firstPacketWritten {
+		m.startPCR = time.Now()
+		m.startPTS = pts
+	}
+
+	pts = pts - m.startPTS + pcrOffset
+
 	// switch segment or initialize the first segment
 	if m.videoTrack == nil {
 		if m.currentSegment.firstPacketWritten {
 			if m.audioAUCount >= segmentMinAUCount &&
-				m.currentSegment.duration() >= m.hlsSegmentDuration {
+				(pts-m.currentSegment.startPTS) >= m.hlsSegmentDuration {
 				m.audioAUCount = 0
+				m.currentSegment.endPTS = pts
 				m.streamPlaylist.pushSegment(m.currentSegment)
-				m.currentSegment = newMuxerTSSegment(m.videoTrack, m.tsmuxer)
+				m.currentSegment = newMuxerTSSegment(m.videoTrack, m.writer)
 			}
-		} else {
-			m.startPCR = time.Now()
-			m.startPTS = pts
 		}
 	} else {
 		if !m.currentSegment.firstPacketWritten {
@@ -166,8 +147,6 @@ func (m *muxerTSGenerator) writeAAC(pts time.Duration, aus [][]byte) error {
 		}
 	}
 
-	pts = pts - m.startPTS + pcrOffset
-
 	for _, au := range aus {
 		enc, err := aac.EncodeADTS([]*aac.ADTSPacket{
 			{

internal/hls/muxer_ts_segment.go

@@ -12,23 +12,23 @@ import (
 
 type muxerTSSegment struct {
 	videoTrack *gortsplib.Track
-	tsmuxer    *astits.Muxer
+	writer     *muxerTSWriter
 
 	name               string
 	buf                bytes.Buffer
 	firstPacketWritten bool
-	minPTS             time.Duration
+	startPTS           time.Duration
-	maxPTS             time.Duration
+	endPTS             time.Duration
 	pcrSendCounter     int
 }
 
 func newMuxerTSSegment(
 	videoTrack *gortsplib.Track,
-	tsmuxer *astits.Muxer,
+	writer *muxerTSWriter,
 ) *muxerTSSegment {
 	t := &muxerTSSegment{
 		videoTrack: videoTrack,
-		tsmuxer:    tsmuxer,
+		writer:     writer,
 		name:       strconv.FormatInt(time.Now().Unix(), 10),
 	}
 
@@ -37,11 +37,13 @@ func newMuxerTSSegment(
 	// - AdaptationField != nil
 	// - RandomAccessIndicator = true
 
+	writer.currentSegment = t
+
 	return t
 }
 
 func (t *muxerTSSegment) duration() time.Duration {
-	return t.maxPTS - t.minPTS
+	return t.endPTS - t.startPTS
 }
 
 func (t *muxerTSSegment) write(p []byte) (int, error) {
@@ -60,15 +62,7 @@ func (t *muxerTSSegment) writeH264(
 	enc []byte) error {
 	if !t.firstPacketWritten {
 		t.firstPacketWritten = true
-		t.minPTS = pts
+		t.startPTS = pts
-		t.maxPTS = pts
-	} else {
-		if pts < t.minPTS {
-			t.minPTS = pts
-		}
-		if pts > t.maxPTS {
-			t.maxPTS = pts
-		}
 	}
 
 	var af *astits.PacketAdaptationField
@@ -104,7 +98,7 @@ func (t *muxerTSSegment) writeH264(
 		oh.PTS = &astits.ClockReference{Base: int64(pts.Seconds() * 90000)}
 	}
 
-	_, err := t.tsmuxer.WriteData(&astits.MuxerData{
+	_, err := t.writer.WriteData(&astits.MuxerData{
 		PID:             256,
 		AdaptationField: af,
 		PES: &astits.PESData{
@@ -125,15 +119,7 @@ func (t *muxerTSSegment) writeAAC(
 	if t.videoTrack == nil {
 		if !t.firstPacketWritten {
 			t.firstPacketWritten = true
-			t.minPTS = pts
+			t.startPTS = pts
-			t.maxPTS = pts
-		} else {
-			if pts < t.minPTS {
-				t.minPTS = pts
-			}
-			if pts > t.maxPTS {
-				t.maxPTS = pts
-			}
 		}
 	}
 
@@ -151,7 +137,7 @@ func (t *muxerTSSegment) writeAAC(
 		}
 	}
 
-	_, err := t.tsmuxer.WriteData(&astits.MuxerData{
+	_, err := t.writer.WriteData(&astits.MuxerData{
 		PID:             257,
 		AdaptationField: af,
 		PES: &astits.PESData{

internal/hls/muxer_ts_writer.go

@@ -0,0 +1,51 @@
+package hls
+
+import (
+	"context"
+
+	"github.com/aler9/gortsplib"
+	"github.com/asticode/go-astits"
+)
+
+type muxerTSWriter struct {
+	innerMuxer     *astits.Muxer
+	currentSegment *muxerTSSegment
+}
+
+func newMuxerTSWriter(
+	videoTrack *gortsplib.Track,
+	audioTrack *gortsplib.Track) *muxerTSWriter {
+	w := &muxerTSWriter{}
+
+	w.innerMuxer = astits.NewMuxer(context.Background(), w)
+
+	if videoTrack != nil {
+		w.innerMuxer.AddElementaryStream(astits.PMTElementaryStream{
+			ElementaryPID: 256,
+			StreamType:    astits.StreamTypeH264Video,
+		})
+	}
+
+	if audioTrack != nil {
+		w.innerMuxer.AddElementaryStream(astits.PMTElementaryStream{
+			ElementaryPID: 257,
+			StreamType:    astits.StreamTypeAACAudio,
+		})
+	}
+
+	if videoTrack != nil {
+		w.innerMuxer.SetPCRPID(256)
+	} else {
+		w.innerMuxer.SetPCRPID(257)
+	}
+
+	return w
+}
+
+func (mt *muxerTSWriter) Write(p []byte) (int, error) {
+	return mt.currentSegment.write(p)
+}
+
+func (mt *muxerTSWriter) WriteData(d *astits.MuxerData) (int, error) {
+	return mt.innerMuxer.WriteData(d)
+}