Protobuf Python UPB Free Threading support.

Add obj_cache lock to pass current free threading tests on python upb.

PiperOrigin-RevId: 864903528

Author: anandolee

python/descriptor_pool.c

@@ -49,7 +49,7 @@ static PyObject* PyUpb_DescriptorPool_DoCreateWithCache(
   pool->symtab = upb_DefPool_New();
   pool->db = db;
   Py_XINCREF(pool->db);
-  PyUpb_WeakMap_Add(obj_cache, pool->symtab, &pool->ob_base);
+  PyUpb_KnownObjCache_Add(obj_cache, pool->symtab, &pool->ob_base);
   return &pool->ob_base;
 }
 

python/google/protobuf/internal/thread_safe_test.py

@@ -75,10 +75,6 @@ def DoNothing():
 
     self.RunThreads(thread_size, DoNothing)
 
-  @unittest.skipIf(
-      api_implementation.Type() != 'cpp',
-      'Only cpp supports free threading for now',
-  )
   def testDescriptorPoolMap(self):
     thread_size = 20
     self.success_count = 0

python/protobuf.c

@@ -16,8 +16,44 @@
 #include "python/repeated.h"
 #include "python/unknown_fields.h"
 
+// Mutex wrapper when GIL-disabled. Zero-cost when GIL-enabled.
+// NOTE: Protobuf Free-threading support is still experimental.
+// TODO: Remove mutex and make upb python lock free.
+#ifdef Py_GIL_DISABLED
+#ifdef _POSIX_THREADS
+#define ENABLE_MUTEX 1
+#include <pthread.h>
+#else
+#error "GIL is disabled but _POSIX_THREADS isn't available"
+#endif  // _POSIX_THREADS
+#endif  // Py_GIL_DISABLED
+
+typedef struct {
+#ifdef ENABLE_MUTEX
+  pthread_mutex_t mutex;
+#endif
+} FreeThreadingMutex;
+
+#ifdef ENABLE_MUTEX
+static FreeThreadingMutex obj_cache_mutex = {PTHREAD_MUTEX_INITIALIZER};
+#else
+static FreeThreadingMutex obj_cache_mutex = {};
+#endif
+
 static upb_Arena* PyUpb_NewArena(void);
 
+void FreeThreadingLock(FreeThreadingMutex* thread_mutex) {
+#ifdef ENABLE_MUTEX
+  pthread_mutex_lock(&(thread_mutex->mutex));
+#endif
+}
+
+void FreeThreadingUnlock(FreeThreadingMutex* thread_mutex) {
+#ifdef ENABLE_MUTEX
+  pthread_mutex_unlock(&(thread_mutex->mutex));
+#endif
+}
+
 static void PyUpb_ModuleDealloc(void* module) {
   PyUpb_ModuleState* s = PyModule_GetState(module);
   PyUpb_WeakMap_Free(s->obj_cache);
@@ -197,15 +233,26 @@ void PyUpb_ObjCache_Add(const void* key, PyObject* py_obj) {
   if (!cache) {
     return;
   }
+  FreeThreadingLock(&obj_cache_mutex);
+  PyUpb_WeakMap_Add(cache, key, py_obj);
+  FreeThreadingUnlock(&obj_cache_mutex);
+}
+
+void PyUpb_KnownObjCache_Add(PyUpb_WeakMap* cache, const void* key,
+                             PyObject* py_obj) {
+  FreeThreadingLock(&obj_cache_mutex);
   PyUpb_WeakMap_Add(cache, key, py_obj);
+  FreeThreadingUnlock(&obj_cache_mutex);
 }
 
 void PyUpb_ObjCache_Delete(const void* key) {
   PyUpb_WeakMap* cache = PyUpb_ObjCache_MaybeInstance();
   if (!cache) {
     return;
   }
+  FreeThreadingLock(&obj_cache_mutex);
   PyUpb_WeakMap_Delete(cache, key);
+  FreeThreadingUnlock(&obj_cache_mutex);
 }
 
 PyObject* PyUpb_ObjCache_Get(const void* key) {
@@ -231,7 +278,10 @@ PyObject* PyUpb_ObjCache_Get(const void* key) {
 #endif
     return NULL;
   }
-  return PyUpb_WeakMap_Get(cache, key);
+  FreeThreadingLock(&obj_cache_mutex);
+  PyObject* obj = PyUpb_WeakMap_Get(cache, key);
+  FreeThreadingUnlock(&obj_cache_mutex);
+  return obj;
 }
 
 // -----------------------------------------------------------------------------
@@ -449,7 +499,9 @@ PyMODINIT_FUNC PyInit__message(void) {
 
   state->allow_oversize_protos = false;
   state->wkt_bases = NULL;
+  FreeThreadingLock(&obj_cache_mutex);
   state->obj_cache = PyUpb_WeakMap_New();
+  FreeThreadingUnlock(&obj_cache_mutex);
   state->c_descriptor_symtab = NULL;
 
   if (!PyUpb_InitDescriptorContainers(m) || !PyUpb_InitDescriptorPool(m) ||

python/protobuf.h

@@ -148,6 +148,8 @@ void PyUpb_WeakMap_DeleteIter(PyUpb_WeakMap* map, intptr_t* iter);
 // The object cache is a global WeakMap for mapping upb objects to the
 // corresponding wrapper.
 void PyUpb_ObjCache_Add(const void* key, PyObject* py_obj);
+void PyUpb_KnownObjCache_Add(PyUpb_WeakMap* cache, const void* key,
+                             PyObject* py_obj);
 void PyUpb_ObjCache_Delete(const void* key);
 PyObject* PyUpb_ObjCache_Get(const void* key);  // returns NULL if not present.
 PyUpb_WeakMap* PyUpb_ObjCache_Instance(void);