diff --git a/muxer.go b/muxer.go index 2663fff..3727ec1 100644 --- a/muxer.go +++ b/muxer.go @@ -17,6 +17,7 @@ type Track struct { spsHasWritten bool pcrHasWritten bool + mux *Muxer streamId uint tsw *TSWriter dataBuf *iovec @@ -74,7 +75,7 @@ func (self *Track) WriteH264NALU(sync bool, duration int, nalu []byte) (err erro data.Prepend(self.getPesHeader(0)) self.tsw.RandomAccessIndicator = sync self.setPCR() - if err = self.tsw.WriteIovec(data); err != nil { + if err = self.tsw.WriteIovecTo(self.mux.W, data); err != nil { return } @@ -87,7 +88,7 @@ func (self *Track) WriteADTSAACFrame(duration int, frame []byte) (err error) { self.dataBuf.Prepend(self.getPesHeader(self.dataBuf.Len)) self.tsw.RandomAccessIndicator = true self.setPCR() - if err = self.tsw.WriteIovec(self.dataBuf); err != nil { + if err = self.tsw.WriteIovecTo(self.mux.W, self.dataBuf); err != nil { return } self.dataBuf = nil @@ -100,10 +101,10 @@ func (self *Track) WriteADTSAACFrame(duration int, frame []byte) (err error) { return } -func newTrack(w io.Writer, pid uint, streamId uint) (track *Track) { +func newTrack(mux *Muxer, pid uint, streamId uint) (track *Track) { track = &Track{ + mux: mux, tsw: &TSWriter{ - W: w, PID: pid, DiscontinuityIndicator: true, }, @@ -118,6 +119,8 @@ type Muxer struct { tswPAT *TSWriter tswPMT *TSWriter elemStreams []ElementaryStreamInfo + TrackH264 *Track + Tracks []*Track } func (self *Muxer) AddAACTrack() (track *Track) { @@ -125,9 +128,10 @@ func (self *Muxer) AddAACTrack() (track *Track) { self.elemStreams, ElementaryStreamInfo{StreamType: ElementaryStreamTypeAdtsAAC, ElementaryPID: 0x101}, ) - track = newTrack(self.W, 0x101, StreamIdAAC) + track = newTrack(self, 0x101, StreamIdAAC) track.pcrHasWritten = true track.cacheSize = 3000 + self.Tracks = append(self.Tracks, track) return } @@ -136,7 +140,9 @@ func (self *Muxer) AddH264Track() (track *Track) { self.elemStreams, ElementaryStreamInfo{StreamType: ElementaryStreamTypeH264, ElementaryPID: 0x100}, ) - track = newTrack(self.W, 0x100, StreamIdH264) + track = newTrack(self, 0x100, StreamIdH264) + self.TrackH264 = track + self.Tracks = append(self.Tracks, track) return } @@ -157,21 +163,25 @@ func (self *Muxer) WriteHeader() (err error) { WritePMT(bufPMT, pmt) tswPMT := &TSWriter{ - W: self.W, PID: 0x1000, DiscontinuityIndicator: true, } tswPAT := &TSWriter{ - W: self.W, PID: 0, DiscontinuityIndicator: true, } - if err = tswPAT.Write(bufPAT.Bytes()); err != nil { + if err = tswPAT.WriteTo(self.W, bufPAT.Bytes()); err != nil { return } - if err = tswPMT.Write(bufPMT.Bytes()); err != nil { + if err = tswPMT.WriteTo(self.W, bufPMT.Bytes()); err != nil { return } + + for _, track := range(self.Tracks) { + track.spsHasWritten = false + track.pcrHasWritten = false + } + return } diff --git a/writer.go b/writer.go index 4d237f3..3e7468f 100644 --- a/writer.go +++ b/writer.go @@ -174,11 +174,22 @@ func (self *TSWriter) EnableVecWriter() { } func (self *TSWriter) WriteIovec(data *iovec) (err error) { - w := self.W if self.vecw != nil { - w = self.vecw + if err = self.WriteIovecTo(self.vecw, data); err != nil { + return + } + if err = self.vecw.Flush(); err != nil { + return + } + } else { + if err = self.WriteIovecTo(self.W, data); err != nil { + return + } } + return +} +func (self *TSWriter) WriteIovecTo(w io.Writer, data *iovec) (err error) { for i := 0; data.Len > 0; i++ { header := TSHeader{ PID: self.PID, @@ -216,16 +227,15 @@ func (self *TSWriter) WriteIovec(data *iovec) (err error) { self.ContinuityCounter++ } - - if self.vecw != nil { - if err = self.vecw.Flush(); err != nil { - return - } - } - return } +func (self *TSWriter) WriteTo(w io.Writer, data []byte) (err error) { + iov := &iovec{} + iov.Append(data) + return self.WriteIovecTo(w, iov) +} + func (self *TSWriter) Write(data []byte) (err error) { iov := &iovec{} iov.Append(data)