link_future() now raises an exception if it is called from the greenthread of

the aiogreen event loop.

4 files changed, 61 insertions, 7 deletions

diff --git a/aiogreen.py b/aiogreen.py
index 43c9ff2..1a7fdde 100644
--- a/aiogreen.py
+++ b/aiogreen.py
@@ -191,6 +191,8 @@ class _Selector(asyncio.selectors._BaseSelectorImpl):
 
 class EventLoop(asyncio.SelectorEventLoop):
     def __init__(self):
+        self._greenthread = None
+
         # Store a reference to the hub to ensure
         # that we always use the same hub
         self._hub = eventlet.hubs.get_hub()
@@ -234,6 +236,7 @@ class EventLoop(asyncio.SelectorEventLoop):
             self._hub.debug_blocking_resolution = self.slow_callback_duration
 
     def run_forever(self):
+        self._greenthread = eventlet.getcurrent()
         try:
             super(EventLoop, self).run_forever()
         finally:
@@ -241,6 +244,7 @@ class EventLoop(asyncio.SelectorEventLoop):
                 # eventlet event loop is still running: cancel the current
                 # detection of blocking tasks
                 signal.alarm(0)
+            self._greenthread = None
 
     def time(self):
         return self._hub.clock()
@@ -299,7 +303,14 @@ def link_future(future):
 
     Wait for a future or a task from a greenthread.
     Return the result or raise the exception of the future.
+
+    The function must not be called from the greenthread
+    of the aiogreen event loop.
     """
+    if future._loop._greenthread == eventlet.getcurrent():
+        raise RuntimeError("link_future() must not be called from "
+                           "the greenthread of the aiogreen event loop")
+
     event = eventlet.event.Event()
     def done(fut):
         try:
diff --git a/doc/changelog.rst b/doc/changelog.rst
index e43da04..8e211f0 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -7,6 +7,8 @@ Version 0.3 (development version)
 * :func:`wrap_greenthread` now raises an exception if the greenthread is
   running or already finished. In debug mode, the exception is not more logged
   to sys.stderr for greenthreads.
+* :func:`link_future` now raises an exception if it is called from the
+  greenthread of the aiogreen event loop.
 * Fix eventlet detection of blocking tasks: cancel the alarm when the aiogreen
   event loop stops.
 
diff --git a/doc/using.rst b/doc/using.rst
index c26f11b..be1edb2 100644
--- a/doc/using.rst
+++ b/doc/using.rst
@@ -93,6 +93,9 @@ aiogreen specific functions:
    Wait for a future (or a task) from a greenthread.
    Return the result or raise the exception of the future.
 
+   The function must not be called from the greenthread of the aiogreen event
+   loop.
+
    Example of greenthread waiting for a trollius task. The ``progress()``
    callback is called regulary to see that the event loop in not blocked::
 
diff --git a/tests/test_eventlet.py b/tests/test_eventlet.py
index 80b3080..271fdaa 100644
--- a/tests/test_eventlet.py
+++ b/tests/test_eventlet.py
@@ -163,13 +163,6 @@ class EventletTests(tests.TestCase):
         self.loop.run_forever()
         self.assertEqual(result, ["spawn", "spawn_after"])
 
-    def test_greenthread_link_future(self):
-        result = []
-        self.loop.call_soon(eventlet.spawn,
-                            greenthread_link_future, result, self.loop)
-        self.loop.run_forever()
-        self.assertEqual(result, [1, 10, 2, 20, 'error', 4])
-
     def test_set_debug(self):
         hub = eventlet.hubs.get_hub()
         self.assertIs(self.loop._hub, hub)
@@ -186,6 +179,51 @@ class EventletTests(tests.TestCase):
             self.assertEqual(hub.debug_blocking, False)
 
 
+class LinkFutureTests(tests.TestCase):
+    def test_greenthread_link_future(self):
+        result = []
+        self.loop.call_soon(eventlet.spawn,
+                            greenthread_link_future, result, self.loop)
+        self.loop.run_forever()
+        self.assertEqual(result, [1, 10, 2, 20, 'error', 4])
+
+    def test_link_future_not_running(self):
+        result = []
+        fut = asyncio.Future(loop=self.loop)
+        event = eventlet.event.Event()
+
+        def func(event, fut):
+            event.send('link')
+            value = aiogreen.link_future(fut)
+            result.append(value)
+            self.loop.stop()
+
+        eventlet.spawn(func, event, fut)
+        event.wait()
+
+        self.loop.call_soon(fut.set_result, 21)
+        self.loop.run_forever()
+        self.assertEqual(result, [21])
+
+    def test_link_future_from_loop(self):
+        result = []
+
+        def func(fut):
+            try:
+                value = aiogreen.link_future(fut)
+            except Exception as exc:
+                result.append('error')
+            else:
+                result.append(value)
+            self.loop.stop()
+
+        fut = asyncio.Future(loop=self.loop)
+        self.loop.call_soon(func, fut)
+        self.loop.call_soon(fut.set_result, 'unused')
+        self.loop.run_forever()
+        self.assertEqual(result, ['error'])
+
+
 class WrapGreenthreadTests(tests.TestCase):
     def test_wrap_greenthread(self):
         def func():