Make PyModule instance available on CinderX ModuleState

Summary: Introducing this so we can bump the reference count on the module if something is alive and relies on its content. Specifically the memory backing the free-list for generators to start with.

Reviewed By: alexmalyshev

Differential Revision: D78516908

fbshipit-source-id: a3e3700b96f60ea4c4cbce5cf502069f4669404a

Author: jbower-fb

cinderx/_cinderx-lib.cpp

@@ -1065,7 +1065,6 @@ void module_free(void* raw_mod) {
   }
 
   state->shutdown();
-  cinderx::setModuleState(nullptr);
 }
 
 // Called when the interpreter is shutting down, allows us to do some aggressive
@@ -1337,7 +1336,7 @@ static int _cinderx_exec(PyObject* m) {
   }
   state->setSymbolizer(symbolizer);
 
-  cinderx::setModuleState(state);
+  cinderx::setModule(m);
 
   CiExc_StaticTypeError =
       PyErr_NewException("cinderx.StaticTypeError", PyExc_TypeError, nullptr);

cinderx/module_state.cpp

@@ -2,6 +2,8 @@
 
 #include "cinderx/module_state.h"
 
+#include "cinderx/Common/log.h"
+
 namespace cinderx {
 
 int ModuleState::traverse(visitproc visit, void* arg) {
@@ -15,16 +17,25 @@ int ModuleState::clear() {
   return 0;
 }
 
+static ModuleState* s_cinderx_state;
+
 void ModuleState::shutdown() {
   cache_manager_.reset();
   runtime_.reset();
   symbolizer_.reset();
+  JIT_DCHECK(
+      this == s_cinderx_state,
+      "Global module state pointer inconsistent with this module state {} != "
+      "{}",
+      reinterpret_cast<void*>(this),
+      reinterpret_cast<void*>(s_cinderx_state));
+  s_cinderx_state = nullptr;
 }
 
-static ModuleState* s_cinderx_state;
-
-void setModuleState(ModuleState* state) {
+void setModule(PyObject* m) {
+  auto state = reinterpret_cast<cinderx::ModuleState*>(PyModule_GetState(m));
   s_cinderx_state = state;
+  state->setModule(m);
 }
 
 ModuleState* getModuleState() {

cinderx/module_state.h

@@ -136,6 +136,17 @@ class ModuleState {
     initialized_ = init;
   }
 
+  void setModule(BorrowedRef<> module) {
+    cinderx_module_ = module;
+  }
+
+  // Returns the PyModule instance for the CinderX module. This can be useful if
+  // we have live data backed by the module, in which case we can increase the
+  // refcount of the module to prevent it from being freed prematurely.
+  BorrowedRef<> module() const {
+    return cinderx_module_;
+  }
+
  private:
   std::unique_ptr<jit::IGlobalCacheManager> cache_manager_;
   std::unique_ptr<jit::ICodeAllocator> code_allocator_;
@@ -152,9 +163,10 @@ class ModuleState {
   jit::UnorderedSet<BorrowedRef<PyFunctionObject>> perf_trampoline_worklist_;
 
   bool initialized_{false};
+  BorrowedRef<> cinderx_module_;
 };
 
-void setModuleState(ModuleState* state);
+void setModule(PyObject* module);
 ModuleState* getModuleState();
 
 } // namespace cinderx