Rewrite once buffer for keyframes

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)
+}