使用 asyncio 的简单 Python TCP 分叉服务器

Question

我想要做什么

我正在尝试模拟以下简单 socat(1) 命令的行为：

socat tcp-listen:SOME_PORT,fork,reuseaddr exec:'SOME_PROGRAM'

上述命令创建了一个分叉 TCP 服务器，它为每个连接分叉并执行 SOME_PROGRAM，将所述命令的 stdin 和 stdout 重定向到 TCP 套接字。

这是我想要实现的目标：

1. 使用asyncio 创建一个简单的 TCP 服务器来处理多个并发连接。
2. 每当收到连接时，将SOME_PROGRAM 作为子进程启动。
3. 将从套接字接收到的任何数据传递到SOME_PROGRAM 的标准输入。
4. 将从SOME_PROGRAM 的标准输出接收到的任何数据传递到套接字。
5. SOME_PROGRAM 退出时，将告别消息连同退出代码一起写入套接字并关闭连接。

我想在纯 Python 中执行此操作，而不使用使用 asyncio 模块的外部库。

到目前为止我所拥有的

这是我目前写的代码：

import asyncio

class ServerProtocol(asyncio.Protocol):
    def connection_made(self, transport):
        self.client_addr   = transport.get_extra_info('peername')
        self.transport     = transport
        self.child_process = None

        print('Connection with {} enstablished'.format(self.client_addr))

        asyncio.ensure_future(self._create_subprocess())

    def connection_lost(self, exception):
        print('Connection with {} closed.'.format(self.client_addr))

        if self.child_process.returncode is not None:
            self.child_process.terminate()

    def data_received(self, data):
        print('Data received: {!r}'.format(data))

        # Make sure the process has been spawned
        # Does this even make sense? Looks so awkward to me...
        while self.child_process is None:
            continue

        # Write any received data to child_process' stdin
        self.child_process.stdin.write(data)

    async def _create_subprocess(self):
        self.child_process = await asyncio.create_subprocess_exec(
            *TARGET_PROGRAM,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE
        )

        # Start reading child stdout
        asyncio.ensure_future(self._pipe_child_stdout())

        # Ideally I would register some callback here so that when
        # child_process exits I can write to the socket a goodbye
        # message and close the connection, but I don't know how
        # I could do that...

    async def _pipe_child_stdout(self):
        # This does not seem to work, this function returns b'', that is an
        # empty buffer, AFTER the process exits...
        data = await self.child_process.stdout.read(100) # Arbitrary buffer size

        print('Child process data: {!r}'.format(data))

        if data:
            # Send to socket
            self.transport.write(data)
            # Reschedule to read more data
            asyncio.ensure_future(self._pipe_child_stdout())


SERVER_PORT    = 6666
TARGET_PROGRAM = ['./test']

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    coro = loop.create_server(ServerProtocol, '0.0.0.0', SERVER_PORT)
    server = loop.run_until_complete(coro)

    print('Serving on {}'.format(server.sockets[0].getsockname()))

    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass

    server.close()
    loop.run_until_complete(server.wait_closed())
    loop.close()

还有我试图作为子进程运行的./test 程序：

#!/usr/bin/env python3

import sys

if sys.stdin.read(2) == 'a\n':
    sys.stdout.write('Good!\n')
else:
    sys.exit(1)

if sys.stdin.read(2) == 'b\n':
    sys.stdout.write('Wonderful!\n')
else:
    sys.exit(1)

sys.exit(0)

不幸的是，上面的代码并没有真正起作用，我有点不知道接下来要尝试什么。

按预期工作：

• 子进程已正确生成，并且似乎也正确接收了来自套接字的输入，因为我可以从htop 看到它，而且我也可以看到，只要我发送b\n 它就会终止。李>

什么不符合预期：

基本上其他的...

