msoxzw/udp.py

## udp.py
import argparse
import asyncio
import struct

_I = struct.Struct('!I')


class UDPClient(asyncio.DatagramProtocol):

    def __init__(self, size, time):
        self.loop = asyncio.get_running_loop()
        self.size = size
        self.time = time

    def datagram_received(self, data, addr):
        self.size -= len(data)
        if not self.size:
            self.time.set_result(self.loop.time())


class UDPServer(asyncio.DatagramProtocol):

    def __init__(self, packet):
        self.packet = packet
        self.transport = None

    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data, addr):
        count, = _I.unpack(data)
        for _ in range(count):
            self.transport.sendto(self.packet, addr)


async def benchmark(count=65535, packet_size=1024, host='127.0.0.1', port=9999):
    loop = asyncio.get_running_loop()
    finish_time = loop.create_future()
    size = count * packet_size

    server, _ = await loop.create_datagram_endpoint(
        lambda: UDPServer(bytes(packet_size)),
        local_addr=(host, port),
    )
    client, _ = await loop.create_datagram_endpoint(
        lambda: UDPClient(size, finish_time),
        remote_addr=(host, port),
    )
    client.sendto(_I.pack(count), (host, port))
    start_time = loop.time()

    try:
        print('Transfer speed: {:.3f} Mbps'.format(
            size * 8 / 1000 / 1000 / (await finish_time - start_time)))
    finally:
        client.close()
        server.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('count', type=int, nargs='?', default=65536)
    args = parser.parse_args()
    asyncio.run(benchmark(args.count))
	import argparse
	import asyncio
	import struct

	_I = struct.Struct('!I')


	class UDPClient(asyncio.DatagramProtocol):

	def __init__(self, size, time):
	self.loop = asyncio.get_running_loop()
	self.size = size
	self.time = time

	def datagram_received(self, data, addr):
	self.size -= len(data)
	if not self.size:
	self.time.set_result(self.loop.time())


	class UDPServer(asyncio.DatagramProtocol):

	def __init__(self, packet):
	self.packet = packet
	self.transport = None

	def connection_made(self, transport):
	self.transport = transport

	def datagram_received(self, data, addr):
	count, = _I.unpack(data)
	for _ in range(count):
	self.transport.sendto(self.packet, addr)


	async def benchmark(count=65535, packet_size=1024, host='127.0.0.1', port=9999):
	loop = asyncio.get_running_loop()
	finish_time = loop.create_future()
	size = count * packet_size

	server, _ = await loop.create_datagram_endpoint(
	lambda: UDPServer(bytes(packet_size)),
	local_addr=(host, port),
	)
	client, _ = await loop.create_datagram_endpoint(
	lambda: UDPClient(size, finish_time),
	remote_addr=(host, port),
	)
	client.sendto(_I.pack(count), (host, port))
	start_time = loop.time()

	try:
	print('Transfer speed: {:.3f} Mbps'.format(
	size * 8 / 1000 / 1000 / (await finish_time - start_time)))
	finally:
	client.close()
	server.close()

	if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('count', type=int, nargs='?', default=65536)
	args = parser.parse_args()
	asyncio.run(benchmark(args.count))