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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
package defs
import (
"bytes"
"crypto/rand"
"fmt"
"io"
"log"
"sync"
"time"
)
// BytesCounter implements io.Reader and io.Writer interface, for counting bytes being read/written in HTTP requests
type BytesCounter struct {
start time.Time
pos int
total uint64
payload []byte
reader io.ReadSeeker
mebi bool
uploadSize int
lock *sync.Mutex
}
func NewCounter() *BytesCounter {
return &BytesCounter{
lock: &sync.Mutex{},
}
}
// Write implements io.Writer
func (c *BytesCounter) Write(p []byte) (int, error) {
n := len(p)
c.lock.Lock()
c.total += uint64(n)
c.lock.Unlock()
return n, nil
}
// Read implements io.Reader
func (c *BytesCounter) Read(p []byte) (int, error) {
n, err := c.reader.Read(p)
c.lock.Lock()
c.total += uint64(n)
c.pos += n
if c.pos == c.uploadSize {
c.resetReader()
}
c.lock.Unlock()
return n, err
}
// SetBase sets the base for dividing bytes into megabyte or mebibyte
func (c *BytesCounter) SetMebi(mebi bool) {
c.mebi = mebi
}
// SetUploadSize sets the size of payload being uploaded
func (c *BytesCounter) SetUploadSize(uploadSize int) {
c.uploadSize = uploadSize * 1024
}
// AvgBytes returns the average bytes/second
func (c *BytesCounter) AvgBytes() float64 {
return float64(c.total) / time.Now().Sub(c.start).Seconds()
}
// AvgMbps returns the average mbits/second
func (c *BytesCounter) AvgMbps() float64 {
var base float64 = 125000
if c.mebi {
base = 131072
}
return c.AvgBytes() / base
}
// AvgHumanize returns the average bytes/kilobytes/megabytes/gigabytes (or bytes/kibibytes/mebibytes/gibibytes) per second
func (c *BytesCounter) AvgHumanize() string {
val := c.AvgBytes()
var base float64 = 1000
if c.mebi {
base = 1024
}
if val < base {
return fmt.Sprintf("%.2f bytes/s", val)
} else if val/base < base {
return fmt.Sprintf("%.2f KB/s", val/base)
} else if val/base/base < base {
return fmt.Sprintf("%.2f MB/s", val/base/base)
} else {
return fmt.Sprintf("%.2f GB/s", val/base/base/base)
}
}
// GenerateBlob generates a random byte array of `uploadSize` in the `payload` field, and sets the `reader` field to
// read from it
func (c *BytesCounter) GenerateBlob() {
c.payload = getRandomData(c.uploadSize)
c.reader = bytes.NewReader(c.payload)
}
// resetReader resets the `reader` field to 0 position
func (c *BytesCounter) resetReader() (int64, error) {
c.pos = 0
return c.reader.Seek(0, 0)
}
// Start will set the `start` field to current time
func (c *BytesCounter) Start() {
c.start = time.Now()
}
// Total returns the total bytes read/written
func (c *BytesCounter) Total() uint64 {
return c.total
}
// CurrentSpeed returns the current bytes/second
func (c *BytesCounter) CurrentSpeed() float64 {
return float64(c.total) / time.Now().Sub(c.start).Seconds()
}
// SeekWrapper is a wrapper around io.Reader to give it a noop io.Seeker interface
type SeekWrapper struct {
io.Reader
}
// Seek implements the io.Seeker interface
func (r *SeekWrapper) Seek(offset int64, whence int) (int64, error) {
return 0, nil
}
// getAvg returns the average value of an float64 array
func getAvg(vals []float64) float64 {
var total float64
for _, v := range vals {
total += v
}
return total / float64(len(vals))
}
// getRandomData returns an `length` sized array of random bytes
func getRandomData(length int) []byte {
data := make([]byte, length)
if _, err := rand.Read(data); err != nil {
log.Fatalf("Failed to generate random data: %s", err)
}
return data
}
|
