QUIC benchmark

quic.py

import argparse
import asyncio
import ssl
import struct

import aioquic.asyncio
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.events import StreamDataReceived

_I = struct.Struct('!I')


class QuicClient(aioquic.asyncio.QuicConnectionProtocol):

    def quic_event_received(self, quic_event):
        if isinstance(quic_event, StreamDataReceived):
            if quic_event.end_stream:
                self.time.set_result(self._loop.time() - self.start_time)

    def transfer(self, size):
        self._quic.send_stream_data(
            stream_id=self._quic.get_next_available_stream_id(),
            data=_I.pack(size),
            end_stream=True,
        )
        self.transmit()
        self.start_time = self._loop.time()
        self.time = self._loop.create_future()
        return self.time


class QuicServer(aioquic.asyncio.QuicConnectionProtocol):

    def quic_event_received(self, quic_event):
        if isinstance(quic_event, StreamDataReceived):
            size, = _I.unpack(quic_event.data)
            self._quic.send_stream_data(
                stream_id=quic_event.stream_id,
                data=b'\x00' * size,
                end_stream=True,
            )
            self.transmit()


async def benchmark(certfile, keyfile=None, password=None, host='::1', port=9999, size=1048576):
    configuration = QuicConfiguration(is_client=False)
    configuration.load_cert_chain(certfile=certfile, keyfile=keyfile, password=password)
    server = await aioquic.asyncio.serve(
        host, port,
        configuration=configuration,
        create_protocol=QuicServer,
    )

    async with aioquic.asyncio.connect(
        host, port,
        configuration=QuicConfiguration(verify_mode=ssl.CERT_NONE),
        create_protocol=QuicClient,
    ) as client:
        print('Transfer speed: {:.3f} Mbps'.format(
            size * 8 / 1000 / 1000 / await client.transfer(size))
        )

    server.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--certificate', required=True)
    parser.add_argument('-k', '--private-key')
    parser.add_argument('size', nargs='?', type=int, default=16, help='Transfer size (MiB)')
    args = parser.parse_args()
    asyncio.run(
        benchmark(
            certfile=args.certificate,
            keyfile=args.private_key,
            size=args.size * 1048576,
        )
    )