-
Notifications
You must be signed in to change notification settings - Fork 278
/
countingconn.go
57 lines (49 loc) · 1.5 KB
/
countingconn.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
53
54
55
56
57
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0
package worker
import (
"net"
"sync"
)
// countingConn is a `net.Conn` implementation that records the bytes that go
// across Read() and Write(). All other `net.Conn` function calls are a
// pass-through to the underlying `net.Conn`, meaning it's also safe to call
// those functions directly on the underlying object, if you have access to it.
type countingConn struct {
net.Conn
bytesRead int64
bytesWritten int64
// Use mutex for counters as net.Conn methods may be called concurrently
// https://github.com/golang/go/issues/27203#issuecomment-415854958
mu sync.Mutex
}
// BytesRead reports the number of bytes read so far
func (c *countingConn) BytesRead() int64 {
c.mu.Lock()
defer c.mu.Unlock()
return c.bytesRead
}
// BytesWritten reports the number of bytes written so far
func (c *countingConn) BytesWritten() int64 {
c.mu.Lock()
defer c.mu.Unlock()
return c.bytesWritten
}
// Read wraps the embedded conn's Read() and counts the number of bytes read
// (the number of bytes the client sent to us).
func (c *countingConn) Read(in []byte) (int, error) {
n, err := c.Conn.Read(in)
c.mu.Lock()
c.bytesRead += int64(n)
c.mu.Unlock()
return n, err
}
// Write wraps the embedded conn's Write() and counts the number of bytes
// written (the number of bytes we sent to the client).
func (c *countingConn) Write(in []byte) (int, error) {
n, err := c.Conn.Write(in)
c.mu.Lock()
c.bytesWritten += int64(n)
c.mu.Unlock()
return n, err
}