From bdf80a0dd360fa7ec2adb3177d28f83f6f43912a Mon Sep 17 00:00:00 2001 From: aler9 <46489434+aler9@users.noreply.github.com> Date: Sun, 20 Mar 2022 17:58:39 +0100 Subject: [PATCH] hls muxer: simplify code --- internal/hls/muxer_ts_generator.go | 45 ++++++++++++++++++++++---- internal/hls/muxer_ts_segment.go | 12 +++---- internal/hls/muxer_ts_writer.go | 51 ------------------------------ 3 files changed, 44 insertions(+), 64 deletions(-) delete mode 100644 internal/hls/muxer_ts_writer.go diff --git a/internal/hls/muxer_ts_generator.go b/internal/hls/muxer_ts_generator.go index 7ac2c871..1f8e2a13 100644 --- a/internal/hls/muxer_ts_generator.go +++ b/internal/hls/muxer_ts_generator.go @@ -1,11 +1,13 @@ 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 ( @@ -25,6 +27,12 @@ func idrPresent(nalus [][]byte) bool { return false } +type writerFunc func(p []byte) (int, error) + +func (f writerFunc) Write(p []byte) (int, error) { + return f(p) +} + type muxerTSGenerator struct { hlsSegmentCount int hlsSegmentDuration time.Duration @@ -33,7 +41,7 @@ type muxerTSGenerator struct { audioTrack *gortsplib.TrackAAC streamPlaylist *muxerStreamPlaylist - writer *muxerTSWriter + writer *astits.Muxer currentSegment *muxerTSSegment videoDTSEst *h264.DTSEstimator startPCR time.Time @@ -55,7 +63,32 @@ func newMuxerTSGenerator( videoTrack: videoTrack, audioTrack: audioTrack, streamPlaylist: streamPlaylist, - writer: newMuxerTSWriter(videoTrack, audioTrack), + } + + m.writer = astits.NewMuxer( + context.Background(), + writerFunc(func(p []byte) (int, error) { + return m.currentSegment.write(p) + })) + + if videoTrack != nil { + m.writer.AddElementaryStream(astits.PMTElementaryStream{ + ElementaryPID: 256, + StreamType: astits.StreamTypeH264Video, + }) + } + + if audioTrack != nil { + m.writer.AddElementaryStream(astits.PMTElementaryStream{ + ElementaryPID: 257, + StreamType: astits.StreamTypeAACAudio, + }) + } + + if videoTrack != nil { + m.writer.SetPCRPID(256) + } else { + m.writer.SetPCRPID(257) } return m @@ -71,7 +104,7 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error { } // create first segment - m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer) + m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer.WriteData) m.startPCR = time.Now() m.startPTS = pts m.videoDTSEst = h264.NewDTSEstimator() @@ -85,7 +118,7 @@ func (m *muxerTSGenerator) writeH264(pts time.Duration, nalus [][]byte) error { (pts-*m.currentSegment.startPTS) >= m.hlsSegmentDuration { m.currentSegment.endPTS = pts m.streamPlaylist.pushSegment(m.currentSegment) - m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer) + m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer.WriteData) } } @@ -136,7 +169,7 @@ func (m *muxerTSGenerator) writeAAC(pts time.Duration, aus [][]byte) error { if m.videoTrack == nil { if m.currentSegment == nil { // create first segment - m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer) + m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer.WriteData) m.startPCR = time.Now() m.startPTS = pts pts = pcrOffset @@ -149,7 +182,7 @@ func (m *muxerTSGenerator) writeAAC(pts time.Duration, aus [][]byte) error { (pts-*m.currentSegment.startPTS) >= m.hlsSegmentDuration { m.currentSegment.endPTS = pts m.streamPlaylist.pushSegment(m.currentSegment) - m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer) + m.currentSegment = newMuxerTSSegment(m.hlsSegmentMaxSize, m.videoTrack, m.writer.WriteData) } } } else { diff --git a/internal/hls/muxer_ts_segment.go b/internal/hls/muxer_ts_segment.go index 6234d391..a3e1e41c 100644 --- a/internal/hls/muxer_ts_segment.go +++ b/internal/hls/muxer_ts_segment.go @@ -14,7 +14,7 @@ import ( type muxerTSSegment struct { hlsSegmentMaxSize uint64 videoTrack *gortsplib.TrackH264 - writer *muxerTSWriter + writeData func(*astits.MuxerData) (int, error) startTime time.Time name string @@ -28,14 +28,14 @@ type muxerTSSegment struct { func newMuxerTSSegment( hlsSegmentMaxSize uint64, videoTrack *gortsplib.TrackH264, - writer *muxerTSWriter, + writeData func(*astits.MuxerData) (int, error), ) *muxerTSSegment { now := time.Now() t := &muxerTSSegment{ hlsSegmentMaxSize: hlsSegmentMaxSize, videoTrack: videoTrack, - writer: writer, + writeData: writeData, startTime: now, name: strconv.FormatInt(now.Unix(), 10), } @@ -45,8 +45,6 @@ func newMuxerTSSegment( // - AdaptationField != nil // - RandomAccessIndicator = true - writer.currentSegment = t - return t } @@ -103,7 +101,7 @@ func (t *muxerTSSegment) writeH264( oh.PTS = &astits.ClockReference{Base: int64(pts.Seconds() * 90000)} } - _, err := t.writer.WriteData(&astits.MuxerData{ + _, err := t.writeData(&astits.MuxerData{ PID: 256, AdaptationField: af, PES: &astits.PESData{ @@ -147,7 +145,7 @@ func (t *muxerTSSegment) writeAAC( } } - _, err := t.writer.WriteData(&astits.MuxerData{ + _, err := t.writeData(&astits.MuxerData{ PID: 257, AdaptationField: af, PES: &astits.PESData{ diff --git a/internal/hls/muxer_ts_writer.go b/internal/hls/muxer_ts_writer.go deleted file mode 100644 index 23a4d571..00000000 --- a/internal/hls/muxer_ts_writer.go +++ /dev/null @@ -1,51 +0,0 @@ -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.TrackH264, - audioTrack *gortsplib.TrackAAC) *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) -}