Python AF_UNIX socket / socketserver example (datagram and streaming mode)

example_dgram.py

"""
LICENSE: MIT

Example of using unix socket in SOCK_DGRAM mode.

"""

import socket
import socketserver
import time
import threading
import pprint
import os
pprint = pprint.PrettyPrinter().pprint

CLNT_SOCKET_ADR_PREFIX = '/tmp/test.dgram.clientsocket'
SRV_SOCKET_ADR = '/tmp/test.dgram.serversocket'


def server():

    class MyHandler(socketserver.DatagramRequestHandler):

        def handle_message(self, message: str) -> str:
            time.sleep(1)
            return f"Pong to '{message}' by server thread {threading.get_ident()}"

        def handle(self):
            request = self.request[0].decode('utf-8')
            print(f"[server] Got {len(self.request[0])} bytes "
                  f"from {self.client_address}: '''{request}'''")
            response = self.handle_message(request)
            self.wfile.write(response.encode('utf-8'))

    if False:  # Change to switch enable/disable threading in the server part
        class MyUnixDatagramServer(socketserver.UnixDatagramServer):
            pass
    else:
        class MyUnixDatagramServer(socketserver.ThreadingMixIn,
                                   socketserver.UnixDatagramServer):
            pass

    if os.path.exists(SRV_SOCKET_ADR):
        os.unlink(SRV_SOCKET_ADR)

    with MyUnixDatagramServer(SRV_SOCKET_ADR, MyHandler) as s:
        s.serve_forever()


def client(i):
    client_socket_address = f"{CLNT_SOCKET_ADR_PREFIX}" \
                            f".{os.getpid()}.{threading.get_ident()}"
    try:
        client = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
        client.bind(client_socket_address)

        def send_and_rcv(msg):
            client.sendto(msg.encode('utf-8'), SRV_SOCKET_ADR)
            response, address = client.recvfrom(1024)
            response = response.decode('utf-8')
            print(f"[client] Got {len(response)} bytes "
                  f"from {str(address)}: '''{response}'''", flush=True)
            return response

        send_and_rcv(f"Ping A{i}")
        send_and_rcv(f"Ping B{i}")
        # as long as there is ping pong between server and client,
        # there is no need to tokenize the message/response tuples

    finally:
        client.close()
        if os.path.exists(client_socket_address):
            os.unlink(client_socket_address)


def main():
    server_thread = threading.Thread(target=server, daemon=True)
    server_thread.start()
    time.sleep(0.2)

    clients = []
    for i in range(0, 10):
        client_thread = threading.Thread(target=client, kwargs={'i': i})
        client_thread.start()
        clients.append(client_thread)

    for c in clients:
        c.join(timeout=10)

    print(f"[main] any client is_alive = {any(c.is_alive() for c in clients)}")
    exit()


if __name__ == "__main__":
    main()

example_stream.py

"""
LICENSE: MIT

Example of using unix socket in SOCK_STREAM mode.

"""

import socket
import socketserver
import time
import threading
import pprint
import struct
import os
pprint = pprint.PrettyPrinter().pprint

SRV_SOCKET_ADR = '/tmp/test.stream.serversocket'


class EOF(Exception):
    pass


def recv_exact(connection, remaining):
    buf = bytearray(remaining)
    view = memoryview(buf)
    while remaining > 0:
        num_read = connection.recv_into(view, remaining)
        view = view[num_read:]
        remaining -= num_read
        if num_read == 0:  # eof
            if remaining == 0:
                break  # return buf
            elif remaining < len(buf):
                print("[server] connection closed mid send")
            raise EOF

    return buf


# Header: data length as 8 bytes in big endian
HEADER_FORMAT = ">Q"
HEADER_SIZE = 8


def send_str_packet(connection, msg):
    msg = msg.encode('utf-8')
    msg_len_packed = struct.pack(HEADER_FORMAT, len(msg))
    connection.sendall(msg_len_packed)
    connection.sendall(msg)


def recv_str_packet(connection):
    """ May raise EOF """
    msg_len_packed = recv_exact(connection, HEADER_SIZE)
    msg_len = struct.unpack(HEADER_FORMAT, msg_len_packed)[0]
    msg = recv_exact(connection, msg_len)
    return msg.decode('utf-8')


def server():

    class MyHandler(socketserver.StreamRequestHandler):

        def handle_message(self, message: str) -> str:
            time.sleep(1)
            return f"Pong to '{message}'" \
                   f" by server thread {threading.get_ident()}"

        def handle(self):
            try:
                while True:
                    request = recv_str_packet(self.request)
                    print(f"[server] Got {len(request)} bytes: "
                          f"'''{request}'''")
                    response = self.handle_message(request)

                    send_str_packet(self.request, response)
            except EOF:
                return
            finally:
                pass

    if False:  # Change to switch enable/disable threading in the server part
        class MyUnixStreamServer(socketserver.UnixStreamServer):
            pass
    else:
        class MyUnixStreamServer(socketserver.ThreadingMixIn,
                                 socketserver.UnixStreamServer):
            pass

    if os.path.exists(SRV_SOCKET_ADR):
        os.unlink(SRV_SOCKET_ADR)

    with MyUnixStreamServer(SRV_SOCKET_ADR, MyHandler) as s:
        s.serve_forever()


def client(i):
    client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    client.connect(SRV_SOCKET_ADR)

    def send_and_rcv(msg):
        send_str_packet(client, msg)
        response = recv_str_packet(client)
        print(f"[client] Got {len(response)} bytes: '''{response}'''")
        return response

    send_and_rcv(f"Ping A{i}")
    send_and_rcv(f"Ping B{i}")

    client.close()


