python socket

connect_event.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-

"""
Tornado source code write comments in `iostream._handle_events` which links is https://github.com/tornadoweb/tornado/blob/master/tornado/iostream.py#L695: 

```
# Most IOLoops will report a write failed connect
# with the WRITE event, but SelectIOLoop reports a
# READ as well so we must check for connecting before
# either.
```

When call `connect` with unblocking option, the ioloop(select, poll, epoll) in different will report different events.

To call connect failed, WRITE and ERROR will be reported in poll where it is as same as epoll, but select will report WRITE and READ events at the same time.

Being reported READ, it doesn't mean that the fd can be read in select, client must try to checking the connection status with `getsockopt`. Note, connect in progress, getsockopt still return err is 0.

This script can be use to make tests for different ioloop in connect case.

Try to run like this:

`$ python connect_event.py -h`

"""


__author__ = 'master'

import time
import os
import sys
import socket
import select
import argparse
import threading

from wsgiref.simple_server import demo_app, make_server

NONE = 0
READ = select.POLLIN
WRITE = select.POLLOUT
ERROR = select.POLLERR | select.POLLHUP


def print_color(color, text):
    print('{}{}\033[0m'.format(color, text))


class BaseIOLoop(object):

    def __init__(self, sock, address):
        self._sock = sock
        self._address = address

    def handle_loop(self):
        raise NotImplementedError

    def handle_connect(self):
        print('CONNECT start')
        try:
            self._sock.connect(self._address)
        except Exception as e:
            print('CONNECT return {e}\n'.format(e=e))
        else:
            print('CONNECT succ \n')

    def handle_connect_ret(self):
        err = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
        if err != 0:
            error = socket.error(err, os.strerror(err))
            print_color('\033[1;31m', '\n=== CONNECT fail {error}'.format(error=error))
        else:
            print_color('\033[1;32m', '\n=== CONNECT succ \n')

    def run(self):
        self.handle_connect()
        self.handle_loop()
        self.handle_connect_ret()


class SelectIOLoop(BaseIOLoop):

    def handle_loop(self):
        r = w = e = [self._sock]
        rt, wt, et = select.select(r, w, e)
        color = '\033[1;34m'
        print_color(color, 'EVENT report: ')
        if rt:
            print_color(color, '>>> READ')
        if wt:
            print_color(color, '>>> WRITE')
        if et:
            print_color(color, '>>> ERROR')


class PollIOLoop(BaseIOLoop):

    def __init__(self, sock, address):
        super(PollIOLoop, self).__init__(sock, address)
        self._poll = select.poll()
        self._poll.register(self._sock.fileno(), WRITE)

    def handle_loop(self):
        color = '\033[1;34m'
        events = self._poll.poll(1)
        for fd, event in events:
            print_color(color, 'EVENT report: ')
            if event & READ:
                print_color(color, '>>> READ')
            if event & WRITE:
                print_color(color, '>>> WRITE')
            if event & ERROR:
                print_color(color, '>>> ERROR')


class EpollIOLoop(PollIOLoop):

    def __init__(self, sock, address):
        super(EpollIOLoop, self).__init__(sock, address)
        self._poll = select.epoll()
        self._poll.register(self._sock.fileno(), WRITE)


def run_server(host, port):
    httpd = make_server(host, port, demo_app)
    sa = httpd.socket.getsockname()
    print("Serving HTTP on", sa[0], "port", sa[1], "...\n")
    httpd.handle_request()
    httpd.server_close()


def main(connected, loop):
    address = ('', 8000)

    if connected:
        t = threading.Thread(target=run_server, args=address)
        t.setDaemon(True)
        t.start()
        time.sleep(3)

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setblocking(False)
    if loop == 'select':
        SelectIOLoop(sock, address).run()
    elif loop == 'poll':
        PollIOLoop(sock, address).run()
    elif loop == 'epoll':
        EpollIOLoop(sock, address).run()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--connected", help="connect succ options", action="store_true")
    parser.add_argument("-l", "--loop", help="choose loop", choices=["select", "poll", "epoll"], required=True)
    args = parser.parse_args()
    color = '\033[1;36;40m'
    loop_str = args.loop.upper()
    if args.connected:
        print_color(color, '=== {} Connect succ ===               '.format(loop_str))
    else:
        print_color(color, '=== {} Connect fail ===               '.format(loop_str))

    print('\n')
    try:
        main(args.connected, args.loop)
    except KeyError as e:
        print(e)
        print("exit...")
        sys.exit(True)