【网络编程实践】2.3.5.5 基于IO复用（非阻塞IO）实现的 netcat

使用非阻塞IO可以有效避免上述情况的发生。但非阻塞IO在编程上要比阻塞IO更难，并且在程序的维护上比较痛苦。一般使用非阻塞IO编程时建议使用一些封装好的网络库比较容易编写。

下面是python实现的基于IO复用的非阻塞IO的netcat 。

#!/usr/bin/python

import errno
import fcntl
import os
import select
import socket
import sys
# 设置非阻塞
def setNonBlocking(fd):
    flags = fcntl.fcntl(fd, fcntl.F_GETFL)
    fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)

# 非阻塞的写数据
def nonBlockingWrite(fd, data):
    try:
        nw = os.write(fd, data)
        return nw
    except OSError as e:
        if e.errno == errno.EWOULDBLOCK:
            return -1


def relay(sock):
    socketEvents = select.POLLIN
    poll = select.poll()
    poll.register(sock, socketEvents)
    poll.register(sys.stdin, select.POLLIN)

    setNonBlocking(sock)
    # setNonBlocking(sys.stdin)
    # setNonBlocking(sys.stdout)

    done = False
    socketOutputBuffer = ''
    while not done:
        events = poll.poll(10000)  # 10 seconds
        for fileno, event in events:
            if event & select.POLLIN:
                if fileno == sock.fileno():		# 套接字可读（该文件描述符号 等于 网络链接套接字文件描述符）
                    data = sock.recv(8192)
                    if data:
                        nw = sys.stdout.write(data)  # stdout does support non-blocking write, though
                    else:
                        done = True
                else:
                    assert fileno == sys.stdin.fileno()	# 标准输入可读
                    data = os.read(fileno, 8192)		# 将输入读到date中
                    if data:
                        assert len(socketOutputBuffer) == 0	# 读之前确认buf中没有数据，否则可能会造成数据乱序
                        nw = nonBlockingWrite(sock.fileno(), data)	# 写数据
	                        if nw < len(data):	# 如果数据没有写完，需要保存剩余的数据到 scoketOutputBuffer中
                            	if nw < 0:
                                	nw = 0
	                            socketOutputBuffer = data[nw:]	# 暂存没有写完的数据
	                            socketEvents |= select.POLLOUT	# 开始关注 pollout 事件（此时socketEvents 同时有sock的读写事件）
	                            poll.register(sock, socketEvents)	# 重新注册sock，关注sock的读（接收对端数据）和写（本端有未发送完的数据）
	                            poll.unregister(sys.stdin)	# 不再关注stdin事件
                    else:	# 如果没有标准输入，则表示可以关闭连接了
                        sock.shutdown(socket.SHUT_WR)
                        poll.unregister(sys.stdin)
            if event & select.POLLOUT:
                if fileno == sock.fileno():	# 判断是否是网络套接字上的事件
                    assert len(socketOutputBuffer) > 0	# 断言，只有当socketOutputBuffer有剩余的时候，才需要向sock写数据
                    nw = nonBlockingWrite(sock.fileno(), socketOutputBuffer)# 向sock文件描述符写数据
                    if nw < len(socketOutputBuffer):	# 如果仍然没写完，则保留余下的数据，继续重复上述操作
                        assert nw > 0
                        socketOutputBuffer = socketOutputBuffer[nw:]
                    else:								# 如果数据写完了，则套接字上的写事件不需要了
                        socketOutputBuffer = ''
                        socketEvents &= ~select.POLLOUT	# 取消套接字写事件标记，此时的socketEvents == select.POLLIN
                        poll.register(sock, socketEvents)# 重新注册套接字事件（关注套接字上的读事件）
                        poll.register(sys.stdin, select.POLLIN) # 关注，标准输入上的读事件



def main(argv):
    if len(argv) < 3:
        binary = argv[0]
        print "Usage:\n  %s -l port\n  %s host port" % (argv[0], argv[0])
        print (sys.stdout.write)
        return
    port = int(argv[2])
    if argv[1] == "-l":
        # server
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server_socket.bind(('', port))
        server_socket.listen(5)
        (client_socket, client_address) = server_socket.accept()
        server_socket.close()	# 关闭监听套接字
        relay(client_socket)
    else:
        # client
        sock = socket.create_connection((argv[1], port))
        relay(sock)


if __name__ == "__main__":
    main(sys.argv)

因为设置为非阻塞IO后，写操作会有无法将全部数据写完的情况发生（short write），因此在发生数据未写完时，我们再注册一个写事件，让程序将未写完的数据全部写完。

程序结构如下：

测试结果：可以看到这次，即使 nc 上有输入也不会阻塞程序。