udp_test/udp_test_host.py

#!/usr/bin/python
"""UDP test host.

This is the host-side implementation of a UDP test tool to evaluate the speed
of a connection using bursts of UDP traffic
"""
import select
import socket
import struct
import sys
import time

import options

try:
  import monotime
except ImportError:
  pass
try:
  monotime = time.monotime
except AttributeError:
  monotime = time.time

optspec = """
udp_test_host.py [options...]
--
p,port=         Specify a port to use. [5001]
s,size=         Specify the packet size [1472]
"""


def main():
  o = options.Options(optspec)
  opt, _, _ = o.parse(sys.argv[1:])

  port = int(opt.port)
  send_rate = 1000.0  # Hz
  current_index = 0
  timeout_val = 10

  sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  sock.bind(('', port))

  client_addr = ('', 0)
  last_client_time = 0
  last_packets_sent = 0
  last_packets_recv = 0
  remainder_time = 0

  cycle_counter = 0
  cycle_time = monotime()
  msg_length = int(opt.size)
  counter_msg = bytearray('7' * msg_length)
  while True:
    # Account for processing time
    sock_status, _, _ = select.select([sock], [], [], remainder_time)
    if sock_status:
      host_time = monotime()
      data, addr = sock.recvfrom(65536)
      magic, pack_type, server_count, client_count, new_time = struct.unpack(
          '!4s4siid', data)
      if magic != 'TPDO' or pack_type != 'UPDT':
        print 'Invalid Packet'
        continue
      if addr != client_addr:
        current_index = 0
        send_rate = 1000.0  # Hz

        client_addr = addr
        cycle_time = monotime()
        cycle_counter = 0
        # Keeps track of last time client was heard from
        last_packets_sent = last_packets_recv = 0
        last_client_time = new_time
        host_time = monotime()
        print 'connection established with %s\n' % client_addr[0]
      else:
        # Calculate metrics for user
        cycle_length = new_time - last_client_time
        measured_pack_sent = (server_count - last_packets_sent) / cycle_length
        measured_send_rate = measured_pack_sent * msg_length * 8 / 1000000
        measured_recv_rate = ((client_count - last_packets_recv) /
                              cycle_length * msg_length * 8 / 1000000)
        if server_count <= last_packets_sent:
          success_rate = 0
        else:
          success_rate = (float(client_count - last_packets_recv) /
                          (server_count - last_packets_sent))

        print ('\nAttempted Rate: %.2fMbps, Sending Rate: %.2fMbps, '
               'Receive Rate: %.2fMbps, Success Rate: %.2f%%'
               % (send_rate * msg_length * 8 / 1000000,
                  measured_send_rate,
                  measured_recv_rate,
                  success_rate * 100))

        cycle_time = monotime()
        cycle_counter = 0
        if float(success_rate) > 0.9:
          send_rate = 1.1 * measured_pack_sent
        elif float(success_rate) < 0.8:
          send_rate = 0.9 * measured_pack_sent
        else:
          send_rate = measured_pack_sent

        # Prevent the host from not sending any packets.
        if send_rate <= 0:
          send_rate = 10
        last_packets_recv = client_count
        last_packets_sent = server_count
        last_client_time = new_time

    elif client_addr != ('', 0):
      # No packets from client received
      burst_counter = 0
      while True:
        remainder_time = (cycle_counter + 1) / send_rate - (monotime() -
                                                            cycle_time)
        if remainder_time > 0:
          break
        cycle_counter += 1
        current_index += 1
        counter_msg[0:12] = struct.pack('!4s4si', 'TPDO', 'RNNG', current_index)
        sock.sendto(counter_msg, client_addr)
        burst_counter += 1
        if not current_index % (int(send_rate/10) + 1):
          sys.stdout.write('.')
          sys.stdout.flush()
        if monotime() - host_time > timeout_val:
          client_addr = ('', 0)
          sys.stderr.write('\nconnection lost with client\n')
          remainder_time = 0
          break
        if burst_counter > 1000:
          remainder_time = 0
          break

if __name__ == '__main__':
  main()