Rewrite once buffer for keyframes

internal/mp4/mp4.go

@@ -1,7 +1,6 @@
 package mp4
 
 import (
-	"io"
 	"net/http"
 	"strconv"
 	"strings"
@@ -56,34 +55,25 @@ func handlerKeyframe(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	wr := &once{} // init and first frame
-	_, _ = cons.WriteTo(wr)
+	once := &core.OnceBuffer{} // init and first frame
+	_, _ = cons.WriteTo(once)
 
 	stream.RemoveConsumer(cons)
 
 	// Apple Safari won't show frame without length
 	header := w.Header()
-	header.Set("Content-Length", strconv.Itoa(len(wr.buf)))
+	header.Set("Content-Length", strconv.Itoa(once.Len()))
 	header.Set("Content-Type", mp4.ContentType(cons.Codecs()))
 
 	if filename := query.Get("filename"); filename != "" {
 		header.Set("Content-Disposition", `attachment; filename="`+filename+`"`)
 	}
 
-	if _, err := w.Write(wr.buf); err != nil {
+	if _, err := once.WriteTo(w); err != nil {
 		log.Error().Err(err).Caller().Send()
 	}
 }
 
-type once struct {
-	buf []byte
-}
-
-func (o *once) Write(p []byte) (n int, err error) {
-	o.buf = p
-	return 0, io.EOF
-}
-
 func handlerMP4(w http.ResponseWriter, r *http.Request) {
 	log.Trace().Msgf("[mp4] %s %+v", r.Method, r.Header)
 

pkg/core/writebuffer.go

@@ -6,6 +6,10 @@ import (
 	"sync"
 )
 
+// WriteBuffer by defaul Write(s) to bytes.Buffer.
+// But after WriteTo to new io.Writer - calls Reset.
+// Reset will flush current buffer data to new writer and starts to Write to new io.Writer
+// WriteTo will be locked until Write fails or Close will be called.
 type WriteBuffer struct {
 	io.Writer
 	err   error
@@ -52,7 +56,7 @@ func (w *WriteBuffer) Close() error {
 func (w *WriteBuffer) Reset(wr io.Writer) {
 	w.mu.Lock()
 	w.add()
-	if buf, ok := wr.(*bytes.Buffer); ok {
+	if buf, ok := w.Writer.(*bytes.Buffer); ok && buf.Len() != 0 {
 		if _, err := io.Copy(wr, buf); err != nil {
 			w.err = err
 			w.done()
@@ -81,3 +85,27 @@ func (w *WriteBuffer) done() {
 		w.wg.Done()
 	}
 }
+
+// OnceBuffer will catch only first message
+type OnceBuffer struct {
+	buf []byte
+}
+
+func (o *OnceBuffer) Write(p []byte) (n int, err error) {
+	if o.buf == nil {
+		o.buf = p
+	}
+	return 0, io.EOF
+}
+
+func (o *OnceBuffer) WriteTo(w io.Writer) (n int64, err error) {
+	return io.Copy(w, bytes.NewReader(o.buf))
+}
+
+func (o *OnceBuffer) Buffer() []byte {
+	return o.buf
+}
+
+func (o *OnceBuffer) Len() int {
+	return len(o.buf)
+}