how to add timeout with asyncio UDP client

Question

my code is like this, it basically just copy the example code of UDP Echo Client in Python
offical documentation.

import asyncio
from asyncio import DatagramProtocol


class EchoClientProtocol(DatagramProtocol):
    def __init__(self, message, on_con_lost):
        self.message = message
        self.on_con_lost = on_con_lost
        self.transport = None

    def connection_made(self, transport):
        self.transport = transport
        print('Send:', self.message)
        self.transport.sendto(self.message.encode())

    def datagram_received(self, data, addr):
        print("Received:", data.decode())

        print("Close the socket")
        self.transport.close()

    def error_received(self, exc):
        print('Error received:', exc)

    def connection_lost(self, exc):
        print("Connection closed")
        self.on_con_lost.set_result(True)


async def main():

    loop = asyncio.get_running_loop()

    on_con_lost = loop.create_future()
    message = "Hello World!"
    transport, protocol = await loop.create_datagram_endpoint(
        lambda: EchoClientProtocol(message, on_con_lost),
        remote_addr=('127.0.0.1', 20001)
    )

    try:
        await on_con_lost
    finally:
        transport.close()


asyncio.run(main())

It works well.But when service is not available, client will be blocked. I tried add asyncio.wait_for to fix it, but it's not working.

transport, protocol = await asyncio.wait_for(loop.create_datagram_endpoint(
    lambda: EchoClientProtocol(message, on_con_lost),
    remote_addr=('127.0.0.1', 20001)
), 5)

I have a guess that maybe because I am not totally understanrd this usage of the on_con_lost parameter, How to do it?

Answer 1

Here's what I think is happening:

You are launching your datagram endpoint. You then have the launcher block on the disconnect future. For the disconnect future to trigger, the connection_lost function must be called. But if there was no connection in the first place, where connection_made isnt called, then that future will not become ready. What you really want is to add the timeout to awaiting the future. If you get a timeout expired exception then shutdown the connection. The way your current code works, you launch the endpoint and once that times out the code moves on to wait for the disconnect future which now cannot become ready.