使用非阻塞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 上有输入也不会阻塞程序。