-
-
Notifications
You must be signed in to change notification settings - Fork 2
/
byte_read_writer.go
52 lines (43 loc) · 880 Bytes
/
byte_read_writer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
package metered
import (
"io"
"sync/atomic"
)
// Implements io.ReadWriter
type ReadWriter struct {
rw io.ReadWriter
rx uint64
tx uint64
}
func NewReadWriter(rw io.ReadWriter) *ReadWriter {
return &ReadWriter{
rw: rw,
}
}
func (rw *ReadWriter) Read(p []byte) (n int, err error) {
n, err = rw.rw.Read(p)
atomic.AddUint64(&rw.rx, uint64(n))
return
}
func (rw *ReadWriter) Write(p []byte) (n int, err error) {
n, err = rw.rw.Write(p)
atomic.AddUint64(&rw.tx, uint64(n))
return
}
func (rw *ReadWriter) Rx() uint64 {
return atomic.LoadUint64(&rw.rx)
}
func (rw *ReadWriter) Tx() uint64 {
return atomic.LoadUint64(&rw.tx)
}
// Implements io.ReadWriteCloser
type ReadWriteCloser struct {
*ReadWriter
io.Closer
}
func NewReadWriteCloser(rw io.ReadWriteCloser) *ReadWriteCloser {
return &ReadWriteCloser{
ReadWriter: NewReadWriter(rw),
Closer: rw,
}
}