main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

var _workers = flag.Int("workers", 10, "the number of worker go routines")
var _targets = flag.String("targets", "127.0.0.1", "the targets to use")
var _startPort = flag.Int("base-port", 8000, "the first port to listen on")
var _portCount = flag.Int("port-count", 10, "the number of ports")
var _deadline = flag.Duration("deadline", 200*time.Millisecond, "the deadline")

type state struct {
	count            int
	failures         int
	deadlineExceeded int
}

func runWorker(worker int, wg *sync.WaitGroup, target string, s *state, blockChan chan struct{}) {
	defer wg.Done()

	// stagger the time of the first transaction.
	time.Sleep(time.Duration(rand.Int31n(999)) * time.Millisecond)

	host, portStr, err := net.SplitHostPort(target)
	if err != nil {
		log.Printf(err.Error())
		return
	}

	ipAddr := net.ParseIP(host)
	if ipAddr == nil {
		log.Printf("invalid ip %s", host)
		return
	}

	port, err := strconv.Atoi(portStr)
	if err != nil {
		log.Printf("invalid ip %s", host)
		return
	}

	tcpAddr := &net.TCPAddr{
		IP:   ipAddr,
		Port: port,
	}

	req := []byte{1, 2, 3, 4}

	for {
		if err := dial(tcpAddr, target, req, s, worker); err != nil {
			log.Printf(err.Error())
		}
		select {
		case <-time.After(time.Second):
		case <-blockChan:
			return
		}
	}
}

func dial(tcpAddr *net.TCPAddr, target string, req []byte, s *state, worker int) error {
	s.count++

	start := time.Now()
	ctx, cancelFn := context.WithTimeout(context.Background(), time.Second)
	defer cancelFn()

	dialer := net.Dialer{}
	conn, err := dialer.DialContext(ctx, "tcp", target)

	// conn, err := net.Dial("tcp", target)
	// conn, err := net.DialTCP("tcp", nil, tcpAddr)
	// conn, err := net.MyDialTCP(ctx, "tcp", nil, tcpAddr)  // uses net.dialTCP
	if err != nil {
		log.Printf("Dial failed %v", err)
		s.failures++
		return err
	}

	connTime := time.Now()

	defer conn.Close()
	if err := conn.SetDeadline(time.Now().Add(time.Second)); err != nil {
		log.Printf("SetDeadline failed %v", err)
		s.failures++
		return err
	}

	if _, err := conn.Write(req); err != nil {
		log.Printf("Write failed %v", err)
		s.failures++
		return err
	}
	writeTime := time.Now()

	res := make([]byte, 16)
	_, err = conn.Read(res)
	if err != nil {
		log.Printf("Read failed %v", err)
		s.deadlineExceeded++
		return err
	}

	now := time.Now()
	totalTime := now.Sub(start)
	if totalTime > *_deadline {
		log.Printf("%s exceeded %s: conn: %s, write %s, read: %s, w: %d", totalTime, *_deadline, connTime.Sub(start), writeTime.Sub(connTime), now.Sub(writeTime), worker)
		s.deadlineExceeded++
	}
	return nil
}

func main() {
	flag.Parse()
	log.Printf("%d workers %s[%d:%d]", *_workers, *_targets, *_startPort, *_startPort+*_portCount)

	targets := strings.Split(*_targets, ",")
	addrs := make([]string, 0, *_portCount*len(targets))

	for _, t := range targets {
		names, err := net.LookupHost(t)
		if err != nil {
			log.Printf("lookup(%s) failed: %v", t, err)
			return
		}

		if len(names) == 0 {
			log.Printf("lookup(%s) no results", t)
			return
		}

		for i := *_startPort; i < *_startPort+*_portCount; i++ {
			addrs = append(addrs, fmt.Sprintf("%s:%d", names[0], i))
		}
	}

	var wg sync.WaitGroup
	states := make([]state, *_workers)
	blockChan := make(chan struct{}, 0)
	start := time.Now()
	for i := 0; i < *_workers; i++ {
		wg.Add(1)
		go runWorker(i, &wg, addrs[i%len(addrs)], &states[i], blockChan)
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	<-c
	end := time.Now()
	close(blockChan)

	wg.Wait()
	count, failures, deadlineExceeded := 0, 0, 0
	for _, state := range states {
		count += state.count
		failures += state.failures
		deadlineExceeded += state.deadlineExceeded
	}

	dur := end.Sub(start)
	fmt.Println("")
	fmt.Printf("GOMAXPROCS=%d, requested conns/sec: %d, deadline: %s\n", runtime.GOMAXPROCS(0), *_workers, *_deadline)
	fmt.Printf("%d connections in %s, %0f conns/sec\n", count, dur, float64(count)/dur.Seconds())
	fmt.Printf("%d connections, %d failures, %d exceeed deadline (%f %%)\n", count, failures, deadlineExceeded, (100 * float32(deadlineExceeded) / float32(count)))
}