package main

import (
	"crypto/rand"
	"encoding/binary"
	"flag"
	"log"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"
)

var N = flag.Int("n", 100000, "Number of queries to make")
var parallel = flag.Int("parallel", 40, "Number of parallel goroutines")
var server = flag.String("server", "127.0.0.1:1053", "Server to query")

func main() {
	flag.Parse()

	go dnsServer()
	time.Sleep(100 * time.Millisecond)

	origID := dns.Id
	cryptID := func() uint16 {
		var v uint16
		err := binary.Read(rand.Reader, binary.BigEndian, &v)
		if err != nil {
			panic(err)
		}
		return v
	}

	dns.Id = origID
	doRun("math/rand", *server, *N, *parallel)
	dns.Id = cryptID
	doRun("crypto/rand", *server, *N, *parallel)
	dns.Id = origID
	doRun("math/rand", *server, *N, *parallel)
	dns.Id = cryptID
	doRun("crypto/rand", *server, *N, *parallel)
	dns.Id = origID
	doRun("math/rand", *server, *N, *parallel)
	dns.Id = cryptID
	doRun("crypto/rand", *server, *N, *parallel)
	dns.Id = origID
	doRun("math/rand", *server, *N, *parallel)
	dns.Id = cryptID
	doRun("crypto/rand", *server, *N, *parallel)
	dns.Id = origID
	doRun("math/rand", *server, *N, *parallel)
	dns.Id = cryptID
	doRun("crypto/rand", *server, *N, *parallel)
	dns.Id = origID
	doRun("math/rand", *server, *N, *parallel)
	dns.Id = cryptID
	doRun("crypto/rand", *server, *N, *parallel)
	dns.Id = origID
	doRun("math/rand", *server, *N, *parallel)
	dns.Id = cryptID
	doRun("crypto/rand", *server, *N, *parallel)
	dns.Id = origID
	doRun("math/rand", *server, *N, *parallel)
	dns.Id = cryptID
	doRun("crypto/rand", *server, *N, *parallel)
}

func doRun(desc string, server string, N, parallel int) {
	client := new(dns.Client)

	begin := time.Now()
	var wg sync.WaitGroup
	wg.Add(parallel)
	var j int32 = int32(N)
	for i := 0; i < parallel; i++ {
		go func() {
			for {
				remaining := atomic.AddInt32(&j, -1)
				if remaining <= 0 {
					wg.Done()
					return
				}
				m := new(dns.Msg)
				m.SetQuestion(dns.Fqdn("example.com."), dns.TypeA)
				_, _, err := client.Exchange(m, server)
				if err != nil {
					log.Print(err)
				}
			}
		}()
	}
	wg.Wait()
	qps := float64(N) / time.Since(begin).Seconds()
	log.Printf("%12s N = %d, parallel = %d, qps = %f", desc, N, parallel, qps)
}

func dnsServer() {
	err := dns.ListenAndServe(":1053", "udp", dns.HandlerFunc(func(w dns.ResponseWriter, r *dns.Msg) {
		m := new(dns.Msg)
		m.SetReply(r)
		m.Compress = false
		m.Rcode = dns.RcodeSuccess
		w.WriteMsg(m)
	}))
	if err != nil {
		log.Fatal(err)
	}
}