如何在UDP的socket.recv()中正确使用KeyboardInterrupt？

我在Windows系统下，使用socket.socketpair()搭配signal.set_wakeup_fd()来实现UDP的socket.recv()与KeyboardInterrupt兼容，示例代码如下。我也了解另一种常用方案：配置带timeout的忙循环来定期检查程序是否被中断。但这两种方案都像是拙劣的权宜之计，请问是否存在“正确”的实现方式？

import select
import signal
import socket

def _udp_listen(address, family=socket.AF_INET, flags=0, sockopts=frozenset()):
    bufsize = getpagesize()
    with socket.socket(family, socket.SOCK_DGRAM) as sock:
        for sockopt in sockopts:
            sock.setsockopt(*sockopt)
        sock.bind(address)
        _coalmine, _canary = socket.socketpair()
        with _canary, _coalmine, _wakeup_fd_ctx(_coalmine.fileno(), strict=True, warn_on_full_buffer=False):
            while True:
                ready = select.select((sock, _canary), (), ())[0]
                if _canary in ready:
                    # There's no need to raise any error ourselves,
                    # since Python itself will raise KeyboardInterrupt
                    # out of select.select() if needed
                    pass
                if sock in ready:
                    yield sock.recv(bufsize, flags)


# -----

from contextlib import contextmanager

try:
    from resource import getpagesize
except ImportError:
    import mmap
    def getpagesize():
        return mmap.PAGESIZE


@contextmanager
def _wakeup_fd_ctx(fd, strict=True, **k):
    _orig_wakeup_fd = signal.set_wakeup_fd(fd, **k)
    _needs_restore = True
    try:
        if _orig_wakeup_fd == -1:
            yield fd
        else:
            # We overwrote the existing handler
            if strict:
                raise RuntimeError(f'wakeup fd already occupied ({_orig_wakeup_fd}). Not sure what to do about that.')
            else:
                signal.set_wakeup_fd(_orig_wakeup_fd)
                _needs_restore = False
                yield _orig_wakeup_fd
    finally:
        if _needs_restore:
            signal.set_wakeup_fd(_orig_wakeup_fd)

核心结论 #

在Windows平台上，你当前使用的socketpair+set_wakeup_fd方案其实已经是相对“正确”的实现方式——因为Windows的信号机制和Unix类系统差异很大，Python在Windows下对KeyboardInterrupt的处理本身就依赖唤醒阻塞的系统调用，而你的方案刚好贴合了这个逻辑。

为什么两种方案看起来像“权宜之计”？ #

这本质是Windows和Unix信号模型的差异导致的：

• 在Unix系统中，recv这类阻塞调用会被信号中断，Python可以直接捕获KeyboardInterrupt；但Windows没有真正的POSIX信号机制，Ctrl+C是通过控制台事件传递的，Python需要主动唤醒阻塞的系统调用才能触发异常。
• 带超时的忙循环方案确实不够优雅，因为它会引入不必要的轮询开销，而且超时时间很难平衡响应速度和资源消耗。

优化你的现有方案 #

你的实现已经很完善，但可以做几个小调整让它更健壮：

• 处理canary socket的读取：当前代码里只是判断_canary在ready列表里，但没有读取数据，时间长了可能会导致缓冲区溢出。可以在if _canary in ready:分支里加上_canary.recv(1)来清空缓冲区。
• 简化上下文管理：如果不需要严格检查已有wakeup fd的场景，可以简化_wakeup_fd_ctx的逻辑，减少不必要的分支。
• 兼容Windows的socketpair限制：Windows下socket.socketpair()在Python 3.5+才支持，且只能创建AF_INET的套接字，如果你需要兼容更早版本，可以用手动创建TCP套接字模拟socketpair的方式。

有没有更“正统”的替代方案？ #

在Windows上，其实没有比唤醒阻塞调用更直接的方式。另一个可选方案是使用线程：

• 把UDP接收逻辑放到单独的线程里，主线程等待用户输入或信号，收到Ctrl+C时通知子线程退出。
• 但线程方案会引入线程同步的复杂度，比如需要用事件对象来终止接收循环，反而不如你当前的方案简洁。

总结来说，你现在用的方案已经是Windows下处理这个问题的最优解之一，算不上“拙劣的权宜之计”——它是贴合Windows平台特性的合理实现。

内容的提问来源于stack exchange，提问作者JamesTheAwesomeDude