add NewMuxer(writeFlusher) for ts/flv

peckjerry · Jul 29, 2016 · c06f934 · c06f934
1 parent 9e12fc5
commit c06f934
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/format/flv/flv.go b/format/flv/flv.go
@@ -312,18 +312,27 @@ func PacketToTag(pkt av.Packet, stream av.CodecData) (_tag flvio.Tag, timestamp
 }
 
 type Muxer struct {
-	bufw *bufio.Writer
+	bufw writeFlusher
 	b []byte
 	streams []av.CodecData
 }
 
-func NewMuxer(w io.Writer) *Muxer {
+type writeFlusher interface {
+	io.Writer
+	Flush() error
+}
+
+func NewMuxerWriteFlusher(w writeFlusher) *Muxer {
 	return &Muxer{
-		bufw: bufio.NewWriterSize(w, pio.RecommendBufioSize),
+		bufw: w,
 		b: make([]byte, 256),
 	}
 }
 
+func NewMuxer(w io.Writer) *Muxer {
+	return NewMuxerWriteFlusher(bufio.NewWriterSize(w, pio.RecommendBufioSize))
+}
+
 var CodecTypes = []av.CodecType{av.H264, av.AAC, av.SPEEX}
 
 func (self *Muxer) WriteHeader(streams []av.CodecData) (err error) {

diff --git a/format/ts/muxer.go b/format/ts/muxer.go
@@ -15,7 +15,7 @@ import (
 var CodecTypes = []av.CodecType{av.H264, av.AAC}
 
 type Muxer struct {
-	w                        *bufio.Writer
+	w                        writeFlusher
 	streams                  []*Stream
 	PaddingToMakeCounterCont bool
 
@@ -29,9 +29,14 @@ type Muxer struct {
 	tswpat, tswpmt *tsio.TSWriter
 }
 
-func NewMuxer(w io.Writer) *Muxer {
+type writeFlusher interface {
+	io.Writer
+	Flush() error
+}
+
+func NewMuxerWriteFlusher(w writeFlusher) *Muxer {
 	return &Muxer{
-		w: bufio.NewWriterSize(w, pio.RecommendBufioSize),
+		w: w,
 		psidata: make([]byte, 188),
 		peshdr: make([]byte, tsio.MaxPESHeaderLength),
 		tshdr: make([]byte, tsio.MaxTSHeaderLength),
@@ -43,6 +48,10 @@ func NewMuxer(w io.Writer) *Muxer {
 	}
 }
 
+func NewMuxer(w io.Writer) *Muxer {
+	return NewMuxerWriteFlusher(bufio.NewWriterSize(w, pio.RecommendBufioSize))
+}
+
 func (self *Muxer) newStream(codec av.CodecData) (err error) {
 	ok := false
 	for _, c := range CodecTypes {