Files
netbird/relay/testec2/main.go
Zoltán Papp ce383eb259 Add ec2 test
2024-08-14 16:39:40 +02:00

384 lines
9.3 KiB
Go

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