Skip to content
Advertisement

How can I get a socket’s .recv() not to block?

I’m trying to write a simple daemon that listens for orders on a Unix socket. The following works, but the connection.recv(1024) line blocks, meaning I can’t kill the server gracefully:

import socket, os

with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server:
    server.bind("/tmp/sock")
    server.listen()
    connection, __ = server.accept()
    with connection:
        while True:
            data = connection.recv(1024)
            print("Hi!")  # This line isn't executed 'til data is sent
            if data:
                print(data.decode())

Ideally, I’d like to place all of this inside a Thread that checks a self.should_stop property every self.LOOP_TIME seconds, and if that value is set to True, then exit. However, as that .recv() line blocks, there’s no way for my program to be doing anything other than waiting at any given time.

Surely there’s a proper way to do this, but as I’m new to sockets, I have no idea what that is.

Edit

Jeremy Friesner’s answer put me on the right track. I realised that I could allow the thread to block and simply set .should_stop then pass an b"" to the socket so that it’d un-block, see that it should stop, and then exit cleanly. Here’s the end result:

import os
import socket

from pathlib import Path
from shutil import rmtree
from threading import Thread


class MyThreadThing(Thread):

    RUNTIME_DIR = Path(os.getenv("XDG_RUNTIME_DIR", "/tmp")) / "my-project-name"

    def __init__(self):

        super().__init__(daemon=True)

        self.should_stop = False

        if self.RUNTIME_DIR.exists():
            rmtree(self.RUNTIME_DIR)
        self.RUNTIME_DIR.mkdir(0o700)

        self.socket_path = self.RUNTIME_DIR / "my-project.sock"

    def run(self) -> None:

        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:

            s.bind(self.socket_path.as_posix())
            s.listen()

            while True:

                connection, __ = s.accept()
                action = ""

                with connection:

                    while True:

                        received = connection.recv(1024).decode()
                        action += received
                        if not received:
                            break

                    # Handle whatever is in `action`

                if self.should_stop:
                    break

        self.socket_path.unlink()

    def stop(self):
        """
        Trigger this when you want to stop the listener.
        """
        self.should_stop = True
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
            s.connect(self.socket_path.as_posix())
            s.send(b"")

Advertisement

Answer

Using arbitrary-length timeouts is always a bit unsatisfactory — either you set the timeout-value to a relatively long time, in which case your program becomes slow to react to the quit-request, because it is pointlessly waiting for timeout period to expire… or you set the timeout-value to a relatively short time, in which case your program is constantly waking up to see if it should quit, wasting CPU power 24/7 to check for an event which might never arrive.

A more elegant way to deal with the problem is to create a pipe, and send a byte on the pipe when you want your event-loop to exit. Your event loop can simultaneously “watch” both the pipe’s reading-end file-descriptor and your networking-socket(s) via select(), and when that file-descriptor indicates it is ready-for-read, your event loop can respond by exiting. This approach is entirely event-driven, so it requires no CPU wakeups except when there is actually something to do.

Below is an example version of your program that implements a signal-handler for SIGINT (aka pressing Control-C) to sends the please-quit-now byte on the pipe:

import socket, os
import select
import signal, sys

# Any bytes written to (writePipeFD) will become available for reading on (readPipeFD)
readPipeFD, writePipeFD = os.pipe()

# Set up a signal-handler to handle SIGINT (aka Ctrl+C) events by writing a byte to the pipe
def signal_handler(sig, frame):
    print("signal_handler() is executing -- SIGINT detected!")
    os.write(writePipeFD, b"")  # doesn't matter what we write; a single 0-byte will do
signal.signal(signal.SIGINT, signal_handler)

with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as serverSock:
    serverSock.bind("/tmp/sock")
    serverSock.listen()

    # Wait for incoming connection (or the please-quit signal, whichever comes first)
    connection = None
    while True:
       readReady,writeReady,exceptReady = select.select([readPipeFD,serverSock], [], [])
       if (readPipeFD in readReady):
          print("accept-loop: Someone wrote a byte to the pipe; time to go away!");
          break
       if (connection in readReady):
          connection, __ = serverSock.accept()
          break

    # Read data from incoming connection (or the please-quit signal, whichever comes first)
    if connection:
       with connection:
          while True:
             readReady,writeReady,exceptReady = select.select([readPipeFD,connection], [], [])
             if (readPipeFD in readReady):
                print("Connection-loop: Someone wrote a byte to the pipe; time to go away!");
                break
             if (connection in readReady):
                data = connection.recv(1024)
                print("Hi!")  # This line isn't executed 'til data is sent
                if data:
                   print(data.decode())

    print("Bye!")
Advertisement