@@ -3353,6 +3353,53 @@ type_freeze(PyObject *module, PyObject *args)
33533353    Py_RETURN_NONE ;
33543354}
33553355
3356+ struct  atexit_data  {
3357+     int  called ;
3358+     PyThreadState  * tstate ;
3359+     PyInterpreterState  * interp ;
3360+ };
3361+ 
3362+ static  void 
3363+ atexit_callback (void  * data )
3364+ {
3365+     struct  atexit_data  * at_data  =  (struct  atexit_data  * )data ;
3366+     // Ensure that the callback is from the same interpreter 
3367+     assert (PyThreadState_Get () ==  at_data -> tstate );
3368+     assert (PyInterpreterState_Get () ==  at_data -> interp );
3369+     ++ at_data -> called ;
3370+ }
3371+ 
3372+ static  PyObject  * 
3373+ test_atexit (PyObject  * self , PyObject  * Py_UNUSED (args ))
3374+ {
3375+     PyThreadState  * oldts  =  PyThreadState_Swap (NULL );
3376+     PyThreadState  * tstate  =  Py_NewInterpreter ();
3377+ 
3378+     struct  atexit_data  data  =  {0 };
3379+     data .tstate  =  PyThreadState_Get ();
3380+     data .interp  =  PyInterpreterState_Get ();
3381+ 
3382+     int  amount  =  10 ;
3383+     for  (int  i  =  0 ; i  <  amount ; ++ i )
3384+     {
3385+         int  res  =  PyUnstable_AtExit (tstate -> interp , atexit_callback , (void  * )& data );
3386+         if  (res  <  0 ) {
3387+             Py_EndInterpreter (tstate );
3388+             PyThreadState_Swap (oldts );
3389+             PyErr_SetString (PyExc_RuntimeError , "atexit callback failed" );
3390+             return  NULL ;
3391+         }
3392+     }
3393+ 
3394+     Py_EndInterpreter (tstate );
3395+     PyThreadState_Swap (oldts );
3396+ 
3397+     if  (data .called  !=  amount ) {
3398+         PyErr_SetString (PyExc_RuntimeError , "atexit callback not called" );
3399+         return  NULL ;
3400+     }
3401+     Py_RETURN_NONE ;
3402+ }
33563403
33573404static  PyMethodDef  TestMethods [] =  {
33583405    {"set_errno" ,               set_errno ,                       METH_VARARGS },
@@ -3495,6 +3542,7 @@ static PyMethodDef TestMethods[] = {
34953542    {"test_critical_sections" , test_critical_sections , METH_NOARGS },
34963543    {"finalize_thread_hang" , finalize_thread_hang , METH_O , NULL },
34973544    {"type_freeze" , type_freeze , METH_VARARGS },
3545+     {"test_atexit" , test_atexit , METH_NOARGS },
34983546    {NULL , NULL } /* sentinel */ 
34993547};
35003548
0 commit comments