5
\$\begingroup\$

I wrote a server-client application in Python where the server sends a never-ending stream of messages to connecting clients.

I was in need for such an architecture since I want to retrieve and process data from a public streaming endpoint, however the company limits the number of active connections to one per user.

With this setup, I can start my server ONCE and then forward the data to any number of connecting clients. This allows me to process data in different ways simultaneously (e.g. one client might write items to a DB, while another calculates some metrics etc.). Further, it allows me to develop a new module and connect it once it is ready for deployment, without interfering with my production modules.


Some details:

I send individual variable-length items by prefixing them with their length and then read this number of bytes of on the client side.

To avoid that a buggy client brings down the server, I have a capped message queue for each client and the client gets disconnected once the queue reaches a defined limit.


It's my first time writing network code in Python (and in general), so I'm looking for ANY advice or suggestions.

Thanks!


Server:

import socket
import struct
import time

from threading import Thread
from queue import Queue

MAX_CLIENTS = 16
MAX_CLIENT_QSIZE = 600

HEADER_FORMAT = struct.Struct('!I')


class _ClientHandler(Thread):

    def __init__(self, sock, address, server):
        Thread.__init__(self, daemon = True)

        self.sock = sock
        self.address = address
        self.server = server # Reference to server

        self.queue = Queue() # Message queue

    def _send_item(self, item):
        """Send item to client. Blocks thread."""

        self.sock.sendall(HEADER_FORMAT.pack(len(item)))
        self.sock.sendall(item.encode('utf-8'))

    def run(self):
        """Pop items from queue and sent to client."""

        while True:
            try:
                item = self.queue.get()
                self._send_item(item)
            except BrokenPipeError:
                self.server.delete_client(self.ident)
            except OSError:
                self.server.delete_client(self.ident)
            except:
                raise

class StreamServer(Thread):

    def __init__(self, address):
        Thread.__init__(self, daemon = True)

        self.address = address
        self._clients = dict()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """Disconnect clients and remove client threads when leaving context."""

        for ident in self._clients.keys():
            self.delete_client(ident)

    def queue_item(self, item):
        """Push item to every client queue."""

        for ident, client in self._clients.items():
            if client.queue.qsize() >= MAX_CLIENT_QSIZE:
                print('Client', client.sock.getsockname(), 'reached maximum qsize')
                self.delete_client(ident)
            client.queue.put(item)

    def get_qsizes(self):
        """Get list of client queue sizes."""

        return [client.queue.qsize() for _, client in self._clients.items()]

    def delete_client(self, ident):
        """Close client socket and delete client thread."""

        self._clients[ident].sock.close()
        del self._clients[ident] # Delete only client thread reference (garbage collection?)

    def run(self):
        """Listen for incomming client connections."""

        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

            sock.bind(self.address)
            sock.listen(4)

            print('Listening at', sock.getsockname())

            while True:
                try:
                    # Refuse connections once client limit reached
                    while len(self._clients) >= MAX_CLIENTS:
                        time.sleep(60)

                    client_sock, client_addr = sock.accept()
                    client_sock.shutdown(socket.SHUT_RD)
                    client = _ClientHandler(client_sock, client_addr, self)

                    print('Connection from', client_sock.getsockname())

                    client.start()
                    self._clients[client.ident] = client
                except:
                    # ...
                    # What exceptions to expect?
                    raise


if __name__ == '__main__':
    with StreamServer(('localhost', 1060)) as server:
        server.start() # Start server thread

        last_time = time.time()
        for item in range(1,1000):
            server.queue_item(str(item))
            time.sleep(1)

            if time.time() - last_time > 10.0:
                print('Client queue sizes:', server.get_qsizes())
                last_time = time.time()

Client:

import socket
import struct

HEADER_FORMAT = struct.Struct('!I')