• 子进程的输出永远不会发送到套接字，实际上根本就不会读取。调用await self.child_process.stdout.read(100) 似乎永远不会终止：相反，它只会在子进程死亡之后终止，结果只是b''（一个空的bytes 对象）。
• 我无法理解子进程何时终止：如上所述，我想在发生这种情况时向套接字发送“再见”消息以及self.child_process.returncode，但我不知道如何以有意义的方式做到这一点。

我尝试了什么：

• 我尝试使用asyncio.loop.subprocess_exec() 而不是asyncio.create_subprocess_exec() 创建子进程。这解决了知道进程何时终止的问题，因为我可以实例化asyncio.SubprocessProtocol 的子类并使用它的process_exited() 方法，但是根本没有帮助我，因为如果我这样做了这样一来，我就不想再与流程'stdin 或stdout 交谈了！也就是说，我没有要与之交互的 Process 对象...
• 我尝试过使用asyncio.loop.connect_write_pipe() 和loop.connect_read_pipe()，但没有成功。

问题

那么，有人可以帮我弄清楚我做错了什么吗？必须有一种方法可以使这项工作顺利进行。当我第一次开始时，我正在寻找一种方法来轻松地使用一些管道重定向，但我不知道这是否可能在这一点上。是吗？看起来应该是。

我可以在 15 分钟内使用 fork()、exec() 和 dup2() 用 C 语言编写这个程序，所以我一定缺少一些东西！任何帮助表示赞赏。

Answer 1

您的代码有两个直接的实现问题：

• 服务器在将接收到的数据传输到子进程之前去除空格。这将删除尾随的换行符，因此如果 TCP 客户端发送"a\n"，子进程将只收到"a"。这样子进程永远不会遇到预期的"a\n" 字符串，并且它总是在读取两个字节后终止。这解释了来自子流程的空字符串 (EOF)。（在对该问题的后续编辑中删除了剥离。）
• 子进程不会刷新其输出，因此服务器不会收到任何写入。只有在子进程退出或填满其输出缓冲区（以千字节为单位）时，才会看到写入操作，并且在显示简短的调试消息时需要一段时间才能填满。

另一个问题是在设计层面。如 cmets 中所述，除非您的明确意图是实现新的 asyncio 协议，否则 recommended 坚持使用更高级别的 stream-based API，在本例中为 start_server 函数。 SubprocessProtocol、connect_write_pipe 和 connect_read_pipe 等更低级别的功能也不是您希望在应用程序代码中使用的东西。此答案的其余部分假设基于流的实现。

start_server 接受一个协程，该协程将在客户端连接时作为新任务生成。它使用两个异步流参数调用，一个用于读取，一个用于写入。协程包含与客户端通信的逻辑；在你的情况下，它会产生子进程并在它和客户端之间传输数据。

请注意，套接字和子进程之间的双向数据传输无法通过读取后写入的简单循环来实现。例如，考虑这个循环：

# INCORRECT: can deadlock (and also doesn't detect EOF)
child = await asyncio.create_subprocess_exec(...)
while True:
    proc_data = await child.stdout.read(1024)  # (1)
    sock_writer.write(proc_data)
    sock_data = await sock_reader.read(1024)
    child.stdin.write(sock_data)               # (2)

这种循环容易出现死锁。如果子进程正在响应它从 TCP 客户端接收到的数据，它有时只会在接收到一些输入后才提供输出。这将无限期地阻塞 (1) 处的循环，因为它只能在向孩子发送 sock_data 之后才能从孩子的 stdout 获取数据，这将在稍后的 (2) 处发生。实际上，（1）等待（2），反之亦然，构成死锁。请注意，颠倒传输顺序将无济于事，因为如果 TCP 客户端正在处理服务器子进程的输出，则循环将死锁。

使用 asyncio 可以轻松避免这种死锁：只需并行生成两个协程，一个将数据从套接字传输到子进程的标准输入，另一个从子进程的标准输入传输数据标准输出到套接字。例如：

