Merge pull request datastax#950 from datastax/python-973

PYTHON-973: Use global variable for libev loops so it can be subclassed

Author: mambocab

CHANGELOG.rst

@@ -4,6 +4,10 @@ Bug Fixes
 ---------
 * Tokenmap.get_replicas returns the wrong value if token coincides with the end of the range (PYTHON-978)
 
+Other
+-----
+* Use global variable for libev loops so it can be subclassed (PYTHON-973)
+
 3.14.0
 ======
 April 17, 2018

cassandra/io/libevreactor.py

@@ -41,12 +41,9 @@
 log = logging.getLogger(__name__)
 
 
-def _cleanup(loop_weakref):
-    try:
-        loop = loop_weakref()
-    except ReferenceError:
-        return
-    loop._cleanup()
+def _cleanup(loop):
+    if loop:
+        loop._cleanup()
 
 
 class LibevLoop(object):
@@ -84,8 +81,6 @@ def __init__(self):
         self._timers = TimerManager()
         self._loop_timer = libev.Timer(self._loop, self._on_loop_timer)
 
-        atexit.register(partial(_cleanup, weakref.ref(self)))
-
     def maybe_start(self):
         should_start = False
         with self._lock:
@@ -228,36 +223,41 @@ def _loop_will_run(self, prepare):
             self._notifier.send()
 
 
+_global_loop = None
+atexit.register(partial(_cleanup, _global_loop))
+
+
 class LibevConnection(Connection):
     """
     An implementation of :class:`.Connection` that uses libev for its event loop.
     """
-    _libevloop = None
     _write_watcher_is_active = False
     _read_watcher = None
     _write_watcher = None
     _socket = None
 
     @classmethod
     def initialize_reactor(cls):
-        if not cls._libevloop:
-            cls._libevloop = LibevLoop()
+        global _global_loop
+        if not _global_loop:
+            _global_loop = LibevLoop()
         else:
-            if cls._libevloop._pid != os.getpid():
+            if _global_loop._pid != os.getpid():
                 log.debug("Detected fork, clearing and reinitializing reactor state")
                 cls.handle_fork()
-                cls._libevloop = LibevLoop()
+                _global_loop = LibevLoop()
 
     @classmethod
     def handle_fork(cls):
-        if cls._libevloop:
-            cls._libevloop._cleanup()
-            cls._libevloop = None
+        global _global_loop
+        if _global_loop:
+            _global_loop._cleanup()
+            _global_loop = None
 
     @classmethod
     def create_timer(cls, timeout, callback):
         timer = Timer(timeout, callback)
-        cls._libevloop.add_timer(timer)
+        _global_loop.add_timer(timer)
         return timer
 
     def __init__(self, *args, **kwargs):
@@ -268,16 +268,16 @@ def __init__(self, *args, **kwargs):
         self._connect_socket()
         self._socket.setblocking(0)
 
-        with self._libevloop._lock:
-            self._read_watcher = libev.IO(self._socket.fileno(), libev.EV_READ, self._libevloop._loop, self.handle_read)
-            self._write_watcher = libev.IO(self._socket.fileno(), libev.EV_WRITE, self._libevloop._loop, self.handle_write)
+        with _global_loop._lock:
+            self._read_watcher = libev.IO(self._socket.fileno(), libev.EV_READ, _global_loop._loop, self.handle_read)
+            self._write_watcher = libev.IO(self._socket.fileno(), libev.EV_WRITE, _global_loop._loop, self.handle_write)
 
         self._send_options_message()
 
-        self._libevloop.connection_created(self)
+        _global_loop.connection_created(self)
 
         # start the global event loop if needed
-        self._libevloop.maybe_start()
+        _global_loop.maybe_start()
 
     def close(self):
         with self.lock:
@@ -286,7 +286,8 @@ def close(self):
             self.is_closed = True
 
         log.debug("Closing connection (%s) to %s", id(self), self.host)
-        self._libevloop.connection_destroyed(self)
+
+        _global_loop.connection_destroyed(self)
         self._socket.close()
         log.debug("Closed socket to %s", self.host)
 
@@ -367,4 +368,4 @@ def push(self, data):
 
         with self._deque_lock:
             self.deque.extend(chunks)
-            self._libevloop.notify()
+            _global_loop.notify()

tests/integration/standard/test_connection.py

