Skip to content

Commit

Permalink
muxer: support multiple audio tracks (bluenviron/mediamtx#2728)
Browse files Browse the repository at this point in the history
  • Loading branch information
@aler9
aler9 committed Sep 22, 2024
1 parent d9c229d commit 2c50b2b
Show file tree
Hide file tree
Showing 4 changed files with 184 additions and 139 deletions.
146 changes: 90 additions & 56 deletions muxer.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,14 @@ func (w *switchableWriter) Write(p []byte) (int, error) {
return w.w.Write(p)
}

func isVideo(codec codecs.Codec) bool {
switch codec.(type) {
case *codecs.AV1, *codecs.VP9, *codecs.H265, *codecs.H264:
return true
}
return false
}

// a prefix is needed to prevent usage of cached segments
// from previous muxing sessions.
func generatePrefix() (string, error) {
Expand Down Expand Up @@ -117,12 +125,10 @@ type fmp4AugmentedSample struct {
// Muxer is a HLS muxer.
type Muxer struct {
//
// parameters (all optional except VideoTrack or AudioTrack).
// parameters (all optional except Tracks).
//
// video track.
VideoTrack *Track
// audio track.
AudioTrack *Track
// tracks.
Tracks []*Track
// Variant to use.
// It defaults to MuxerVariantLowLatency
Variant MuxerVariant
Expand Down Expand Up @@ -189,21 +195,43 @@ func (m *Muxer) Start() error {
m.SegmentMaxSize = 50 * 1024 * 1024
}

if m.VideoTrack == nil && m.AudioTrack == nil {
return fmt.Errorf("one between VideoTrack and AudioTrack is required")
if len(m.Tracks) == 0 {
return fmt.Errorf("at least one track must be provided")
}

hasVideo := false

if m.Variant == MuxerVariantMPEGTS {
if m.VideoTrack != nil {
if _, ok := m.VideoTrack.Codec.(*codecs.H264); !ok {
return fmt.Errorf(
"the MPEG-TS variant of HLS only supports H264 video. Use the fMP4 or Low-Latency variants instead")
hasAudio := false

for _, track := range m.Tracks {
if isVideo(track.Codec) {
if hasVideo {
return fmt.Errorf("the MPEG-TS variant of HLS supports only one video track")
}
if _, ok := track.Codec.(*codecs.H264); !ok {
return fmt.Errorf(
"the MPEG-TS variant of HLS only supports H264 video. Use the fMP4 or Low-Latency variants instead")
}
hasVideo = true
} else {
if hasAudio {
return fmt.Errorf("the MPEG-TS variant of HLS supports only one audio track")
}
if _, ok := track.Codec.(*codecs.MPEG4Audio); !ok {
return fmt.Errorf(
"the MPEG-TS variant of HLS only supports MPEG-4 Audio. Use the fMP4 or Low-Latency variants instead")
}
hasAudio = true
}
}
if m.AudioTrack != nil {
if _, ok := m.AudioTrack.Codec.(*codecs.MPEG4Audio); !ok {
return fmt.Errorf(
"the MPEG-TS variant of HLS only supports MPEG-4 Audio. Use the fMP4 or Low-Latency variants instead")
} else {
for _, track := range m.Tracks {
if isVideo(track.Codec) {
if hasVideo {
return fmt.Errorf("only one video track is currently supported")
}
hasVideo = true
}
}
}
Expand Down Expand Up @@ -233,26 +261,15 @@ func (m *Muxer) Start() error {
}
m.server.initialize()

if m.VideoTrack != nil {
track := &muxerTrack{
Track: m.VideoTrack,
variant: m.Variant,
isLeading: true,
}
track.initialize()
m.mtracks = append(m.mtracks, track)
m.mtracksByTrack[m.VideoTrack] = track
}

if m.AudioTrack != nil {
track := &muxerTrack{
Track: m.AudioTrack,
for i, track := range m.Tracks {
mtrack := &muxerTrack{
Track: track,
variant: m.Variant,
isLeading: m.VideoTrack == nil,
isLeading: isVideo(track.Codec) || (!hasVideo && i == 0),
}
track.initialize()
m.mtracks = append(m.mtracks, track)
m.mtracksByTrack[m.AudioTrack] = track
mtrack.initialize()
m.mtracks = append(m.mtracks, mtrack)
m.mtracksByTrack[track] = mtrack
}

if m.Variant == MuxerVariantMPEGTS {
Expand All @@ -275,25 +292,32 @@ func (m *Muxer) Start() error {
m.streams = append(m.streams, stream)

default:
if m.VideoTrack != nil {
videoStream := &muxerStream{
muxer: m,
tracks: []*muxerTrack{m.mtracksByTrack[m.VideoTrack]},
id: "video",
defaultRenditionPicked := false

for i, track := range m.mtracks {
var id string
if isVideo(track.Codec) {
id = "video" + strconv.FormatInt(int64(i+1), 10)
} else {
id = "audio" + strconv.FormatInt(int64(i+1), 10)
}

isRendition := !track.isLeading

defaultRendition := isRendition && !defaultRenditionPicked
if defaultRendition {
defaultRenditionPicked = true
}
videoStream.initialize()
m.streams = append(m.streams, videoStream)
}

if m.AudioTrack != nil {
audioStream := &muxerStream{
muxer: m,
tracks: []*muxerTrack{m.mtracksByTrack[m.AudioTrack]},
id: "audio",
isRendition: m.VideoTrack != nil,
stream := &muxerStream{
muxer: m,
tracks: []*muxerTrack{track},
id: id,
isRendition: isRendition,
defaultRendition: defaultRendition,
}
audioStream.initialize()
m.streams = append(m.streams, audioStream)
stream.initialize()
m.streams = append(m.streams, stream)
}
}

Expand Down Expand Up @@ -327,52 +351,62 @@ func (m *Muxer) Close() {

// WriteAV1 writes an AV1 temporal unit.
func (m *Muxer) WriteAV1(
track *Track,
ntp time.Time,
pts time.Duration,
tu [][]byte,
) error {
return m.segmenter.writeAV1(ntp, pts, tu)
return m.segmenter.writeAV1(m.mtracksByTrack[track], ntp, pts, tu)
}

// WriteVP9 writes a VP9 frame.
func (m *Muxer) WriteVP9(
track *Track,
ntp time.Time,
pts time.Duration,
frame []byte,
) error {
return m.segmenter.writeVP9(ntp, pts, frame)
return m.segmenter.writeVP9(m.mtracksByTrack[track], ntp, pts, frame)
}

// WriteH265 writes an H265 access unit.
func (m *Muxer) WriteH265(
track *Track,
ntp time.Time,
pts time.Duration,
au [][]byte,
) error {
return m.segmenter.writeH265(ntp, pts, au)
return m.segmenter.writeH265(m.mtracksByTrack[track], ntp, pts, au)
}

// WriteH264 writes an H264 access unit.
func (m *Muxer) WriteH264(
track *Track,
ntp time.Time,
pts time.Duration,
au [][]byte,
) error {
return m.segmenter.writeH264(ntp, pts, au)
return m.segmenter.writeH264(m.mtracksByTrack[track], ntp, pts, au)
}

// WriteOpus writes Opus packets.
func (m *Muxer) WriteOpus(
track *Track,
ntp time.Time,
pts time.Duration,
packets [][]byte,
) error {
return m.segmenter.writeOpus(ntp, pts, packets)
return m.segmenter.writeOpus(m.mtracksByTrack[track], ntp, pts, packets)
}

// WriteMPEG4Audio writes MPEG-4 Audio access units.
func (m *Muxer) WriteMPEG4Audio(ntp time.Time, pts time.Duration, aus [][]byte) error {
return m.segmenter.writeMPEG4Audio(ntp, pts, aus)
func (m *Muxer) WriteMPEG4Audio(
track *Track,
ntp time.Time,
pts time.Duration,
aus [][]byte,
) error {
return m.segmenter.writeMPEG4Audio(m.mtracksByTrack[track], ntp, pts, aus)
}

// Handle handles a HTTP request.
Expand Down
32 changes: 15 additions & 17 deletions muxer_segmenter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,12 +31,11 @@ func (s *muxerSegmenter) initialize() {
}

func (s *muxerSegmenter) writeAV1(
track *muxerTrack,
ntp time.Time,
pts time.Duration,
tu [][]byte,
) error {
track := s.muxer.mtracksByTrack[s.muxer.VideoTrack]

codec := track.Codec.(*codecs.AV1)
randomAccess := false

Expand Down Expand Up @@ -86,12 +85,11 @@ func (s *muxerSegmenter) writeAV1(
}

func (s *muxerSegmenter) writeVP9(
track *muxerTrack,
ntp time.Time,
pts time.Duration,
frame []byte,
) error {
track := s.muxer.mtracksByTrack[s.muxer.VideoTrack]

var h vp9.Header
err := h.Unmarshal(frame)
if err != nil {
Expand Down Expand Up @@ -163,12 +161,11 @@ func (s *muxerSegmenter) writeVP9(
}

func (s *muxerSegmenter) writeH265(
track *muxerTrack,
ntp time.Time,
pts time.Duration,
au [][]byte,
) error {
track := s.muxer.mtracksByTrack[s.muxer.VideoTrack]

randomAccess := false
codec := track.Codec.(*codecs.H265)

Expand Down Expand Up @@ -245,12 +242,11 @@ func (s *muxerSegmenter) writeH265(
}

func (s *muxerSegmenter) writeH264(
track *muxerTrack,
ntp time.Time,
pts time.Duration,
au [][]byte,
) error {
track := s.muxer.mtracksByTrack[s.muxer.VideoTrack]

randomAccess := false
codec := track.Codec.(*codecs.H264)
nonIDRPresent := false
Expand Down Expand Up @@ -356,12 +352,11 @@ func (s *muxerSegmenter) writeH264(
}

func (s *muxerSegmenter) writeOpus(
track *muxerTrack,
ntp time.Time,
pts time.Duration,
packets [][]byte,
) error {
track := s.muxer.mtracksByTrack[s.muxer.AudioTrack]

if s.muxer.Variant == MuxerVariantMPEGTS {
return fmt.Errorf("unimplemented")
} else {
Expand Down Expand Up @@ -389,11 +384,14 @@ func (s *muxerSegmenter) writeOpus(
}
}

func (s *muxerSegmenter) writeMPEG4Audio(ntp time.Time, pts time.Duration, aus [][]byte) error {
track := s.muxer.mtracksByTrack[s.muxer.AudioTrack]

func (s *muxerSegmenter) writeMPEG4Audio(
track *muxerTrack,
ntp time.Time,
pts time.Duration,
aus [][]byte,
) error {
if s.muxer.Variant == MuxerVariantMPEGTS {
if s.muxer.VideoTrack == nil {
if track.isLeading {
if track.stream.nextSegment == nil {
err := s.muxer.createFirstSegment(pts, ntp)
if err != nil {
Expand All @@ -420,7 +418,7 @@ func (s *muxerSegmenter) writeMPEG4Audio(ntp time.Time, pts time.Duration, aus [

return nil
} else {
sampleRate := time.Duration(s.muxer.AudioTrack.Codec.(*codecs.MPEG4Audio).Config.SampleRate)
sampleRate := time.Duration(track.Codec.(*codecs.MPEG4Audio).Config.SampleRate)

for i, au := range aus {
auNTP := ntp.Add(time.Duration(i) * mpeg4audio.SamplesPerAccessUnit *
Expand Down Expand Up @@ -551,7 +549,7 @@ func (s *muxerSegmenter) fmp4WriteAudio(
duration := track.fmp4NextSample.dts - sample.dts
sample.Duration = uint32(durationGoToMp4(duration, track.fmp4TimeScale))

if s.muxer.VideoTrack == nil {
if track.isLeading {
// create first segment
if track.stream.nextSegment == nil {
err := s.muxer.createFirstSegment(sample.dts, sample.ntp)
Expand All @@ -576,7 +574,7 @@ func (s *muxerSegmenter) fmp4WriteAudio(
}

// switch segment
if s.muxer.VideoTrack == nil &&
if track.isLeading &&
(track.fmp4NextSample.dts-track.stream.nextSegment.(*muxerSegmentFMP4).startDTS) >= s.muxer.SegmentMinDuration {
err = s.muxer.rotateSegments(track.fmp4NextSample.dts, track.fmp4NextSample.ntp, false)
if err != nil {
Expand Down
Loading

0 comments on commit 2c50b2b

Please sign in to comment.