udpstresstest.go

package main

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"sync"
	"time"
)

const numClients = 1000
const dataLen = 80
const numMessages = 5

func main() {
	addrChan := make(chan net.Addr, 1)
	go func() {
		err := runServer(addrChan)
		panic(err)
	}()

	serverAddr := <-addrChan
	var wg sync.WaitGroup
	for i := 0; i < numClients; i++ {
		wg.Add(1)
		go func() {
			err := runClient(serverAddr)
			if err != nil {
				panic(err)
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

// the server echos everything it gets
func runServer(addrChan chan<- net.Addr) error {
	addr, err := net.ResolveUDPAddr("udp", "localhost:0")
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return err
	}
	defer conn.Close()
	fmt.Println("Listening on", conn.LocalAddr())
	addrChan <- conn.LocalAddr()

	for {
		b := make([]byte, 100)
		n, addr, err := conn.ReadFrom(b)
		if err != nil {
			return err
		}
		b = b[:n]
		_, err = conn.WriteTo(b, addr)
		if err != nil {
			return err
		}
	}
}

// send numMessages messages to the server and wait for the echo
func runClient(remoteAddr net.Addr) error {
	clientAddr, err := net.ResolveUDPAddr("udp", "localhost:0")
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", clientAddr)
	if err != nil {
		return err
	}
	defer conn.Close()
	conn.SetDeadline(time.Now().Add(5 * time.Second))

	data := bytes.Repeat([]byte{'f'}, dataLen)
	go func() {
		for i := 0; i < numMessages; i++ {
			_, err = conn.WriteTo(data, remoteAddr)
			if err != nil {
				panic(err)
			}
			time.Sleep(time.Millisecond)
		}
	}()

	for i := 0; i < numMessages; i++ {
		b := make([]byte, 100)
		n, _, err := conn.ReadFrom(b)
		if err != nil {
			fmt.Println("Error while reading message", i)
			return err
		}
		b = b[:n]
		if !bytes.Equal(b, data) {
			return errors.New("not equal")
		}
	}
	fmt.Printf("Successfully transferred %d bytes of data from %s\n", dataLen, conn.LocalAddr().String())
	return nil
}