# correct: deadlock-free (and detects EOF)
async def _transfer(src, dest):
    while True:
        data = await src.read(1024)
        if data == b'':
            break
        dest.write(data)

child = await asyncio.create_subprocess_exec(...)
loop.create_task(_transfer(child.stdout, sock_writer))
loop.create_task(_transfer(sock_reader, child.stdin))
await child.wait()

此设置与第一个while 循环之间的区别在于，这两个传输相互独立。死锁不会发生，因为从套接字读取永远不会等待从子进程读取，反之亦然。

应用于问题，整个服务器将如下所示：

import asyncio

class ProcServer:
    async def _transfer(self, src, dest):
        while True:
            data = await src.read(1024)
            if data == b'':
                break
            dest.write(data)

    async def _handle_client(self, r, w):
        loop = asyncio.get_event_loop()
        print(f'Connection from {w.get_extra_info("peername")}')
        child = await asyncio.create_subprocess_exec(
            *TARGET_PROGRAM, stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE)
        sock_to_child = loop.create_task(self._transfer(r, child.stdin))
        child_to_sock = loop.create_task(self._transfer(child.stdout, w))
        await child.wait()
        sock_to_child.cancel()
        child_to_sock.cancel()
        w.write(b'Process exited with status %d\n' % child.returncode)
        w.close()

    async def start_serving(self):
        await asyncio.start_server(self._handle_client,
                                   '0.0.0.0', SERVER_PORT)

SERVER_PORT    = 6666
TARGET_PROGRAM = ['./test']

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    server = ProcServer()
    loop.run_until_complete(server.start_serving())
    loop.run_forever()

随附的test 程序也必须修改为在每个sys.stdout.write() 之后调用sys.stdout.flush()，否则消息会停留在其stdio 缓冲区中，而不是发送到父级。

当我第一次开始时，我正在寻找一种方法来轻松地使用管道重定向，但我不知道此时是否有可能。是吗？看起来应该是。

在类 Unix 系统上，当然可以将套接字重定向到派生的子进程，以便子进程直接与客户端对话。 （旧的inetd Unix 服务器是这样工作的。）但是 asyncio 不支持这种操作模式，原因有两个：

• 它不适用于 Python 和 asyncio 支持的所有系统，尤其是在 Windows 上。
• 它与核心 asyncio 功能不兼容，例如传输/协议和流，这些功能假定对底层套接字的所有权和独占访问权。

即使您不关心可移植性，请考虑第二点：您可能需要处理或记录 TCP 客户端和子进程之间交换的数据，如果它们被焊接在一起，您就无法做到这一点内核级别。此外，与仅处理不透明的子进程相比，在异步协程中实现超时和取消要容易得多。

如果不可移植性和无法控制通信对您的用例来说很好，那么您可能首先不需要 asyncio - 没有什么可以阻止您生成一个运行经典阻塞服务器的线程来处理每个客户端os.fork、os.dup2 和 os.execlp 的序列与您在 C 中编写的相同。

编辑

正如 OP 在评论中指出的那样，原始代码通过杀死子进程来处理 TCP 客户端断开连接。在流层，连接丢失由流反映文件结束或引发异常。在上面的代码中，可以通过将通用的self._transfer() 替换为处理这种情况的更具体的协程来轻松应对连接丢失。例如，而不是：

sock_to_child = loop.create_task(self._transfer(r, child.stdin))

...可以写：

sock_to_child = loop.create_task(self._sock_to_child(r, child))

并像这样定义_sock_to_child（未经测试）：

async def _sock_to_child(self, reader, child):
    try:
        await self._transfer(reader, child.stdin)
    except IOError as e:
        # IO errors are an expected part of the workflow,
        # we don't want to propagate them
        print('exception:', e)
    child.kill()

如果子进程的寿命超过 TCP 客户端，child.kill() 行可能永远不会执行，因为协程将被 _handle_client 取消，同时在 src.read() 内的 _transfer() 中暂停。