mirror of
https://github.com/bluenviron/mediamtx.git
synced 2025-12-25 04:22:00 -08:00
36 lines
614 B
Go
36 lines
614 B
Go
package bytecounter
|
|
|
|
import (
|
|
"io"
|
|
"sync/atomic"
|
|
)
|
|
|
|
// Reader allows to count read bytes.
|
|
type Reader struct {
|
|
r io.Reader
|
|
count uint64
|
|
}
|
|
|
|
// NewReader allocates a Reader.
|
|
func NewReader(r io.Reader) *Reader {
|
|
return &Reader{
|
|
r: r,
|
|
}
|
|
}
|
|
|
|
// 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 uint64) {
|
|
atomic.StoreUint64(&r.count, v)
|
|
}
|