impl

shadowsocks/eventloop.py

@@ -0,0 +1,174 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+# Copyright (c) 2014 clowwindy
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# from ssloop
+# https://github.com/clowwindy/ssloop
+
+
+import select
+from collections import defaultdict
+
+
+__all__ = ['EventLoop', 'MODE_NULL', 'MODE_IN', 'MODE_OUT', 'MODE_ERR',
+           'MODE_HUP', 'MODE_NVAL']
+
+MODE_NULL = 0x00
+MODE_IN = 0x01
+MODE_OUT = 0x04
+MODE_ERR = 0x08
+MODE_HUP = 0x10
+MODE_NVAL = 0x20
+
+
+class EpollLoop(object):
+
+    def __init__(self):
+        self._epoll = select.epoll()
+
+    def poll(self, timeout):
+        return self._epoll.poll(timeout)
+
+    def add_fd(self, fd, mode):
+        self._epoll.register(fd, mode)
+
+    def remove_fd(self, fd):
+        self._epoll.unregister(fd)
+
+    def modify_fd(self, fd, mode):
+        self._epoll.modify(fd, mode)
+
+
+class KqueueLoop(object):
+
+    MAX_EVENTS = 1024
+
+    def __init__(self):
+        self._kqueue = select.kqueue()
+        self._fds = {}
+
+    def _control(self, fd, mode, flags):
+        events = []
+        if mode & MODE_IN:
+            events.append(select.kevent(fd, select.KQ_FILTER_READ, flags))
+        if mode & MODE_OUT:
+            events.append(select.kevent(fd, select.KQ_FILTER_WRITE, flags))
+        for e in events:
+            self._kqueue.control([e], 0)
+
+    def poll(self, timeout):
+        if timeout < 0:
+            timeout = None  # kqueue behaviour
+        events = self._kqueue.control(None, KqueueLoop.MAX_EVENTS, timeout)
+        results = defaultdict(lambda: MODE_NULL)
+        for e in events:
+            fd = e.ident
+            if e.filter == select.KQ_FILTER_READ:
+                results[fd] |= MODE_IN
+            elif e.filter == select.KQ_FILTER_WRITE:
+                results[fd] |= MODE_OUT
+        return results.iteritems()
+
+    def add_fd(self, fd, mode):
+        self._fds[fd] = mode
+        self._control(fd, mode, select.KQ_EV_ADD)
+
+    def remove_fd(self, fd):
+        self._control(fd, self._fds[fd], select.KQ_EV_DELETE)
+        del self._fds[fd]
+
+    def modify_fd(self, fd, mode):
+        self.remove_fd(fd)
+        self.add_fd(fd, mode)
+
+
+class SelectLoop(object):
+
+    def __init__(self):
+        self._r_list = set()
+        self._w_list = set()
+        self._x_list = set()
+
+    def poll(self, timeout):
+        r, w, x = select.select(self._r_list, self._w_list, self._x_list,
+                                timeout)
+        results = defaultdict(lambda: MODE_NULL)
+        for p in [(r, MODE_IN), (w, MODE_OUT), (x, MODE_ERR)]:
+            for fd in p[0]:
+                results[fd] |= p[1]
+        return results.items()
+
+    def add_fd(self, fd, mode):
+        if mode & MODE_IN:
+            self._r_list.add(fd)
+        if mode & MODE_OUT:
+            self._w_list.add(fd)
+        if mode & MODE_ERR:
+            self._x_list.add(fd)
+
+    def remove_fd(self, fd):
+        if fd in self._r_list:
+            self._r_list.remove(fd)
+        if fd in self._w_list:
+            self._w_list.remove(fd)
+        if fd in self._x_list:
+            self._x_list.remove(fd)
+
+    def modify_fd(self, fd, mode):
+        self.remove_fd(fd)
+        self.add_fd(fd, mode)
+
+
+class EventLoop(object):
+    def __init__(self):
+        if hasattr(select, 'epoll'):
+            self._impl = EpollLoop()
+        elif hasattr(select, 'kqueue'):
+            self._impl = KqueueLoop()
+        elif hasattr(select, 'select'):
+            self._impl = SelectLoop()
+        else:
+            raise Exception('can not find any available functions in select '
+                            'package')
+        self._fd_to_f = defaultdict(list)
+
+    def poll(self, timeout=None):
+        events = self._impl.poll(timeout)
+        return ((self._fd_to_f[fd], event) for fd, event in events)
+
+    def add(self, f, mode):
+        fd = f.fileno()
+        self._fd_to_f[fd].append(f)
+        self._impl.add_fd(fd, mode)
+
+    def remove(self, f):
+        fd = f.fileno()
+        a = self._fd_to_f[fd]
+        if len(a) <= 1:
+            self._fd_to_f[fd] = None
+        else:
+            a.remove(f)
+        self._impl.remove_fd(fd)
+
+    def modify(self, f, mode):
+        fd = f.fileno()
+        self._impl.modify_fd(fd, mode)