Test and fix for #3381

2 files changed, 71 insertions, 1 deletions

diff --git a/lib_pypy/greenlet.py b/lib_pypy/greenlet.py
index 4bea4d1a52..4c25e4986b 100644
--- a/lib_pypy/greenlet.py
+++ b/lib_pypy/greenlet.py
@@ -38,7 +38,8 @@ class greenlet(_continulet):
 
     def __new__(cls, *args, **kwds):
         self = _continulet.__new__(cls)
-        self.parent = getcurrent()
+        self.parent = getcurrent()     # creates '_tls.thread_id' if needed
+        self.__thread_id = _tls.thread_id
         return self
 
     def __init__(self, run=None, parent=None):
@@ -46,6 +47,7 @@ class greenlet(_continulet):
             self.run = run
         if parent is not None:
             self.parent = parent
+            self.__thread_id = parent.__thread_id
 
     def switch(self, *args, **kwds):
         "Switch execution to this greenlet, optionally passing the values "
@@ -58,6 +60,8 @@ class greenlet(_continulet):
 
     def __switch(target, methodname, *baseargs):
         current = getcurrent()
+        if current.__thread_id is not target.__thread_id:
+            raise error("cannot switch to a different thread")
         #
         while not (target.__main or _continulet.is_pending(target)):
             # inlined __nonzero__ ^^^ in case it's overridden
@@ -173,6 +177,7 @@ _tls = _local()
 def _green_create_main():
     # create the main greenlet for this thread
     _tls.current = None
+    _tls.thread_id = object()
     gmain = greenlet.__new__(greenlet)
     gmain._greenlet__main = True
     gmain._greenlet__started = True
diff --git a/pypy/module/test_lib_pypy/test_greenlet_thread.py b/pypy/module/test_lib_pypy/test_greenlet_thread.py
new file mode 100644
index 0000000000..82ea61ed1d
--- /dev/null
+++ b/pypy/module/test_lib_pypy/test_greenlet_thread.py
@@ -0,0 +1,65 @@
+import py
+try:
+    import thread, time
+    from lib_pypy import greenlet
+    #import greenlet
+except ImportError as e:
+    py.test.skip(e)
+
+
+class TestThread:
+
+    def test_cannot_switch_to_main_of_different_thread(self):
+        mains = []
+        mains.append(greenlet.getcurrent())
+        lock = thread.allocate_lock()
+        lock.acquire()
+        got_exception = []
+        #
+        def run_thread():
+            main = greenlet.getcurrent()
+            assert main not in mains
+            mains.append(main)
+            try:
+                mains[0].switch()
+            except Exception as e:
+                got_exception.append(e)
+            lock.release()
+        #
+        thread.start_new_thread(run_thread, ())
+        lock.acquire()
+        assert isinstance(got_exception[0], greenlet.error)
+
+    def test_nonstarted_greenlet_is_still_attached_to_thread(self):
+        subs = []
+        lock = thread.allocate_lock()
+        lock.acquire()
+        got_exception = []
+        #
+        def run_thread():
+            g = greenlet.greenlet(lambda *args: None)
+            subs.append(g)
+            lock.release()
+            time.sleep(1)
+        #
+        thread.start_new_thread(run_thread, ())
+        lock.acquire()
+        [g] = subs
+        py.test.raises(greenlet.error, g.switch)
+
+    def test_noninited_greenlet_is_still_attached_to_thread(self):
+        subs = []
+        lock = thread.allocate_lock()
+        lock.acquire()
+        got_exception = []
+        #
+        def run_thread():
+            g = greenlet.greenlet.__new__(greenlet.greenlet)
+            subs.append(g)
+            lock.release()
+            time.sleep(1)
+        #
+        thread.start_new_thread(run_thread, ())
+        lock.acquire()
+        [g] = subs
+        py.test.raises(greenlet.error, g.switch)