add TSMuxer.WriteTo func

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 {
 		return
 	}
+	if err = tswPMT.WriteTo(self.W, bufPMT.Bytes()); err != nil {
+		return
+	}
+
+	for _, track := range(self.Tracks) {
+		track.spsHasWritten = false
+		track.pcrHasWritten = false
+	}
+
 	return
 }
 

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,14 +227,13 @@ 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) {