@@ -38,6 +38,7 @@
     requiresmallclockgranularity, greaterthancass20
 try:
     from cassandra.io.libevreactor import LibevConnection
+    import cassandra.io.libevreactor
 except ImportError:
     LibevConnection = None
 
@@ -392,26 +393,13 @@ def test_connect_timeout(self):
                 break
         self.assertTrue(exception_thrown)
 
-
-class AsyncoreConnectionTests(ConnectionTests, unittest.TestCase):
-
-    klass = AsyncoreConnection
-
-    def setUp(self):
-        if is_monkey_patched():
-            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)
 
@@ -420,14 +408,35 @@ class C2(AsyncoreConnection):
         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)
+        self.assertEqual(len(get_eventloop_threads(self.event_loop_name)), 1)
+
+
+def get_eventloop_threads(name):
+    import threading
+    event_loops_threads = [thread for thread in threading.enumerate() if name == thread.name]
+
+    return event_loops_threads
+
+
+class AsyncoreConnectionTests(ConnectionTests, unittest.TestCase):
+
+    klass = AsyncoreConnection
+    event_loop_name = "cassandra_driver_event_loop"
+
+    def setUp(self):
+        if is_monkey_patched():
+            raise unittest.SkipTest("Can't test asyncore with monkey patching")
+        ConnectionTests.setUp(self)
+
+    def clean_global_loop(self):
+        cassandra.io.asyncorereactor._global_loop._cleanup()
+        cassandra.io.asyncorereactor._global_loop = None
 
 
 class LibevConnectionTests(ConnectionTests, unittest.TestCase):
 
     klass = LibevConnection
+    event_loop_name = "event_loop"
 
     def setUp(self):
         if is_monkey_patched():
@@ -436,3 +445,7 @@ def setUp(self):
             raise unittest.SkipTest(
                 'libev does not appear to be installed properly')
         ConnectionTests.setUp(self)
+
+    def clean_global_loop(self):
+        cassandra.io.libevreactor._global_loop._cleanup()
+        cassandra.io.libevreactor._global_loop = None

tests/unit/io/test_libevreactor.py

@@ -35,7 +35,6 @@ class LibevConnectionTest(ReactorTestMixin, unittest.TestCase):
 
     connection_class = LibevConnection
     socket_attr_name = '_socket'
-    loop_attr_name = '_libevloop'
     null_handle_function_args = None, 0
 
     def setUp(self):
@@ -62,7 +61,7 @@ def test_watchers_are_finished(self):
         Test for asserting that watchers are closed in LibevConnection
 
         This test simulates a process termination without calling cluster.shutdown(), which would trigger
-        LibevConnection._libevloop._cleanup. It will check the watchers have been closed
+        _global_loop._cleanup. It will check the watchers have been closed
         Finally it will restore the LibevConnection reactor so it doesn't affect
         the rest of the tests
 
@@ -72,24 +71,25 @@ def test_watchers_are_finished(self):
 
         @test_category connection
         """
-        with patch.object(LibevConnection._libevloop, "_thread"),\
-             patch.object(LibevConnection._libevloop, "notify"):
+        from cassandra.io.libevreactor import _global_loop
+        with patch.object(_global_loop, "_thread"),\
+             patch.object(_global_loop, "notify"):
 
             self.make_connection()
 
             # We have to make a copy because the connections shouldn't
             # be alive when we verify them
-            live_connections = set(LibevConnection._libevloop._live_conns)
+            live_connections = set(_global_loop._live_conns)
 
             # This simulates the process ending without cluster.shutdown()
             # being called, then with atexit _cleanup for libevreactor would
             # be called
-            libev__cleanup(weakref.ref(LibevConnection._libevloop))
+            libev__cleanup(_global_loop)
             for conn in live_connections:
                 self.assertTrue(conn._write_watcher.stop.mock_calls)
                 self.assertTrue(conn._read_watcher.stop.mock_calls)
 
-        LibevConnection._libevloop._shutdown = False
+        _global_loop._shutdown = False
 
 
 class LibevTimerPatcher(unittest.TestCase):
@@ -125,7 +125,8 @@ def create_timer(self):
 
     @property
     def _timers(self):
-        return self.connection._libevloop._timers
+        from cassandra.io.libevreactor import _global_loop
+        return _global_loop._timers
 
     def make_connection(self):
         c = LibevConnection('1.2.3.4', cql_version='3.0.1')