Use readers/writers instead of fdmap

Author: ask

celery/worker/consumer.py

@@ -84,6 +84,7 @@
 from Queue import Empty
 
 from kombu.utils.encoding import safe_repr
+from kombu.utils.eventio import READ, WRITE, ERR
 
 from celery.app import app_or_default
 from celery.datastructures import AttributeDict
@@ -370,7 +371,7 @@ def consume_messages(self, sleep=sleep, min=min, Empty=Empty):
             qos = self.qos
             update_qos = qos.update
             update_readers = hub.update_readers
-            fdmap = hub.fdmap
+            readers, writers = hub.readers, hub.writers
             poll = hub.poller.poll
             fire_timers = hub.fire_timers
             scheduled = hub.timer._queue
@@ -431,12 +432,18 @@ def on_task_received(body, message):
                     update_qos()
 
                 update_readers(on_poll_start())
-                if fdmap:
+                if readers or writers:
                     connection.more_to_read = True
                     while connection.more_to_read:
                         for fileno, event in poll(time_to_sleep) or ():
                             try:
-                                fdmap[fileno](fileno, event)
+                                if event & READ:
+                                    readers[fileno](fileno, event)
+                                if event & WRITE:
+                                    writers[fileno](fileno, event)
+                                if event & ERR:
+                                    readers[fileno](fileno, event)
+                                    writers[fileno](fileno, event)
                             except Empty:
                                 break
                             except socket.error:

celery/worker/hub.py

@@ -51,9 +51,11 @@ class Hub(object):
     READ, WRITE, ERR = READ, WRITE, ERR
 
     def __init__(self, timer=None):
-        self.fdmap = {}
+        self.readers = {}
+        self.writers = {}
         self.timer = Schedule() if timer is None else timer
         self.on_init = []
+        self.on_close = []
         self.on_task = []
 
     def start(self):
@@ -85,11 +87,12 @@ def fire_timers(self, min_delay=1, max_delay=10, max_timers=10):
 
     def add(self, fd, callback, flags):
         self.poller.register(fd, flags)
-        try:
-            fileno = fd.fileno()
-        except AttributeError:
-            fileno = fd
-        self.fdmap[fileno] = callback
+        if not isinstance(fd, int):
+            fd = fd.fileno()
+        if flags & READ:
+            self.readers[fd] = callback
+        if flags & WRITE:
+            self.writers[fd] = callback
 
     def add_reader(self, fd, callback):
         return self.add(fd, callback, READ | ERR)
@@ -110,7 +113,10 @@ def remove(self, fd):
             pass
 
     def close(self):
-        [self.remove(fd) for fd in self.fdmap.keys()]
+        [self.remove(fd) for fd in self.readers.keys()]
+        [self.remove(fd) for fd in self.writers.keys()]
+        for callback in self.on_close:
+            callback(self)
 
     @cached_property
     def scheduler(self):