mirror of
https://github.com/netbirdio/netbird.git
synced 2026-04-25 19:56:46 +00:00
Add ec2 test
This commit is contained in:
383
relay/testec2/main.go
Normal file
383
relay/testec2/main.go
Normal file
@@ -0,0 +1,383 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"flag"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/relay/testec2/tun"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
)
|
||||
|
||||
var (
|
||||
dataSize = 1024 * 1024 * 50 // 50MB
|
||||
pairs = []int{1, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100}
|
||||
signalListenAddress = ":8081"
|
||||
|
||||
relaySrvAddress string
|
||||
turnSrvAddress string
|
||||
signalURL string
|
||||
udpListener string
|
||||
)
|
||||
|
||||
type testResult struct {
|
||||
numOfPairs int
|
||||
duration time.Duration
|
||||
speed float64
|
||||
}
|
||||
|
||||
func (tr testResult) Speed() string {
|
||||
speed := tr.speed
|
||||
var unit string
|
||||
|
||||
switch {
|
||||
case speed < 1024:
|
||||
unit = "B/s"
|
||||
case speed < 1048576:
|
||||
speed /= 1024
|
||||
unit = "KB/s"
|
||||
case speed < 1073741824:
|
||||
speed /= 1048576
|
||||
unit = "MB/s"
|
||||
default:
|
||||
speed /= 1073741824
|
||||
unit = "GB/s"
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%.2f %s", speed, unit)
|
||||
}
|
||||
|
||||
func seedRandomData(size int) ([]byte, error) {
|
||||
token := make([]byte, size)
|
||||
_, err := rand.Read(token)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return token, nil
|
||||
}
|
||||
|
||||
func avg(transferDuration []time.Duration) (time.Duration, float64) {
|
||||
var totalDuration time.Duration
|
||||
for _, d := range transferDuration {
|
||||
totalDuration += d
|
||||
}
|
||||
avgDuration := totalDuration / time.Duration(len(transferDuration))
|
||||
bps := float64(dataSize) / avgDuration.Seconds()
|
||||
return avgDuration, bps
|
||||
}
|
||||
|
||||
func RelayReceiverMain() []testResult {
|
||||
testResults := make([]testResult, 0, len(pairs))
|
||||
for _, p := range pairs {
|
||||
tr := testResult{numOfPairs: p}
|
||||
td := relayReceive(relaySrvAddress, p)
|
||||
tr.duration, tr.speed = avg(td)
|
||||
|
||||
testResults = append(testResults, tr)
|
||||
}
|
||||
|
||||
return testResults
|
||||
}
|
||||
|
||||
func RelaySenderMain() {
|
||||
log.Infof("starting sender")
|
||||
log.Infof("starting seed phase")
|
||||
|
||||
testData, err := seedRandomData(dataSize)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to seed random data: %s", err)
|
||||
}
|
||||
|
||||
log.Infof("data size: %d", len(testData))
|
||||
|
||||
for n, p := range pairs {
|
||||
log.Infof("running test with %d pairs", p)
|
||||
relayTransfer(relaySrvAddress, testData, p)
|
||||
|
||||
// grant time to prepare new receivers
|
||||
if n < len(pairs)-1 {
|
||||
time.Sleep(3 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// TRUNSenderMain is the sender
|
||||
// - allocate turn clients
|
||||
// - send relayed addresses to signal server in batch
|
||||
// - wait for signal server to send back addresses in a map
|
||||
// - send test data to each address in parallel
|
||||
func TRUNSenderMain() {
|
||||
log.Infof("starting TURN sender test")
|
||||
|
||||
log.Infof("starting seed random data: %d", dataSize)
|
||||
testData, err := seedRandomData(dataSize)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to seed random data: %s", err)
|
||||
}
|
||||
|
||||
ss := SignalClient{signalURL}
|
||||
|
||||
for _, p := range pairs {
|
||||
log.Infof("running test with %d pairs", p)
|
||||
|
||||
turnConns := make(map[string]*TurnConn)
|
||||
addresses := make([]string, 0, len(pairs))
|
||||
for i := 0; i < p; i++ {
|
||||
tc := AllocateTurnClient(turnSrvAddress)
|
||||
log.Infof("allocated turn client: %s", tc.Address().String())
|
||||
turnConns[tc.Address().String()] = tc
|
||||
addresses = append(addresses, tc.Address().String())
|
||||
}
|
||||
|
||||
log.Infof("send addresses via signal server: %d", len(addresses))
|
||||
clientAddresses, err := ss.SendAddress(addresses)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to send address: %s", err)
|
||||
}
|
||||
log.Infof("received addresses: %v", clientAddresses.Address)
|
||||
|
||||
var i int
|
||||
devices := make([]*tun.Device, 0, len(clientAddresses.Address))
|
||||
for k, v := range clientAddresses.Address {
|
||||
tc, ok := turnConns[k]
|
||||
if !ok {
|
||||
log.Fatalf("failed to find turn conn: %s", k)
|
||||
}
|
||||
|
||||
addr, err := net.ResolveUDPAddr("udp", v)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to resolve udp address: %s", err)
|
||||
}
|
||||
device := &tun.Device{
|
||||
Name: fmt.Sprintf("mtun-sender-%d", i),
|
||||
IP: fmt.Sprintf("10.0.%d.1", i),
|
||||
PConn: tc.relayConn,
|
||||
DstAddr: addr,
|
||||
}
|
||||
|
||||
err = device.Up()
|
||||
if err != nil {
|
||||
log.Fatalf("failed to bring up device: %s", err)
|
||||
}
|
||||
|
||||
devices = append(devices, device)
|
||||
i++
|
||||
}
|
||||
|
||||
log.Infof("waiting for tcpListeners to be ready")
|
||||
time.Sleep(2 * time.Second)
|
||||
|
||||
tcpConns := make([]net.Conn, 0, len(devices))
|
||||
for i := range devices {
|
||||
addr := fmt.Sprintf("10.0.%d.2:9999", i)
|
||||
log.Infof("dialing: %s", addr)
|
||||
tcpConn, err := net.Dial("tcp", addr)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to dial tcp: %s", err)
|
||||
}
|
||||
tcpConns = append(tcpConns, tcpConn)
|
||||
}
|
||||
|
||||
log.Infof("start test data transfer for %d pairs", len(devices))
|
||||
testDataLen := len(testData)
|
||||
wg := sync.WaitGroup{}
|
||||
for i, tcpConn := range tcpConns {
|
||||
log.Infof("sending test data to device: %d", i)
|
||||
wg.Add(1)
|
||||
go func(i int, tcpConn net.Conn) {
|
||||
defer wg.Done()
|
||||
defer tcpConn.Close()
|
||||
|
||||
log.Infof("start to sending test data: %s", tcpConn.RemoteAddr())
|
||||
|
||||
si := NewStartInidication(time.Now(), testDataLen)
|
||||
_, err = tcpConn.Write(si)
|
||||
if err != nil {
|
||||
log.Errorf("failed to write to tcp: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
pieceSize := 1024
|
||||
for j := 0; j < testDataLen; j += pieceSize {
|
||||
end := j + pieceSize
|
||||
if end > testDataLen {
|
||||
end = testDataLen
|
||||
}
|
||||
_, writeErr := tcpConn.Write(testData[j:end])
|
||||
if writeErr != nil {
|
||||
log.Errorf("failed to write to tcp conn: %s", writeErr)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
time.Sleep(3 * time.Second)
|
||||
}(i, tcpConn)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
for _, d := range devices {
|
||||
_ = d.Close()
|
||||
}
|
||||
|
||||
log.Infof("test finished with %d pairs", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TURNReaderMain() []testResult {
|
||||
log.Infof("starting TURN receiver test")
|
||||
si := NewSignalService()
|
||||
go func() {
|
||||
log.Infof("starting signal server")
|
||||
err := si.Listen(signalListenAddress)
|
||||
if err != nil {
|
||||
log.Errorf("failed to listen: %s", err)
|
||||
}
|
||||
}()
|
||||
|
||||
testResults := make([]testResult, 0, len(pairs))
|
||||
for range pairs {
|
||||
log.Infof("waiting for addresses")
|
||||
addresses := <-si.AddressesChan
|
||||
log.Infof("received addresses: %d", len(addresses))
|
||||
|
||||
conns := make([]*net.UDPConn, 0, len(addresses))
|
||||
clientAddresses := make(map[string]string, len(addresses))
|
||||
devices := make([]*tun.Device, 0, len(conns))
|
||||
for i, addr := range addresses {
|
||||
localAddr, err := net.ResolveUDPAddr("udp", udpListener)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to resolve UDP address: %s", err)
|
||||
}
|
||||
|
||||
conn, err := net.ListenUDP("udp", localAddr)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to create UDP connection: %s", err)
|
||||
}
|
||||
|
||||
conns = append(conns, conn)
|
||||
clientAddresses[addr] = conn.LocalAddr().String()
|
||||
|
||||
dstAddr, err := net.ResolveUDPAddr("udp", addr)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to resolve address: %s", err)
|
||||
}
|
||||
|
||||
device := &tun.Device{
|
||||
Name: fmt.Sprintf("mtun-%d", i),
|
||||
IP: fmt.Sprintf("10.0.%d.2", i),
|
||||
PConn: conn,
|
||||
DstAddr: dstAddr,
|
||||
}
|
||||
|
||||
err = device.Up()
|
||||
if err != nil {
|
||||
log.Fatalf("failed to bring up device: %s, %s", device.Name, err)
|
||||
}
|
||||
devices = append(devices, device)
|
||||
}
|
||||
|
||||
// send back local addresses
|
||||
log.Infof("response addresses back: %d", len(clientAddresses))
|
||||
si.ClientAddressChan <- clientAddresses
|
||||
|
||||
durations := make(chan time.Duration, len(conns))
|
||||
for _, d := range devices {
|
||||
go func(d *tun.Device) {
|
||||
tcpListener, err := net.Listen("tcp", d.IP+":9999")
|
||||
if err != nil {
|
||||
log.Fatalf("failed to listen on tcp: %s", err)
|
||||
}
|
||||
defer tcpListener.Close()
|
||||
log := log.WithField("device", tcpListener.Addr())
|
||||
|
||||
tcpConn, err := tcpListener.Accept()
|
||||
if err != nil {
|
||||
log.Fatalf("failed to accept connection: %s", err)
|
||||
}
|
||||
log.Infof("remote peer connected")
|
||||
|
||||
buf := make([]byte, 103)
|
||||
n, err := tcpConn.Read(buf)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to read from tcp: %s", err)
|
||||
}
|
||||
|
||||
si := DecodeStartIndication(buf[:n])
|
||||
log.Infof("received start indication: %v, %d", si, n)
|
||||
|
||||
buf = make([]byte, 8192)
|
||||
i, err := tcpConn.Read(buf)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to read from tcp: %s", err)
|
||||
}
|
||||
now := time.Now()
|
||||
|
||||
for i < si.TransferSize {
|
||||
n, err := tcpConn.Read(buf)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to read from tcp: %s", err)
|
||||
}
|
||||
i += n
|
||||
}
|
||||
log.Infof("finished reading")
|
||||
durations <- time.Since(now)
|
||||
}(d)
|
||||
}
|
||||
|
||||
durationsList := make([]time.Duration, 0, len(conns))
|
||||
for d := range durations {
|
||||
durationsList = append(durationsList, d)
|
||||
if len(durationsList) == len(conns) {
|
||||
close(durations)
|
||||
}
|
||||
}
|
||||
|
||||
avgDuration, avgSpeed := avg(durationsList)
|
||||
ts := testResult{
|
||||
numOfPairs: len(conns),
|
||||
duration: avgDuration,
|
||||
speed: avgSpeed,
|
||||
}
|
||||
testResults = append(testResults, ts)
|
||||
|
||||
for _, d := range devices {
|
||||
_ = d.Close()
|
||||
}
|
||||
}
|
||||
|
||||
return testResults
|
||||
}
|
||||
|
||||
func main() {
|
||||
var mode string
|
||||
|
||||
_ = util.InitLog("debug", "console")
|
||||
flag.StringVar(&mode, "mode", "sender", "sender or receiver mode")
|
||||
flag.Parse()
|
||||
|
||||
relaySrvAddress = os.Getenv("TEST_RELAY_SERVER") // rel://ip:port
|
||||
turnSrvAddress = os.Getenv("TEST_TURN_SERVER") // ip:3478
|
||||
signalURL = os.Getenv("TEST_SIGNAL_URL") // http://receiver_ip:8081
|
||||
udpListener = os.Getenv("TEST_UDP_LISTENER") // IP:0
|
||||
|
||||
if mode == "receiver" {
|
||||
relayResult := RelayReceiverMain()
|
||||
turnResults := TURNReaderMain()
|
||||
for i := 0; i < len(turnResults); i++ {
|
||||
log.Infof("pairs: %d,\tRelay speed:\t%s,\trelay duration:\t%s", relayResult[i].numOfPairs, relayResult[i].Speed(), relayResult[i].duration)
|
||||
log.Infof("pairs: %d,\tTURN speed:\t%s,\tturn duration:\t%s", turnResults[i].numOfPairs, turnResults[i].Speed(), turnResults[i].duration)
|
||||
}
|
||||
} else {
|
||||
RelaySenderMain()
|
||||
// grant time for receiver to start
|
||||
time.Sleep(3 * time.Second)
|
||||
TRUNSenderMain()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user