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
|
package ratelimit
import (
"bytes"
"sync"
"testing"
"time"
"github.com/NebulousLabs/fastrand"
)
// TestRLSimpleWriteRead tests a simple rate-limited write and read operation.
func TestRLSimpleWriteRead(t *testing.T) {
// Set limits
packetSize := uint64(64)
bps := int64(1000)
rl := NewRateLimit(bps, bps, packetSize)
// Create a io.ReadWriter.
rw := bytes.NewBuffer(make([]byte, 0))
// Wrap it into a rate limited ReadWriter.
c := make(chan struct{})
defer close(c)
rlc := NewRLReadWriter(rw, rl, c)
// Create 1mb to write.
data := fastrand.Bytes(1000)
// Write data while measuring time.
start := time.Now()
n, err := rlc.Write(data)
d := time.Since(start)
// Check for errors
if n < len(data) {
t.Error("Not whole data was written")
}
if err != nil {
t.Error("Failed to write data", err)
}
// Check the duration. We need to subtract packetSize since the time will
// be off by one packet. That's because the last written packet will finish
// faster than anticipated.
if d.Seconds() < float64(uint64(len(data))-packetSize)/float64(bps) {
t.Error("Write didn't take long enough", d.Seconds())
}
// Read data back from file while measuring time.
readData := make([]byte, len(data))
start = time.Now()
n, err = rlc.Read(readData)
d = time.Since(start)
// Check for errors
if n < len(data) {
t.Error("Not whole data was read")
}
if err != nil {
t.Error("Failed to read data", err)
}
// Check the duration again. Should be the same time.
if d.Seconds() < float64(uint64(len(data))-packetSize)/float64(bps) {
t.Error("Read didn't take long enough", d.Seconds())
}
// Check if the read data is the same as the written one.
if bytes.Compare(readData, data) != 0 {
t.Error("Read data doesn't match written data")
}
}
// TestRLParallelWriteRead tests a parallel rate-limited write and read operations.
func TestRLParallelWriteRead(t *testing.T) {
// Set limits
bps := int64(1000)
bytesToWrite := int(bps)
rl := NewRateLimit(bps, bps, 4096)
// f creates a rate limited buffer, writes some data to it and reads it
// afterwards.
f := func() {
// Create a io.ReadWriter.
rw := bytes.NewBuffer(make([]byte, 0))
// Wrap it into a rate limited ReadWriter.
c := make(chan struct{})
defer close(c)
rlc := NewRLReadWriter(rw, rl, c)
// Create 1mb to write.
data := fastrand.Bytes(bytesToWrite)
// Write data while measuring time.
n, err := rlc.Write(data)
// Check for errors
if n < len(data) {
t.Error("Not whole data was written")
}
if err != nil {
t.Error("Failed to write data", err)
}
// Read data back from file while measuring time.
readData := make([]byte, len(data))
n, err = rlc.Read(readData)
// Check for errors
if n < len(data) {
t.Error("Not whole data was read")
}
if err != nil {
t.Error("Failed to read data", err)
}
// Check if the read data is the same as the written one.
if bytes.Compare(readData, data) != 0 {
t.Error("Read data doesn't match written data")
}
}
// Start a few threads and wait for them to finish.
var wg sync.WaitGroup
start := time.Now()
numThreads := 10
for i := 0; i < numThreads; i++ {
wg.Add(1)
go func() {
f()
wg.Done()
}()
}
wg.Wait()
d := time.Since(start)
// d should be around 9 seconds. Each threads reads and writes exactly 1
// second of data but the first thread can write instantly.
if d.Seconds() < float64(numThreads-1) || d.Seconds() > float64(numThreads) {
t.Fatalf("Test should run between %v and %v seconds but was %v", numThreads-1, numThreads, d.Seconds())
}
}
|
