PYTHON-697 Asyncore reactor should use a global variable instead of a class variable for the event loop

Jaume Marhuenda · Jaume Marhuenda · commit 6bce25c918f4 · 2018-05-18T09:45:21.000-04:00
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -15,6 +15,7 @@ Features
 * Allow filter queries with fields that have an index managed outside of cqlengine (PYTHON-966)
 * Twisted SSL Support (PYTHON-343)
 * Support IS NOT NULL operator in cqlengine (PYTHON-968)
+* Asyncore reactors should use a global variable instead of a class variable for the event loop (PYTHON-697)
 
 Other
 -----
diff --git a/cassandra/io/asyncorereactor.py b/cassandra/io/asyncorereactor.py
@@ -43,13 +43,9 @@
 
 _dispatcher_map = {}
 
-def _cleanup(loop_weakref):
-    try:
-        loop = loop_weakref()
-    except ReferenceError:
-        return
-
-    loop._cleanup()
+def _cleanup(loop):
+    if loop:
+        loop._cleanup()
 
 
 class WaitableTimer(Timer):
@@ -228,8 +224,6 @@ def __init__(self):
             dispatcher = _BusyWaitDispatcher()
         self._loop_dispatcher = dispatcher
 
-        atexit.register(partial(_cleanup, weakref.ref(self)))
-
     def maybe_start(self):
         should_start = False
         did_acquire = False
@@ -299,40 +293,43 @@ def _cleanup(self):
         log.debug("Dispatchers were closed")
 
 
+_global_loop = None
+atexit.register(partial(_cleanup, _global_loop))
+
+
 class AsyncoreConnection(Connection, asyncore.dispatcher):
     """
     An implementation of :class:`.Connection` that uses the ``asyncore``
     module in the Python standard library for its event loop.
     """
 
-    _loop = None
-
     _writable = False
     _readable = False
 
     @classmethod
     def initialize_reactor(cls):
-        if not cls._loop:
-            cls._loop = AsyncoreLoop()
+        global _global_loop
+        if not _global_loop:
+            _global_loop = AsyncoreLoop()
         else:
             current_pid = os.getpid()
-            if cls._loop._pid != current_pid:
+            if _global_loop._pid != current_pid:
                 log.debug("Detected fork, clearing and reinitializing reactor state")
                 cls.handle_fork()
-                cls._loop = AsyncoreLoop()
+                _global_loop = AsyncoreLoop()
 
     @classmethod
     def handle_fork(cls):
-        global _dispatcher_map
+        global _dispatcher_map, _global_loop
         _dispatcher_map = {}
-        if cls._loop:
-            cls._loop._cleanup()
-            cls._loop = None
+        if _global_loop:
+            _global_loop._cleanup()
+            _global_loop = None
 
     @classmethod
     def create_timer(cls, timeout, callback):
         timer = Timer(timeout, callback)
-        cls._loop.add_timer(timer)
+        _global_loop.add_timer(timer)
         return timer
 
     def __init__(self, *args, **kwargs):
@@ -344,14 +341,14 @@ def __init__(self, *args, **kwargs):
         self._connect_socket()
 
         # start the event loop if needed
-        self._loop.maybe_start()
+        _global_loop.maybe_start()
 
         init_handler = WaitableTimer(
             timeout=0,
             callback=partial(asyncore.dispatcher.__init__,
                              self, self._socket, _dispatcher_map)
         )
-        self._loop.add_timer(init_handler)
+        _global_loop.add_timer(init_handler)
         init_handler.wait(kwargs["connect_timeout"])
 
         self._writable = True
@@ -451,7 +448,7 @@ def push(self, data):
         with self.deque_lock:
             self.deque.extend(chunks)
             self._writable = True
-        self._loop.wake_loop()
+        _global_loop.wake_loop()
 
     def writable(self):
         return self._writable
diff --git a/tests/integration/standard/test_connection.py b/tests/integration/standard/test_connection.py
@@ -20,12 +20,13 @@
 from functools import partial
 from six.moves import range
 import sys
+import threading
 from threading import Thread, Event
 import time
-import weakref
 
 from cassandra import ConsistencyLevel, OperationTimedOut
 from cassandra.cluster import NoHostAvailable, ConnectionShutdown, Cluster
+import cassandra.io.asyncorereactor
 from cassandra.io.asyncorereactor import AsyncoreConnection
 from cassandra.protocol import QueryMessage
 from cassandra.connection import Connection
@@ -401,6 +402,28 @@ def setUp(self):
             raise unittest.SkipTest("Can't test asyncore with monkey patching")
         ConnectionTests.setUp(self)
 
+    def test_subclasses_share_loop(self):
+        class C1(AsyncoreConnection):
+            pass
+
+        class C2(AsyncoreConnection):
+            pass
+
+        cassandra.io.asyncorereactor._global_loop._cleanup()
+        cassandra.io.asyncorereactor._global_loop = None
+
+        clusterC1 = Cluster(connection_class=C1)
+        clusterC1.connect(wait_for_all_pools=True)
+
+        clusterC2 = Cluster(connection_class=C2)
+        clusterC2.connect(wait_for_all_pools=True)
+        self.addCleanup(clusterC1.shutdown)
+        self.addCleanup(clusterC2.shutdown)
+
+        event_loops_threads = [thread for thread in threading.enumerate() if
+                               thread.name == "cassandra_driver_event_loop"]
+        self.assertEqual(len(event_loops_threads), 1)
+
 
 class LibevConnectionTests(ConnectionTests, unittest.TestCase):
 
diff --git a/tests/unit/io/test_asyncorereactor.py b/tests/unit/io/test_asyncorereactor.py
@@ -18,6 +18,7 @@
 
 from mock import patch
 import socket
+import cassandra.io.asyncorereactor
 from cassandra.io.asyncorereactor import AsyncoreConnection
 from tests import is_monkey_patched
 from tests.unit.io.utils import ReactorTestMixin, TimerTestMixin, noop_if_monkey_patched
@@ -76,7 +77,7 @@ def create_timer(self):
 
     @property
     def _timers(self):
-        return self.connection._loop._timers
+        return cassandra.io.asyncorereactor._global_loop._timers
 
     def setUp(self):
         if is_monkey_patched():