class StreamClient():

    def __init__(self, address):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        self.sock.connect(address)
        self.sock.shutdown(socket.SHUT_WR)
        print('Connected to', self.sock.getsockname())

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.sock.close()

    def _recvall(self, size):
        data = b''
        while size:
            block = self.sock.recv(size)

            if not block:
                raise EOFError('Socket closed during message transmission')

            data += block
            size -= len(block)
        return data

    def __iter__(self):
        while True:
            try:
                (item_size,) = HEADER_FORMAT.unpack(self._recvall(HEADER_FORMAT.size))
                yield self._recvall(item_size).decode('utf-8')
            except EOFError:
                return
            except:
                # ...
                # What exceptions to expect?
                raise


if __name__ == '__main__':
    with StreamClient(('localhost', 1060)) as stream:
        for item in stream:
            print(item)
\$\endgroup\$
2
+50
\$\begingroup\$

This looks very reasonable, but some improvements are possible:

getsockname() returns the local side's address

In your code you write:

client_sock, client_addr = sock.accept()
...
print('Connection from', client_sock.getsockname())

That should have been client_sock.getpeername(). But you don't need to call any extra function, since the address of the peer is already stored in client_addr, so just write:

print('Connection from', client_addr)

Avoid transcoding between bytes and strings

It looks like you call encode('utf-8') in the server, and decode('utf-8') in the client. This sounds a bit redundant. Especially if, as you say, you are receiving data from a socket to a public endpoint, so I assume this already is in a bytes array initially.

Don't add sleep() calls to your code

I see this often in code: adding a sleep() here and there to patch around some issue. The problem is that now your code is suddenly becoming unresponsive for the time you are sleeping. And 60 seconds is a long time. What if within that 60 seconds, all clients drop their connections? You could accept lots of new connections during that time, but you can't because you have to wait until the end of that minute. Either just close the connection if you didn't want to accept it, or wait indefinitely, but have some way for a _ClientHandler to wake up the StreamServer as soon as it notices the peer closed its connection.

Add support for IPv6

You really should support IPv6 nowadays. It's not hard at all. Depending on which platform you are running on, you might get away with just using AF_UNSPEC instead of AF_INET when creating a listening socket, and it will accept connections from either version of IP. Otherwise, create two listening sockets, one for AF_INET and one for AF_INET6.

Add graceful shutdown of all threads

Add a way to shut down threads in a graceful way. To shut down a _ClientHandler, add a special item to the queue that signals that this handler should return from run(). For the StreamServer(), add a flag named running for example that is True by default. In run(), check the flag each iteration of the loop. If it's no longer True, exit the loop. To shutdown in an orderly way, you have to set self.running = False, then create a connection to the listening socket in order to have sock.accept() return.

Don't attempt to just close client sockets or the listening socket. There is no guarantee that this will cause the threads to immediately receive an error.

Consider using non-blocking sockets

Instead of having a queue per client, you could instead make use of the TCP send buffer for each client connection. So instead of first enquing some data, you immediately send copies of the data on each client's socket. Now the problem is to know when the send buffer is full. You can do that by making the socket non-blocking using socket.setblocking(False). When you attempt to write to a non-blocking socket, you will receive an error. If that happens, you can just close the connection.

You can tune the size of the sendbuffer in this way:

socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, size)

Since you no longer need queues, you also don't need one Thread per client anymore.

\$\endgroup\$
  • \$\begingroup\$ Thank you, these are some very helpful remarks. Is there anything more I can do to improve stability? \$\endgroup\$ – bi_scholar Aug 13 at 15:27
  • 1
    \$\begingroup\$ @bi_scholar It looks stable enough to me. But if you have a lot of clients, then at some point you will reach a limit to how many clients you can handle simultaneously. The code will drop client connections when that happens, which is fine in a way (it doesn't crash), but if you need it to be able to handle more clients, then I see two solutions: port the code to C, and/or have one primary StreamServer and several secondary StreamServers that each connect to the primary one. This should provide virtually limitless scalability. \$\endgroup\$ – G. Sliepen Aug 13 at 17:46
  • \$\begingroup\$ That makes sense, thanks! \$\endgroup\$ – bi_scholar Aug 13 at 18:53

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.