Python 3.5 async await and asyncio

client.py

"""
client.py - AsyncIO Server using StreamReader and StreamWriter

This will create 200 client connections to a server running server.py

It will handshake and run similar to this:

Server: HELLO
Client: WORLD

Server: READY
Client: one
Server: ECHO 1: one

...

Client: six
Server: ECHO 6: six

Client: BYE
Server: BYE

"""

import asyncio
import logging

log = logging.getLogger(__name__)

clients = {}  # task -> (reader, writer)


def make_connection(host, port):

    task = asyncio.Task(handle_client(host, port))

    clients[task] = (host, port)

    def client_done(task):
        del clients[task]
        log.info("Client Task Finished")
        if len(clients) == 0:
            log.info("clients is empty, stopping loop.")
            loop = asyncio.get_event_loop()
            loop.stop()

    log.info("New Client Task")
    task.add_done_callback(client_done)


async def handle_client(host, port):
    log.info("Connecting to %s %d", host, port)
    client_reader, client_writer = await asyncio.open_connection(host,
                                                                      port)
    log.info("Connected to %s %d", host, port)
    try:
        # looking for a hello
        # give client a chance to respond, timeout after 10 seconds
        data = await asyncio.wait_for(client_reader.readline(),
                                           timeout=10.0)

        if data is None:
            log.warning("Expected HELLO, received None")
            return

        sdata = data.decode().rstrip().upper()
        log.info("Received %s", sdata)
        if sdata != "HELLO":
            log.warning("Expected HELLO, received '%s'", sdata)
            return

        # send back a WORLD
        client_writer.write("WORLD\n".encode())

        # wait for a READY
        data = await asyncio.wait_for(client_reader.readline(),
                                           timeout=10.0)

        if data is None:
            log.warning("Expected READY, received None")
            return

        sdata = data.decode().rstrip().upper()
        if sdata != "READY":
            log.warning("Expected READY, received '%s'", sdata)
            return

        echostrings = ['one', 'two', 'three', 'four', 'five', 'six']

        for echostring in echostrings:
            # send each string and get a reply, it should be an echo back
            client_writer.write(("%s\n" % echostring).encode())
            data = await asyncio.wait_for(client_reader.readline(),
                                               timeout=10.0)
            if data is None:
                log.warning("Echo received None")
                return

            sdata = data.decode().rstrip()
            log.info(sdata)

        # send BYE to disconnect gracefully
        client_writer.write("BYE\n".encode())

        # receive BYE confirmation
        data = await asyncio.wait_for(client_reader.readline(),
                                           timeout=10.0)

        sdata = data.decode().rstrip().upper()
        log.info("Received '%s'" % sdata)
    finally:
        log.info("Disconnecting from %s %d", host, port)
        client_writer.close()
        log.info("Disconnected from %s %d", host, port)


def main():
    log.info("MAIN begin")
    loop = asyncio.get_event_loop()
    for x in range(200):
        make_connection('localhost', 2991)
    loop.run_forever()
    log.info("MAIN end")

if __name__ == '__main__':
    log = logging.getLogger("")
    formatter = logging.Formatter("%(asctime)s %(levelname)s " +
                                  "[%(module)s:%(lineno)d] %(message)s")
    # setup console logging
    log.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)

    ch.setFormatter(formatter)
    log.addHandler(ch)
    main()

server.py

"""
server.py - AsyncIO Server using StreamReader and StreamWriter

example in another terminal:

    $ nc localhost 2991
    HELLO
    WORLD
    READY
    one
    ECHO 1: one
    two
    ECHO 2: two
    three
    ECHO 3: three
    four
    ECHO 4: four
    five
    ECHO 5: five
    six
    ECHO 6: six
    seven
    ECHO 7: seven
    eight
    ECHO 8: eight
    nine
    ECHO 9: nine
    ten
    ECHO 10: ten
    bye
    BYE

    $
"""

import asyncio
import logging

log = logging.getLogger(__name__)

clients = {}  # task -> (reader, writer)


def accept_client(client_reader, client_writer):
    task = asyncio.Task(handle_client(client_reader, client_writer))
    clients[task] = (client_reader, client_writer)

    def client_done(task):
        del clients[task]
        client_writer.close()
        log.info("End Connection")

    log.info("New Connection")
    task.add_done_callback(client_done)


async def handle_client(client_reader, client_writer):
    # send a hello to let the client know they are connected
    client_writer.write("HELLO\n".encode())

    # give client a chance to respond, timeout after 10 seconds
    data = await asyncio.wait_for(client_reader.readline(),
                                       timeout=10.0)

    if data is None:
        log.warning("Expected WORLD, received None")
        return

    sdata = data.decode().rstrip()
    log.info("Received %s", sdata)
    if sdata != "WORLD":
        log.warning("Expected WORLD, received '%s'", sdata)
        return

    # now be an echo back server until client sends a bye
    i = 0  # sequence number
    # let client know we are ready
    client_writer.write("READY\n".encode())
    while True:
        i = i + 1
        # wait for input from client
        data = await asyncio.wait_for(client_reader.readline(),
                                           timeout=10.0)
        if data is None:
            log.warning("Received no data")
            # exit echo loop and disconnect
            return

        sdata = data.decode().rstrip()
        if sdata.upper() == 'BYE':
            client_writer.write("BYE\n".encode())
            break
        response = ("ECHO %d: %s\n" % (i, sdata))
        client_writer.write(response.encode())


def main():
    loop = asyncio.get_event_loop()
    f = asyncio.start_server(accept_client, host=None, port=2991)
    loop.run_until_complete(f)
    loop.run_forever()

if __name__ == '__main__':
    log = logging.getLogger("")
    formatter = logging.Formatter("%(asctime)s %(levelname)s " +
                                  "[%(module)s:%(lineno)d] %(message)s")
    # setup console logging
    log.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)

    ch.setFormatter(formatter)
    log.addHandler(ch)
    main()