def main():
    server_thread = threading.Thread(target=server, daemon=True)
    server_thread.start()
    time.sleep(0.2)

    clients = []
    for i in range(0, 10):
        client_thread = threading.Thread(target=client, kwargs={'i': i})
        client_thread.start()
        clients.append(client_thread)

    for c in clients:
        c.join(timeout=10)

    print(f"[main] any client is_alive = {any(c.is_alive() for c in clients)}")
    exit()


if __name__ == "__main__":
    main()

output_dgram

[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068575282944: '''Ping A0'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068566890240: '''Ping A1'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068558497536: '''Ping A2'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068345870080: '''Ping A3'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068329084672: '''Ping A4'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068320691968: '''Ping A5'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068303906560: '''Ping A6'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068295513856: '''Ping A7'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140067725104896: '''Ping A9'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140067733497600: '''Ping A8'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A0' by server thread 140068337477376'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068575282944: '''Ping B0'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A1' by server thread 140068312299264'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068566890240: '''Ping B1'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A2' by server thread 140067741890304'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068558497536: '''Ping B2'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A5' by server thread 140067699926784'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A4' by server thread 140067708319488'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A3' by server thread 140067716712192'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068329084672: '''Ping B4'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A7' by server thread 140067205019392'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A6' by server thread 140067691534080'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068320691968: '''Ping B5'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A9' by server thread 140067196626688'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping A8' by server thread 140067188233984'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068295513856: '''Ping B7'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068303906560: '''Ping B6'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140067733497600: '''Ping B8'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140068345870080: '''Ping B3'''
[server] Got 7 bytes from /tmp/test.dgram.clientsocket.1164023.140067725104896: '''Ping B9'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B0' by server thread 140067179841280'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B1' by server thread 140068312299264'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B2' by server thread 140067741890304'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B5' by server thread 140067691534080'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B7' by server thread 140067188233984'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B4' by server thread 140067699926784'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B9' by server thread 140068337477376'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B6' by server thread 140067205019392'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B3' by server thread 140067196626688'''
[client] Got 50 bytes from /tmp/test.dgram.serversocket: '''Pong to 'Ping B8' by server thread 140067708319488'''
[main] any client is_alive = False

output_stream

[server] Got 7 bytes: '''Ping A0'''
[server] Got 7 bytes: '''Ping A1'''
[server] Got 7 bytes: '''Ping A2'''
[server] Got 7 bytes: '''Ping A3'''
[server] Got 7 bytes: '''Ping A4'''
[server] Got 7 bytes: '''Ping A5'''
[server] Got 7 bytes: '''Ping A6'''
[server] Got 7 bytes: '''Ping A7'''
[server] Got 7 bytes: '''Ping A8'''
[server] Got 7 bytes: '''Ping A9'''
[client] Got 50 bytes: '''Pong to 'Ping A0' by server thread 140615371671296'''
[server] Got 7 bytes: '''Ping B0'''
[client] Got 50 bytes: '''Pong to 'Ping A1' by server thread 140615014676224'''
[client] Got 50 bytes: '''Pong to 'Ping A2' by server thread 140614997890816'''
[server] Got 7 bytes: '''Ping B1'''
[server] Got 7 bytes: '''Ping B2'''
[client] Got 50 bytes: '''Pong to 'Ping A3' by server thread 140614981105408'''
[client] Got 50 bytes: '''Pong to 'Ping A4' by server thread 140614544914176'''
[server] Got 7 bytes: '''Ping B3'''
[server] Got 7 bytes: '''Ping B4'''
[client] Got 50 bytes: '''Pong to 'Ping A5' by server thread 140614536521472'''
[client] Got 50 bytes: '''Pong to 'Ping A6' by server thread 140614528128768'''
[server] Got 7 bytes: '''Ping B5'''
[server] Got 7 bytes: '''Ping B6'''
[client] Got 50 bytes: '''Pong to 'Ping A8' by server thread 140614494557952'''
[client] Got 50 bytes: '''Pong to 'Ping A7' by server thread 140614502950656'''
[server] Got 7 bytes: '''Ping B8'''
[server] Got 7 bytes: '''Ping B7'''
[client] Got 50 bytes: '''Pong to 'Ping A9' by server thread 140613940934400'''
[server] Got 7 bytes: '''Ping B9'''
[client] Got 50 bytes: '''Pong to 'Ping B0' by server thread 140615371671296'''
[client] Got 50 bytes: '''Pong to 'Ping B2' by server thread 140614997890816'''
[client] Got 50 bytes: '''Pong to 'Ping B1' by server thread 140615014676224'''
[client] Got 50 bytes: '''Pong to 'Ping B4' by server thread 140614544914176'''
[client] Got 50 bytes: '''Pong to 'Ping B3' by server thread 140614981105408'''
[client] Got 50 bytes: '''Pong to 'Ping B6' by server thread 140614528128768'''
[client] Got 50 bytes: '''Pong to 'Ping B5' by server thread 140614536521472'''
[client] Got 50 bytes: '''Pong to 'Ping B7' by server thread 140614502950656'''
[client] Got 50 bytes: '''Pong to 'Ping B8' by server thread 140614494557952'''
[client] Got 50 bytes: '''Pong to 'Ping B9' by server thread 140613940934400'''
[main] any client is_alive = False