api, metrics: add number of bytes received and sent from/to all entities (#1235)

* API: number of bytes received/sent from/to RTSP connections
* API: number of bytes received/sent from/to RTSP sessions
* API: number of bytes received/sent from/to RTMP connections
* API: number of bytes sent to HLS connections
* API: number of bytes received from paths
* metrics of all the above

internal/rtmp/bytecounter/reader.go

@@ -1,42 +1,36 @@
 package bytecounter
 
 import (
-	"bufio"
 	"io"
+	"sync/atomic"
 )
 
-type readerInner struct {
-	r     io.Reader
-	count uint32
-}
-
-func (r *readerInner) Read(p []byte) (int, error) {
-	n, err := r.r.Read(p)
-	r.count += uint32(n)
-	return n, err
-}
-
 // Reader allows to count read bytes.
 type Reader struct {
-	ri *readerInner
-	*bufio.Reader
+	r     io.Reader
+	count uint64
 }
 
 // NewReader allocates a Reader.
 func NewReader(r io.Reader) *Reader {
-	ri := &readerInner{r: r}
 	return &Reader{
-		ri:     ri,
-		Reader: bufio.NewReader(ri),
+		r: r,
 	}
 }
 
-// Count returns read bytes.
-func (r Reader) Count() uint32 {
-	return r.ri.count
+// Read implements io.Reader.
+func (r *Reader) Read(p []byte) (int, error) {
+	n, err := r.r.Read(p)
+	atomic.AddUint64(&r.count, uint64(n))
+	return n, err
+}
+
+// Count returns received bytes.
+func (r *Reader) Count() uint64 {
+	return atomic.LoadUint64(&r.count)
 }
 
 // SetCount sets read bytes.
-func (r *Reader) SetCount(v uint32) {
-	r.ri.count = v
+func (r *Reader) SetCount(v uint64) {
+	atomic.StoreUint64(&r.count, v)
 }

internal/rtmp/bytecounter/reader_test.go

@@ -19,5 +19,5 @@ func TestReader(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 64, n)
 
-	require.Equal(t, uint32(100+1024), r.Count())
+	require.Equal(t, uint64(100+64), r.Count())
 }

internal/rtmp/bytecounter/writer.go

@@ -2,12 +2,13 @@ package bytecounter
 
 import (
 	"io"
+	"sync/atomic"
 )
 
 // Writer allows to count written bytes.
 type Writer struct {
 	w     io.Writer
-	count uint32
+	count uint64
 }
 
 // NewWriter allocates a Writer.
@@ -20,16 +21,16 @@ func NewWriter(w io.Writer) *Writer {
 // Write implements io.Writer.
 func (w *Writer) Write(p []byte) (int, error) {
 	n, err := w.w.Write(p)
-	w.count += uint32(n)
+	atomic.AddUint64(&w.count, uint64(n))
 	return n, err
 }
 
-// Count returns written bytes.
-func (w Writer) Count() uint32 {
-	return w.count
+// Count returns sent bytes.
+func (w *Writer) Count() uint64 {
+	return atomic.LoadUint64(&w.count)
 }
 
-// SetCount sets written bytes.
-func (w *Writer) SetCount(v uint32) {
-	w.count = v
+// SetCount sets sent bytes.
+func (w *Writer) SetCount(v uint64) {
+	atomic.StoreUint64(&w.count, v)
 }

internal/rtmp/bytecounter/writer_test.go

@@ -14,5 +14,5 @@ func TestWriter(t *testing.T) {
 	w.SetCount(100)
 
 	w.Write(bytes.Repeat([]byte{0x01}, 64))
-	require.Equal(t, uint32(100+64), w.Count())
+	require.Equal(t, uint64(100+64), w.Count())
 }

internal/rtmp/conn.go

@@ -114,10 +114,19 @@ type Conn struct {
 
 // NewConn initializes a connection.
 func NewConn(rw io.ReadWriter) *Conn {
-	c := &Conn{}
-	c.bc = bytecounter.NewReadWriter(rw)
-	c.mrw = message.NewReadWriter(c.bc, false)
-	return c
+	return &Conn{
+		bc: bytecounter.NewReadWriter(rw),
+	}
+}
+
+// BytesReceived returns the number of bytes received.
+func (c *Conn) BytesReceived() uint64 {
+	return c.bc.Reader.Count()
+}
+
+// BytesSent returns the number of bytes sent.
+func (c *Conn) BytesSent() uint64 {
+	return c.bc.Writer.Count()
 }
 
 func (c *Conn) readCommand() (*message.MsgCommandAMF0, error) {
@@ -161,6 +170,8 @@ func (c *Conn) InitializeClient(u *url.URL, isPublishing bool) error {
 		return err
 	}
 
+	c.mrw = message.NewReadWriter(c.bc, false)
+
 	err = c.mrw.Write(&message.MsgSetWindowAckSize{
 		Value: 2500000,
 	})
@@ -319,6 +330,8 @@ func (c *Conn) InitializeServer() (*url.URL, bool, error) {
 		return nil, false, err
 	}
 
+	c.mrw = message.NewReadWriter(c.bc, false)
+
 	cmd, err := c.readCommand()
 	if err != nil {
 		return nil, false, err

internal/rtmp/conn_test.go

@@ -244,6 +244,14 @@ func TestInitializeClient(t *testing.T) {
 			err = conn.InitializeClient(u, ca == "publish")
 			require.NoError(t, err)
 
+			if ca == "read" {
+				require.Equal(t, uint64(3421), conn.BytesReceived())
+				require.Equal(t, uint64(3409), conn.BytesSent())
+			} else {
+				require.Equal(t, uint64(3427), conn.BytesReceived())
+				require.Equal(t, uint64(3466), conn.BytesSent())
+			}
+
 			<-done
 		})
 	}

internal/rtmp/rawmessage/reader.go

@@ -1,6 +1,7 @@
 package rawmessage
 
 import (
+	"bufio"
 	"errors"
 	"fmt"
 	"time"
@@ -22,14 +23,14 @@ type readerChunkStream struct {
 }
 
 func (rc *readerChunkStream) readChunk(c chunk.Chunk, chunkBodySize uint32) error {
-	err := c.Read(rc.mr.r, chunkBodySize)
+	err := c.Read(rc.mr.br, chunkBodySize)
 	if err != nil {
 		return err
 	}
 
 	// check if an ack is needed
 	if rc.mr.ackWindowSize != 0 {
-		count := rc.mr.r.Count()
+		count := uint32(rc.mr.r.Count())
 		diff := count - rc.mr.lastAckCount
 
 		if diff > (rc.mr.ackWindowSize) {
@@ -210,6 +211,7 @@ type Reader struct {
 	r           *bytecounter.Reader
 	onAckNeeded func(uint32) error
 
+	br            *bufio.Reader
 	chunkSize     uint32
 	ackWindowSize uint32
 	lastAckCount  uint32
@@ -223,8 +225,11 @@ type Reader struct {
 
 // NewReader allocates a Reader.
 func NewReader(r *bytecounter.Reader, onAckNeeded func(uint32) error) *Reader {
+	br := bufio.NewReader(r)
+
 	return &Reader{
 		r:            r,
+		br:           br,
 		onAckNeeded:  onAckNeeded,
 		chunkSize:    128,
 		chunkStreams: make(map[byte]*readerChunkStream),
@@ -244,7 +249,7 @@ func (r *Reader) SetWindowAckSize(v uint32) {
 // Read reads a Message.
 func (r *Reader) Read() (*Message, error) {
 	for {
-		byt, err := r.r.ReadByte()
+		byt, err := r.br.ReadByte()
 		if err != nil {
 			return nil, err
 		}
@@ -258,7 +263,7 @@ func (r *Reader) Read() (*Message, error) {
 			r.chunkStreams[chunkStreamID] = rc
 		}
 
-		r.r.UnreadByte()
+		r.br.UnreadByte()
 
 		msg, err := rc.readMessage(typ)
 		if err != nil {

internal/rtmp/rawmessage/reader_test.go

@@ -198,8 +198,7 @@ func TestReader(t *testing.T) {
 	for _, ca := range cases {
 		t.Run(ca.name, func(t *testing.T) {
 			var buf bytes.Buffer
-			bcr := bytecounter.NewReader(&buf)
-			r := NewReader(bcr, func(count uint32) error {
+			r := NewReader(bytecounter.NewReader(&buf), func(count uint32) error {
 				return nil
 			})
 
@@ -224,14 +223,14 @@ func TestReaderAcknowledge(t *testing.T) {
 			onAckCalled := make(chan struct{})
 
 			var buf bytes.Buffer
-			bcr := bytecounter.NewReader(&buf)
-			r := NewReader(bcr, func(count uint32) error {
+			bc := bytecounter.NewReader(&buf)
+			r := NewReader(bc, func(count uint32) error {
 				close(onAckCalled)
 				return nil
 			})
 
 			if ca == "overflow" {
-				bcr.SetCount(4294967096)
+				bc.SetCount(4294967096)
 				r.lastAckCount = 4294967096
 			}
 

internal/rtmp/rawmessage/writer.go

@@ -20,7 +20,7 @@ type writerChunkStream struct {
 func (wc *writerChunkStream) writeChunk(c chunk.Chunk) error {
 	// check if we received an acknowledge
 	if wc.mw.checkAcknowledge && wc.mw.ackWindowSize != 0 {
-		diff := wc.mw.w.Count() - wc.mw.ackValue
+		diff := uint32(wc.mw.w.Count()) - wc.mw.ackValue
 
 		if diff > (wc.mw.ackWindowSize * 3 / 2) {
 			return fmt.Errorf("no acknowledge received within window")