hls muxer: simplify code

This commit is contained in:
aler9 2022-03-20 17:58:39 +01:00
parent 6385fcdea5
commit bdf80a0dd3
3 changed files with 44 additions and 64 deletions

View file

@ -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 {

View file

@ -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{

View file

@ -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)
}