From c06f934d9e2a5f78acb18f45b3ce0a15cfb89553 Mon Sep 17 00:00:00 2001
From: nareix <hi.nareix@gmail.com>
Date: Fri, 29 Jul 2016 16:19:17 +0800
Subject: [PATCH] add NewMuxer(writeFlusher) for ts/flv

---
 format/flv/flv.go  | 15 ++++++++++++---
 format/ts/muxer.go | 15 ++++++++++++---
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/format/flv/flv.go b/format/flv/flv.go
index 681bb2d7..c9e50415 100644
--- 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
index 7b880b4f..2ddb0ff0 100644
--- 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 {