简体   繁体   中英

Python 3 websockets - send message before closing connection

I'm new to Stack Overflow (although have been a long-term "stalker"!) so please be gentle with me!

I'm trying to learn Python, in particular Asyncio using websockets.

Having scoured the web for examples/tutorials I've put together the following tiny chat application, and could use some advice before it gets bulkier (more commands etc) and becomes difficult to refactor.

My main question, is why (when sending the DISCONNECT command) does it need the asyncio.sleep(0) in order to send the disconnection verification message BEFORE closing the connection?

Other than that, am I on the right tracks with the structure here?

I feel that there's too much async/await but I can't quite wrap my head around why.

Staring at tutorials and S/O posts for hours on end doesn't seem to be helping at this point so I thought I'd get some expert advice directly!

Here we go, simple WS server that responds to "nick", "msg", "test" & "disconnect" commands. No prefix required, ie "nick Rachel".

import asyncio
import websockets
import sys

class ChatServer:

    def __init__(self):
        print("Chat Server Starting..")
        self.Clients = set()
        if sys.platform == 'win32':
            self.loop = asyncio.ProactorEventLoop()
            asyncio.set_event_loop(self.loop)
        else:
            self.loop = asyncio.get_event_loop()

    def run(self):
        start_server = websockets.serve(self.listen, '0.0.0.0', 8080)
        try:
            self.loop.run_until_complete(start_server)
            print("Chat Server Running!")
            self.loop.run_forever()
        except:
            print("Chat Server Error!")

    async def listen(self, websocket, path):

        client = Client(websocket=websocket)
        sender_task = asyncio.ensure_future(self.handle_outgoing_queue(client))

        self.Clients.add(client)
        print("+ connection: " + str(len(self.Clients)))

        while True:
            try:
                msg = await websocket.recv()
                if msg is None:
                    break

                await self.handle_message(client, msg)

            except websockets.exceptions.ConnectionClosed:
                break

        self.Clients.remove(client)
        print("- connection: " + str(len(self.Clients)))

    async def handle_outgoing_queue(self, client):
        while client.websocket.open:
            msg = await client.outbox.get()
            await client.websocket.send(msg)


    async def handle_message(self, client, data):

        strdata = data.split(" ")
        _cmd = strdata[0].lower()

        try:
            # Check to see if the command exists. Otherwise, AttributeError is thrown.
            func = getattr(self, "cmd_" + _cmd)

            try:
                await func(client, param, strdata)
            except IndexError:
                await client.send("Not enough parameters!")

        except AttributeError:
            await client.send("Command '%s' does not exist!" % (_cmd))

    # SERVER COMMANDS

    async def cmd_nick(self, client, param, strdata):
        # This command needs a parameter (with at least one character). If not supplied, IndexError is raised
        # Is there a cleaner way of doing this? Otherwise it'll need to reside within all functions that require a param
        test = param[1][0]


        # If we've reached this point there's definitely a parameter supplied
        client.Nick = param[1]
        await client.send("Your nickname is now %s" % (client.Nick))

    async def cmd_msg(self, client, param, strdata):
        # This command needs a parameter (with at least one character). If not supplied, IndexError is raised
        # Is there a cleaner way of doing this? Otherwise it'll need to reside within all functions that require a param
        test = param[1][0]

        # If we've reached this point there's definitely a parameter supplied
        message = strdata.split(" ",1)[1]

        # Before we proceed, do we have a nickname?
        if client.Nick == None:
            await client.send("You must choose a nickname before sending messages!")
            return

        for each in self.Clients:
            await each.send("%s says: %s" % (client.Nick, message))

    async def cmd_test(self, client, param, strdata):
        # This command doesn't need a parameter, so simply let the client know they issued this command successfully.
        await client.send("Test command reply!")

    async def cmd_disconnect(self, client, param, strdata):
        # This command doesn't need a parameter, so simply let the client know they issued this command successfully.
        await client.send("DISCONNECTING")
        await asyncio.sleep(0)      # If this isn't here we don't receive the "disconnecting" message - just an exception in "handle_outgoing_queue" ?
        await client.websocket.close()


class Client():
    def __init__(self, websocket=None):
        self.websocket = websocket
        self.IPAddress = websocket.remote_address[0]
        self.Port = websocket.remote_address[1]

        self.Nick = None
        self.outbox = asyncio.Queue()

    async def send(self, data):
        await self.outbox.put(data)

chat = ChatServer()
chat.run()

Your code uses infinite size Queues , which means .put() calls .put_nowait() and returns immediately. (If you do want to keep these queues in your code, consider using 'None' in the queue as a signal to close a connection and move client.websocket.close() to handle_outgoing_queue() ).

Another issue: Consider replacing for x in seq: await co(x) with await asyncio.wait([co(x) for x in seq]) . Try it with asyncio.sleep(1) to experience a dramatic difference.

I believe a better option will be dropping all outbox Queue s and just relay on the built in asyncio queue and ensure_future . The websockets package already includes Queues in its implementation.

I want to point out that the author of websockets indicated in a post on July 17 of 2017 that websockets used to return None when the connection was closed but that was changed at some point. Instead he suggests that you use a try and deal with the exception. The OP's code shows both a check for None AND a try/except. The None check is needlessly verbose and apparently not even accurate since with the current version, websocket.recv() doesn't return anything when the client closes.

Addressing the "main" question, it looks like a race condition of sorts. Remember that asyncio does it's work by going around and touching all the awaited elements in order to nudge them along. If your 'close connection' command is processed at some point ahead of when your queue is cleared, the client will never get that last message in the queue. Adding the async.sleep adds an extra step to the round robin and probably puts your queue emptying task ahead of your 'close connection'.

Addressing the amount of awaits, it's all about how many asynchronous things you need to have happen to accomplish the goal. If you block at any point you'll stop all the other tasks that you want to keep going.

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM