Avoid acquiring the GIL at the end of nogil functions (GH-3556) (GH-4749)

Closes https://github.com/cython/cython/issues/4637
See See https://github.com/cython/cython/issues/3556

* Acquire the GIL in nogil functions only when strictly needed on function exit, e.g. for cleaning up temp variables from with-gil blocks or adding tracebacks.

See https://github.com/cython/cython/issues/3554

* Make the GIL-avoidance in 7d99b0f0 actually work in nogil functions and not just nogil sections.

See https://github.com/cython/cython/issues/3558

5 files changed, 81 insertions, 35 deletions

diff --git a/Cython/Compiler/Code.py b/Cython/Compiler/Code.py
index 45c5a325c..ed7dceb35 100644
--- a/Cython/Compiler/Code.py
+++ b/Cython/Compiler/Code.py
@@ -2407,8 +2407,8 @@ class CCodeWriter(object):
                 UtilityCode.load_cached("ForceInitThreads", "ModuleSetupCode.c"))
         self.putln('__Pyx_RefNannySetupContext("%s", %d);' % (name, acquire_gil and 1 or 0))
 
-    def put_finish_refcount_context(self):
-        self.putln("__Pyx_RefNannyFinishContext();")
+    def put_finish_refcount_context(self, nogil=False):
+        self.putln("__Pyx_RefNannyFinishContextNogil()" if nogil else "__Pyx_RefNannyFinishContext();")
 
     def put_add_traceback(self, qualified_name, include_cline=True):
         """
diff --git a/Cython/Compiler/Nodes.py b/Cython/Compiler/Nodes.py
index 339b1fa04..01ef96db2 100644
--- a/Cython/Compiler/Nodes.py
+++ b/Cython/Compiler/Nodes.py
@@ -1862,11 +1862,12 @@ class FuncDefNode(StatNode, BlockNode):
 
         use_refnanny = not lenv.nogil or lenv.has_with_gil_block
 
+        gilstate_decl = None
         if acquire_gil or acquire_gil_for_var_decls_only:
             code.put_ensure_gil()
             code.funcstate.gil_owned = True
-        elif lenv.nogil and lenv.has_with_gil_block:
-            code.declare_gilstate()
+        else:
+            gilstate_decl = code.insertion_point()
 
         if profile or linetrace:
             if not self.is_generator:
@@ -1989,6 +1990,19 @@ class FuncDefNode(StatNode, BlockNode):
         code.putln("")
         code.putln("/* function exit code */")
 
+        gil_owned = {
+            'success': code.funcstate.gil_owned,
+            'error': code.funcstate.gil_owned,
+            'gil_state_declared': gilstate_decl is None,
+        }
+        def assure_gil(code_path):
+            if not gil_owned[code_path]:
+                if not gil_owned['gil_state_declared']:
+                    gilstate_decl.declare_gilstate()
+                    gil_owned['gil_state_declared'] = True
+                code.put_ensure_gil(declare_gilstate=False)
+                gil_owned[code_path] = True
+
         # ----- Default return value
         if not self.body.is_terminator:
             if self.return_type.is_pyobject:
@@ -1996,8 +2010,10 @@ class FuncDefNode(StatNode, BlockNode):
                 #    lhs = "(PyObject *)%s" % Naming.retval_cname
                 #else:
                 lhs = Naming.retval_cname
+                assure_gil('success')
                 code.put_init_to_py_none(lhs, self.return_type)
-            else:
+            elif not self.return_type.is_memoryviewslice:
+                # memory view structs receive their default value on initialisation
                 val = self.return_type.default_value
                 if val:
                     code.putln("%s = %s;" % (Naming.retval_cname, val))
@@ -2019,6 +2035,7 @@ class FuncDefNode(StatNode, BlockNode):
                 code.globalstate.use_utility_code(restore_exception_utility_code)
                 code.putln("{ PyObject *__pyx_type, *__pyx_value, *__pyx_tb;")
                 code.putln("__Pyx_PyThreadState_declare")
+                assure_gil('error')
                 code.putln("__Pyx_PyThreadState_assign")
                 code.putln("__Pyx_ErrFetch(&__pyx_type, &__pyx_value, &__pyx_tb);")
                 for entry in used_buffer_entries:
@@ -2038,20 +2055,14 @@ class FuncDefNode(StatNode, BlockNode):
                 # code.globalstate.use_utility_code(get_exception_tuple_utility_code)
                 # code.put_trace_exception()
 
-                if lenv.nogil and not lenv.has_with_gil_block:
-                    code.putln("{")
-                    code.put_ensure_gil()
-
+                assure_gil('error')
                 code.put_add_traceback(self.entry.qualified_name)
-
-                if lenv.nogil and not lenv.has_with_gil_block:
-                    code.put_release_ensured_gil()
-                    code.putln("}")
             else:
                 warning(self.entry.pos,
                         "Unraisable exception in function '%s'." %
                         self.entry.qualified_name, 0)
-                code.put_unraisable(self.entry.qualified_name, lenv.nogil)
+                assure_gil('error')
+                code.put_unraisable(self.entry.qualified_name)
             default_retval = self.return_type.default_value
             if err_val is None and default_retval:
                 err_val = default_retval
@@ -2062,19 +2073,33 @@ class FuncDefNode(StatNode, BlockNode):
                 code.putln("__Pyx_pretend_to_initialize(&%s);" % Naming.retval_cname)
 
             if is_getbuffer_slot:
+                assure_gil('error')
                 self.getbuffer_error_cleanup(code)
 
             # If we are using the non-error cleanup section we should
             # jump past it if we have an error. The if-test below determine
             # whether this section is used.
             if buffers_present or is_getbuffer_slot or self.return_type.is_memoryviewslice:
+                # In the buffer cases, we already called assure_gil('error') and own the GIL.
+                assert gil_owned['error'] or self.return_type.is_memoryviewslice
                 code.put_goto(code.return_from_error_cleanup_label)
+            else:
+                # align error and success GIL state
+                if gil_owned['success']:
+                    assure_gil('error')
+                elif gil_owned['error']:
+                    code.put_release_ensured_gil()
+                    gil_owned['error'] = False
 
         # ----- Non-error return cleanup
         code.put_label(code.return_label)
+        assert gil_owned['error'] == gil_owned['success'], "%s != %s" % (gil_owned['error'], gil_owned['success'])
+
         for entry in used_buffer_entries:
+            assure_gil('success')
             Buffer.put_release_buffer_code(code, entry)
         if is_getbuffer_slot:
+            assure_gil('success')
             self.getbuffer_normal_cleanup(code)
 
         if self.return_type.is_memoryviewslice:
@@ -2084,17 +2109,22 @@ class FuncDefNode(StatNode, BlockNode):
             cond = code.unlikely(self.return_type.error_condition(Naming.retval_cname))
             code.putln(
                 'if (%s) {' % cond)
-            if env.nogil:
+            if not gil_owned['success']:
                 code.put_ensure_gil()
             code.putln(
                 'PyErr_SetString(PyExc_TypeError, "Memoryview return value is not initialized");')
-            if env.nogil:
+            if not gil_owned['success']:
                 code.put_release_ensured_gil()
             code.putln(
                 '}')
 
         # ----- Return cleanup for both error and no-error return
-        code.put_label(code.return_from_error_cleanup_label)
+        if code.label_used(code.return_from_error_cleanup_label):
+            # If we came through the success path, then we took the same GIL decisions as for jumping here.
+            # If we came through the error path and did not jump, then we aligned both paths above.
+            # In the end, all paths are aligned from this point on.
+            assert gil_owned['error'] == gil_owned['success'], "%s != %s" % (gil_owned['error'], gil_owned['success'])
+            code.put_label(code.return_from_error_cleanup_label)
 
         for entry in lenv.var_entries:
             if not entry.used or entry.in_closure:
@@ -2121,6 +2151,7 @@ class FuncDefNode(StatNode, BlockNode):
                 code.put_xdecref_memoryviewslice(entry.cname,
                                                  have_gil=not lenv.nogil)
         if self.needs_closure:
+            assure_gil('success')
             code.put_decref(Naming.cur_scope_cname, lenv.scope_class.type)
 
         # ----- Return
@@ -2136,6 +2167,7 @@ class FuncDefNode(StatNode, BlockNode):
         if self.entry.is_special and self.entry.name == "__hash__":
             # Returning -1 for __hash__ is supposed to signal an error
             # We do as Python instances and coerce -1 into -2.
+            assure_gil('success')
             code.putln("if (unlikely(%s == -1) && !PyErr_Occurred()) %s = -2;" % (
                 Naming.retval_cname, Naming.retval_cname))
 
@@ -2145,16 +2177,16 @@ class FuncDefNode(StatNode, BlockNode):
                 # generators are traced when iterated, not at creation
                 if self.return_type.is_pyobject:
                     code.put_trace_return(
-                        Naming.retval_cname, nogil=not code.funcstate.gil_owned)
+                        Naming.retval_cname, nogil=not gil_owned['success'])
                 else:
                     code.put_trace_return(
-                        "Py_None", nogil=not code.funcstate.gil_owned)
+                        "Py_None", nogil=not gil_owned['success'])
 
         if not lenv.nogil:
             # GIL holding function
-            code.put_finish_refcount_context()
+            code.put_finish_refcount_context(nogil=not gil_owned['success'])
 
-        if acquire_gil or (lenv.nogil and lenv.has_with_gil_block):
+        if acquire_gil or (lenv.nogil and gil_owned['success']):
             # release the GIL (note that with-gil blocks acquire it on exit in their EnsureGILNode)
             code.put_release_ensured_gil()
             code.funcstate.gil_owned = False
diff --git a/Cython/Compiler/ParseTreeTransforms.py b/Cython/Compiler/ParseTreeTransforms.py
index 8e0cbc303..24f475120 100644
--- a/Cython/Compiler/ParseTreeTransforms.py
+++ b/Cython/Compiler/ParseTreeTransforms.py
@@ -1859,19 +1859,6 @@ if VALUE is not None:
 
         return node
 
-    def _handle_nogil_cleanup(self, lenv, node):
-        "Handle cleanup for 'with gil' blocks in nogil functions."
-        if lenv.nogil and lenv.has_with_gil_block:
-            # Acquire the GIL for cleanup in 'nogil' functions, by wrapping
-            # the entire function body in try/finally.
-            # The corresponding release will be taken care of by
-            # Nodes.FuncDefNode.generate_function_definitions()
-            node.body = Nodes.NogilTryFinallyStatNode(
-                node.body.pos,
-                body=node.body,
-                finally_clause=Nodes.EnsureGILNode(node.body.pos),
-                finally_except_clause=Nodes.EnsureGILNode(node.body.pos))
-
     def _handle_fused(self, node):
         if node.is_generator and node.has_fused_arguments:
             node.has_fused_arguments = False
@@ -1912,7 +1899,6 @@ if VALUE is not None:
             node = self._create_fused_function(env, node)
         else:
             node.body.analyse_declarations(lenv)
-            self._handle_nogil_cleanup(lenv, node)
             self._super_visit_FuncDefNode(node)
 
         self.seen_vars_stack.pop()
diff --git a/Cython/Utility/ModuleSetupCode.c b/Cython/Utility/ModuleSetupCode.c
index dee40fb5d..81d9b2c03 100644
--- a/Cython/Utility/ModuleSetupCode.c
+++ b/Cython/Utility/ModuleSetupCode.c
@@ -1284,6 +1284,11 @@ static CYTHON_INLINE int __Pyx_Is_Little_Endian(void)
   #define __Pyx_RefNannySetupContext(name, acquire_gil) \
           __pyx_refnanny = __Pyx_RefNanny->SetupContext((name), __LINE__, __FILE__)
 #endif
+  #define __Pyx_RefNannyFinishContextNogil() { \
+              PyGILState_STATE __pyx_gilstate_save = PyGILState_Ensure(); \
+              __Pyx_RefNannyFinishContext(); \
+              PyGILState_Release(__pyx_gilstate_save); \
+          }
   #define __Pyx_RefNannyFinishContext() \
           __Pyx_RefNanny->FinishContext(&__pyx_refnanny)
   #define __Pyx_INCREF(r)  __Pyx_RefNanny->INCREF(__pyx_refnanny, (PyObject *)(r), __LINE__)
@@ -1297,6 +1302,7 @@ static CYTHON_INLINE int __Pyx_Is_Little_Endian(void)
 #else
   #define __Pyx_RefNannyDeclarations
   #define __Pyx_RefNannySetupContext(name, acquire_gil)
+  #define __Pyx_RefNannyFinishContextNogil()
   #define __Pyx_RefNannyFinishContext()
   #define __Pyx_INCREF(r) Py_INCREF(r)
   #define __Pyx_DECREF(r) Py_DECREF(r)
diff --git a/tests/run/trace_nogil.pyx b/tests/run/trace_nogil.pyx
new file mode 100644
index 000000000..dee443e5b
--- /dev/null
+++ b/tests/run/trace_nogil.pyx
@@ -0,0 +1,22 @@
+# cython: linetrace=True
+
+cdef void foo(int err) nogil except *:
+    with gil:
+        raise ValueError(err)
+
+
+# Test from gh-4637
+def handler(int err):
+    """
+    >>> handler(0)
+    All good
+    >>> handler(1)    # doctest: +ELLIPSIS
+    Traceback (most recent call last):
+        ...
+    ValueError: 1
+    """
+    if (err % 2):
+        with nogil:
+            foo(err)
+    else:
+        print("All good")