speedtestlocal.go

package main

import (
	"context"
	"flag"
	"log"
	"net"
	"time"
)

const (
	BYTE     = 1
	KB       = 1024 * BYTE
	MB       = 1024 * KB
	BUFSIZE  = 1 * KB
	SENDSIZE = 500 * MB
	RECVSIZE = 400 * MB
)

var (
	ADDR   string
	CLIENT bool
	SERVER bool
)

func init() {
	flag.StringVar(&ADDR, "addr", "localhost:8080", "address to send to")
	flag.BoolVar(&CLIENT, "client", false, "perform client operations")
	flag.BoolVar(&SERVER, "server", false, "perform server operations")

	flag.Parse()
}

func server(ctx context.Context, shutdown context.CancelFunc) {
	conn, err := net.ListenUDP("udp", &net.UDPAddr{
		IP:   []byte{0, 0, 0, 0},
		Port: 8080,
	})
	if err != nil {
		log.Println("error listening to udp:", err)
		shutdown()
	}
	defer conn.Close()

	buf := make([]byte, BUFSIZE)

	totalBytes := 0
	var tStart time.Time

	for {
		if ctx.Err() != nil {
			return
		}

		n, _, err := conn.ReadFromUDP(buf)
		if err != nil {
			log.Println("udp read error:", err)
			shutdown()
			return
		}

		if totalBytes == 0 {
			tStart = time.Now()
		}
		totalBytes += n
		if totalBytes >= RECVSIZE {
			dur := time.Now().Sub(tStart)
			log.Printf("%f Mbps: recv. %d bytes in %fs", float64(totalBytes)/MB/dur.Seconds(), totalBytes, dur.Seconds())
		}

		if n >= 1 && buf[0] == 1 {
			log.Println("Done.")
		}
	}
}

func client(ctx context.Context, shutdown context.CancelFunc) {
	conn, err := net.Dial("udp", ADDR)
	if err != nil {
		log.Println("dial error:", err)
		shutdown()
		return
	}
	defer conn.Close()

	buf := make([]byte, BUFSIZE)

	for i := 0; i < SENDSIZE; {
		n, err := conn.Write(buf)
		if err != nil {
			log.Println("write err:", err)
			shutdown()
			return
		}
		i += n
	}

	// send end
	buf[0] = 1
	_, err = conn.Write(buf)
	if err != nil {
		log.Println("error writing end:", err)
	}
}

func main() {
	ctx, shutdown := context.WithCancel(context.Background())
	ctx, shutdown = context.WithTimeout(ctx, 1*time.Minute)

	if CLIENT {
		go client(ctx, shutdown)
	}
	if SERVER {
		go server(ctx, shutdown)
	}
	<-ctx.Done()
}