#lang racket/base

;; Simple stress-tester and performance measurement tool for DNS
;; implementations.

(require racket/udp)
(require racket/set)
(require "../racket-bitsyntax/main.rkt")
(require "api.rkt")
(require "codec.rkt")
(require "test-rrs.rkt")

(require racket/pretty)

(define latencies (make-vector 200 0))
(define latency-pos 0)
(define (record-latency-ms! ms)
  (vector-set! latencies latency-pos ms)
  (set! latency-pos (modulo (+ latency-pos 1) (vector-length latencies)))
  (when (zero? latency-pos)
    (for-each display (list ";; Mean latency "(/ (for/fold
						     ((sum 0))
						     ((i latencies))
						   (+ sum i))
						 (vector-length latencies))"ms\n"))))

(define (stress hostname port-number count rate)
  (define s (udp-open-socket #f #f))

  (define start-time (current-inexact-milliseconds))
  (let loop ((remaining count))
    (define request-message (dns-message (random 65536)
					 'request
					 'query
					 'non-authoritative
					 'not-truncated
					 'recursion-desired
					 'no-recursion-available
					 'no-error
					 (list (question (domain '(#"www" #"google" #"com"))
							 'a
							 'in
							 #f))
					 '()
					 '()
					 '()))

    (define now (current-inexact-milliseconds))
    (define sent-count (- count remaining))
    (define delta-ms (- now start-time))
    (define current-rate (/ sent-count (/ delta-ms 1000.0)))

    (when (> current-rate rate)
      (define target-delta-sec (/ sent-count rate))
      (sleep (- target-delta-sec (/ delta-ms 1000))))

    (when (zero? (modulo sent-count rate))
      (for-each display (list ";; Sent "sent-count" at target "rate"Hz, actual "
			      current-rate"Hz, in "delta-ms"ms\n")))

    (when (positive? remaining)
      (define sent-time (current-inexact-milliseconds))
      (udp-send-to s hostname port-number (dns-message->packet request-message))

      (define buffer (make-bytes 512))
      (define-values (packet-length source-hostname source-port)
	(udp-receive! s buffer))
      (define received-time (current-inexact-milliseconds))
      (define reply (packet->dns-message (sub-bit-string buffer 0 (* 8 packet-length))))

      ;;(pretty-print reply)
      (record-latency-ms! (- received-time sent-time))
      (loop (- remaining 1)))))

(stress "localhost" (test-port-number) 100000 500)