python-websocket开发（二）网络编程

    上一篇博客介绍了python websocket的协议解析，这一篇介绍网络编程，如何用socket接收发原始数据，
    通过上一篇文章的协议解析，形成程序员能够理解的数据，照例先贴代码，再解释。

from socket import *
import selectors
from sdk.web_socket.protocol_handler import WebSocketProtocolHandler
import multiprocessing
import json
import os
from sdk.web_socket.base_interface import BaseInterface
import time


class WebSocketServer:
    def __init__(self, port=8000):
        self.__server_fileobj = socket(AF_INET, SOCK_STREAM)
        self.__server_fileobj.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        self.__server_fileobj.bind(('0.0.0.0', port))
        self.__server_fileobj.listen(5)
        self.__server_fileobj.setblocking(False)  # 设置socket的接口为非阻塞
        print(f"==========WebSocket bind:{port}==========")

        # socket mgr
        self.__conn_mgr = BaseInterface.create_ipc_dict()
        self.__protocol_handler_mgr = BaseInterface.create_ipc_dict()
        self.__heartcheck_mgr = BaseInterface.create_ipc_dict()

        # multiprocess mgr
        self.__ipc_queue = BaseInterface.create_ipc_queue()

        self.connect_callback = BaseInterface.create_ipc_list()
        self.disconnect_callback = BaseInterface.create_ipc_list()

    def register_connect_callback(self, ref):
        self.connect_callback.append(ref)

    def register_disconnect_callback(self, ref):
        self.disconnect_callback.append(ref)

    def start(self, msg_handler=None):
        receive_process = multiprocessing.Process(name="Receive", target=WebSocketServer.read_main_loop, args=(self, msg_handler, BaseInterface.ipc_func_list, BaseInterface.ipc_param_list))
        receive_process.start()

        send_process = multiprocessing.Process(name="Send", target=WebSocketServer.write_main_loop, args=(self, ))
        send_process.start()

    def push_msg(self, fd, msg_type, msg_body):
        try:
            msg_dict = {
                "msg_type": msg_type,
                "msg_body": msg_body
            }
            msg = json.dumps(msg_dict).encode("utf-8")
            exclusive_dict = {
                "fd": fd,
                "msg": msg
            }
            self.__ipc_queue.put(exclusive_dict)
        except Exception:
            print("can't push this msg_type and msg_body")

    def broadcast(self, msg_type, msg_body):
        try:
            for k in self.__conn_mgr.keys():
                self.push_msg(k, msg_type, msg_body)
        except Exception:
            print("broadcast may not complete.")

    @staticmethod
    def accept(server, server_fileobj, sel, mask, msg_handler, fd_conn_mgr):
        # print("accept")
        conn, addr = server_fileobj.accept()
        sel.register(conn, selectors.EVENT_READ, WebSocketServer.read)
        server.__add_conn(conn)
        fd_conn_mgr[conn.fileno()] = conn

    @staticmethod
    def read(server, conn, sel, mask, msg_handler, fd_conn_mgr):
        try:
            handler = server.__get_protocol_handler(conn.fileno())
            if handler:
                body = handler.recv_msg()

                if not body:
                    sel.unregister(conn)
                    server.__remove_conn(conn)
                    conn.close()
                    if conn.fileno() in fd_conn_mgr.keys():
                        del fd_conn_mgr[conn.fileno()]
                else:
                    if msg_handler is not None:
                        msg_handler.handle_msg(server, body, conn.fileno())
        except Exception as e:
            print(e)
            sel.unregister(conn)
            server.__remove_conn(conn)
            conn.close()
            if conn.fileno() in fd_conn_mgr.keys():
                del fd_conn_mgr[conn.fileno()]

    @staticmethod
    def process_info(info=""):
        print(f"=====module name: {info}=====\nparent process: {os.getppid()}\nprocess id: {os.getpid()}")

    @staticmethod
    def read_main_loop(server, msg_handler, func_list, param_list):
        BaseInterface.ipc_func_list = func_list
        BaseInterface.ipc_param_list = param_list
        BaseInterface.sync_ipc_data()

        msg_handler.initialize_pool(server)
        WebSocketServer.process_info("Receive Loop")
        sel = selectors.DefaultSelector()
        sel.register(server.__server_fileobj, selectors.EVENT_READ, WebSocketServer.accept)
        fd_conn_mgr = dict()

        last_heartbreak_time = time.clock()
        while True:
            # 检测所有的fileobj，是否有完成wait data的
            events = sel.select(timeout=1)
            for sel_obj, mask in events:
                # print(f"continue..{sel_obj}  {mask}")
                callback = sel_obj.data
                callback(server, sel_obj.fileobj, sel, mask, msg_handler, fd_conn_mgr)

            # 检测所有fd是否有效
            # for k, v in server.__heartcheck_mgr.items():
            #     server.__heartcheck_mgr[k] -= 1
            #     if v <= 0:
            #         conn = fd_conn_mgr[k]
            #         sel.unregister(conn)
            #         server.__remove_conn(conn)
            #         conn.close()
            #         del fd_conn_mgr[k]

            # 发送心跳检测包 间隔2秒
            # if time.clock() - last_heartbreak_time > 2:
            #     last_heartbreak_time = time.clock()
            #     for k, v in server.__conn_mgr.items():
            #         server.push_msg(k, msg_type=2, msg_body="")

    @staticmethod
    def write_main_loop(server):
        WebSocketServer.process_info("Send Loop")
        while True:
            try:
                exclusive_dict = server.__ipc_queue.get(timeout=1)
                fd = exclusive_dict.get("fd", None)
                msg = exclusive_dict.get("msg", None)
                if fd and msg:
                    handler = server.__get_protocol_handler(fd)
                    if handler:
                        handler.send_msg(msg)
            except Exception:
                pass
            finally:
                pass

    def __add_conn(self, conn):
        self.__conn_mgr[conn.fileno()] = conn
        self.__protocol_handler_mgr[conn.fileno()] = WebSocketProtocolHandler(conn)
        self.__heartcheck_mgr[conn.fileno()] = 5

        print(f"fd: {conn.fileno()} join {conn}")
        self.push_msg(fd=conn.fileno(), msg_type=1, msg_body=conn.fileno())
        if self.connect_callback is not None:
            for f in self.connect_callback:
                f(conn.fileno())

    def __remove_conn(self, conn):
        if conn.fileno() in self.__conn_mgr.keys():
            del self.__conn_mgr[conn.fileno()]
        if conn.fileno() in self.__protocol_handler_mgr.keys():
            del self.__protocol_handler_mgr[conn.fileno()]
        if conn.fileno() in self.__heartcheck_mgr.keys():
            del self.__heartcheck_mgr[conn.fileno()]

        print(f"fd: {conn.fileno()} leave {conn}")
        if self.disconnect_callback is not None:
            for f in self.disconnect_callback:
                f(conn.fileno())

    def __get_protocol_handler(self, fd):
        return self.__protocol_handler_mgr.get(fd, None)

    def reset_heartbreak(self, fd):
        self.__heartcheck_mgr[fd] = 5

整个代码分为3个部分，客户端连接管理、收数据、发数据。
首先是客户端连接管理，这是传统的网络编程
然后是收数据、发数据各一个进程处理，和主进程区分开来。