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())
}
}
|