Python Proto Free Threading tests/experimental

Add experimental and simple tests for free threading support on python fast cpp.

PiperOrigin-RevId: 835343204

Author: anandolee

python/google/protobuf/internal/thread_safe_test.py

@@ -7,12 +7,17 @@
 
 """Unittest for thread safe"""
 
+import sys
 import threading
 import time
 import unittest
 
-from google.protobuf import unittest_pb2
+from google.protobuf import descriptor_pb2
+from google.protobuf import descriptor_pool
+from google.protobuf import message_factory
+from google.protobuf.internal import api_implementation
 
+from google.protobuf import unittest_pb2
 
 class ThreadSafeTest(unittest.TestCase):
 
@@ -35,7 +40,7 @@ def ParseMessage():
     field_des = unittest_pb2.TestAllTypes.DESCRIPTOR.fields_by_name[
         'optional_int32'
     ]
-    count = 5000
+    count = 1000
     for x in range(0, count):
       # delete the _decoders because only the first time parse the field
       # may cause data race.
@@ -51,5 +56,57 @@ def ParseMessage():
     self.assertEqual(count * 2, self.success)
 
 
+class FreeThreadingTest(unittest.TestCase):
+
+  def RunThreads(self, thread_size, func):
+    threads = []
+    for i in range(0, thread_size):
+      threads.append(threading.Thread(target=func))
+    for thread in threads:
+      thread.start()
+    for thread in threads:
+      thread.join()
+
+  def testDoNothing(self):
+    thread_size = 10
+
+    def DoNothing():
+      return
+
+    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
+    lock = threading.Lock()
+
+    def CreatePool():
+      def DoCreate():
+        pool = descriptor_pool.DescriptorPool()
+        file_proto = descriptor_pb2.FileDescriptorProto(name='foo')
+        message_proto = file_proto.message_type.add(name='SomeMessage')
+        message_proto.field.add(
+            name='int_field',
+            number=1,
+            type=descriptor_pb2.FieldDescriptorProto.TYPE_INT32,
+            label=descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL,
+        )
+        pool.Add(file_proto)
+        desc = pool.FindMessageTypeByName('SomeMessage')
+        msg = message_factory.GetMessageClass(desc)()
+        msg.int_field = 1
+
+      DoCreate()
+      with lock:
+        self.success_count += 1
+
+    self.RunThreads(thread_size, CreatePool)
+    self.assertEqual(thread_size, self.success_count)
+
+
 if __name__ == '__main__':
   unittest.main()

python/google/protobuf/pyext/descriptor.cc

@@ -25,16 +25,13 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/absl_check.h"
 #include "absl/strings/string_view.h"
-#ifdef Py_GIL_DISABLED
-// Only include mutex for free-threaded builds
-#include "absl/synchronization/mutex.h"
-#endif
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/dynamic_message.h"
 #include "google/protobuf/internal_feature_helper.h"
 #include "google/protobuf/io/coded_stream.h"
 #include "google/protobuf/pyext/descriptor_containers.h"
 #include "google/protobuf/pyext/descriptor_pool.h"
+#include "google/protobuf/pyext/free_threading_mutex.h"
 #include "google/protobuf/pyext/message.h"
 #include "google/protobuf/pyext/message_factory.h"
 #include "google/protobuf/pyext/scoped_pyobject_ptr.h"
@@ -78,51 +75,6 @@ namespace google {
 namespace protobuf {
 namespace python {
 
-// Zero-cost mutex wrapper that compiles away to nothing in GIL-enabled builds.
-// Similar to nanobind's ft_mutex pattern.
-class ABSL_LOCKABLE ABSL_ATTRIBUTE_WARN_UNUSED FreeThreadingMutex {
- public:
-  FreeThreadingMutex() = default;
-  explicit constexpr FreeThreadingMutex(absl::ConstInitType)
-#ifdef Py_GIL_DISABLED
-      : mutex_(absl::kConstInit)
-#endif
-  {
-  }
-  FreeThreadingMutex(const FreeThreadingMutex&) = delete;
-  FreeThreadingMutex& operator=(const FreeThreadingMutex&) = delete;
-
-#ifndef Py_GIL_DISABLED
-  // GIL-enabled build: no-op mutex (zero cost)
-  void Lock() {}
-  void Unlock() {}
-#else
-  // Free-threaded build: real mutex
-  void Lock() ABSL_EXCLUSIVE_LOCK_FUNCTION() { mutex_.Lock(); }
-  void Unlock() ABSL_UNLOCK_FUNCTION() { mutex_.Unlock(); }
-
- private:
-  absl::Mutex mutex_;
-#endif
-};
-
-// RAII lock guard for FreeThreadingMutex
-class ABSL_SCOPED_LOCKABLE FreeThreadingLockGuard {
- public:
-  explicit FreeThreadingLockGuard(FreeThreadingMutex& mutex)
-      ABSL_EXCLUSIVE_LOCK_FUNCTION(mutex)
-      : mutex_(mutex) {
-    mutex_.Lock();
-  }
-  ~FreeThreadingLockGuard() ABSL_UNLOCK_FUNCTION() { mutex_.Unlock(); }
-
-  FreeThreadingLockGuard(const FreeThreadingLockGuard&) = delete;
-  FreeThreadingLockGuard& operator=(const FreeThreadingLockGuard&) = delete;
-
- private:
-  FreeThreadingMutex& mutex_;
-};
-
 // Mutex to protect interned_descriptors from concurrent access in
 // free-threading Python builds. Zero-cost in GIL-enabled builds.
 // NOTE: Free-threading support is still experimental.

python/google/protobuf/pyext/descriptor_pool.cc

@@ -11,6 +11,8 @@
 #include <utility>
 #include <vector>
 
+#include "absl/base/const_init.h"
+
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
 
@@ -22,6 +24,7 @@
 #include "google/protobuf/pyext/descriptor.h"
 #include "google/protobuf/pyext/descriptor_database.h"
 #include "google/protobuf/pyext/descriptor_pool.h"
+#include "google/protobuf/pyext/free_threading_mutex.h"
 #include "google/protobuf/pyext/message.h"
 #include "google/protobuf/pyext/message_factory.h"
 #include "google/protobuf/pyext/scoped_pyobject_ptr.h"
@@ -46,6 +49,8 @@ namespace python {
 static absl::flat_hash_map<const DescriptorPool*, PyDescriptorPool*>*
     descriptor_pool_map;
 
+static FreeThreadingMutex descriptor_pool_map_mutex(absl::kConstInit);
+
 namespace cdescriptor_pool {
 
 // Collects errors that occur during proto file building to allow them to be
@@ -127,11 +132,14 @@ static PyDescriptorPool* PyDescriptorPool_NewWithUnderlay(
   cpool->is_mutable = true;
   cpool->underlay = underlay;
 
-  if (!descriptor_pool_map->insert(
-      std::make_pair(cpool->pool, cpool)).second) {
-    // Should never happen -- would indicate an internal error / bug.
-    PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
-    return nullptr;
+  {
+    FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
+    if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool))
+             .second) {
+      // Should never happen -- would indicate an internal error / bug.
+      PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
+      return nullptr;
+    }
   }
 
   return cpool;
@@ -162,10 +170,14 @@ static PyDescriptorPool* PyDescriptorPool_NewWithDatabase(
   cpool->pool = pool;
   cpool->is_owned = true;
 
-  if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool)).second) {
-    // Should never happen -- would indicate an internal error / bug.
-    PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
-    return nullptr;
+  {
+    FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
+    if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool))
+             .second) {
+      // Should never happen -- would indicate an internal error / bug.
+      PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
+      return nullptr;
+    }
   }
 
   return cpool;
@@ -191,7 +203,10 @@ static PyObject* New(PyTypeObject* type,
 
 static void Dealloc(PyObject* pself) {
   PyDescriptorPool* self = reinterpret_cast<PyDescriptorPool*>(pself);
-  descriptor_pool_map->erase(self->pool);
+  {
+    FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
+    descriptor_pool_map->erase(self->pool);
+  }
   Py_CLEAR(self->py_message_factory);
   for (auto it = self->descriptor_options->begin();
        it != self->descriptor_options->end(); ++it) {
@@ -692,9 +707,7 @@ bool InitDescriptorPool() {
 
   // Register this pool to be found for C++-generated descriptors.
   descriptor_pool_map->insert(
-      std::make_pair(DescriptorPool::generated_pool(),
-                     python_generated_pool));
-
+      std::make_pair(DescriptorPool::generated_pool(), python_generated_pool));
   return true;
 }
 
@@ -712,6 +725,7 @@ PyDescriptorPool* GetDescriptorPool_FromPool(const DescriptorPool* pool) {
       pool == DescriptorPool::generated_pool()) {
     return python_generated_pool;
   }
+  FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
   auto it = descriptor_pool_map->find(pool);
   if (it == descriptor_pool_map->end()) {
     PyErr_SetString(PyExc_KeyError, "Unknown descriptor pool");
@@ -737,11 +751,14 @@ PyObject* PyDescriptorPool_FromPool(const DescriptorPool* pool) {
   cpool->is_owned = false;
   cpool->is_mutable = false;
   cpool->underlay = nullptr;
-
-  if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool)).second) {
-    // Should never happen -- We already checked the existence above.
-    PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
-    return nullptr;
+  {
+    FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
+    if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool))
+             .second) {
+      // Should never happen -- We already checked the existence above.
+      PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
+      return nullptr;
+    }
   }
 
   return reinterpret_cast<PyObject*>(cpool);

python/google/protobuf/pyext/free_threading_mutex.h

@@ -0,0 +1,73 @@
+// Protocol Buffers - Google's data interchange format
+// Copyright 2025 Google Inc.  All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+#ifndef GOOGLE_PROTOBUF_PYTHON_CPP_FREE_THREADING_MUTEX_H__
+#define GOOGLE_PROTOBUF_PYTHON_CPP_FREE_THREADING_MUTEX_H__
+
+#include "absl/base/attributes.h"
+#include "absl/base/const_init.h"
+#include "absl/base/thread_annotations.h"
+#ifdef Py_GIL_DISABLED
+// Only include mutex for free-threaded builds
+#include "absl/synchronization/mutex.h"
+#endif
+
+namespace google {
+namespace protobuf {
+namespace python {
+
+// Zero-cost mutex wrapper that compiles away to nothing in GIL-enabled builds.
+// Similar to nanobind's ft_mutex pattern.
+// NOTE: Protobuf Free-threading support is still experimental.
+class ABSL_LOCKABLE ABSL_ATTRIBUTE_WARN_UNUSED FreeThreadingMutex {
+ public:
+  FreeThreadingMutex() = default;
+  explicit constexpr FreeThreadingMutex(absl::ConstInitType)
+#ifdef Py_GIL_DISABLED
+      : mutex_(absl::kConstInit)
+#endif
+  {
+  }
+  FreeThreadingMutex(const FreeThreadingMutex&) = delete;
+  FreeThreadingMutex& operator=(const FreeThreadingMutex&) = delete;
+
+#ifndef Py_GIL_DISABLED
+  // GIL-enabled build: no-op mutex (zero cost)
+  void Lock() {}
+  void Unlock() {}
+#else
+  // Free-threaded build: real mutex
+  void Lock() ABSL_EXCLUSIVE_LOCK_FUNCTION() { mutex_.Lock(); }
+  void Unlock() ABSL_UNLOCK_FUNCTION() { mutex_.Unlock(); }
+
+ private:
+  absl::Mutex mutex_;
+#endif
+};
+
+// RAII lock guard for FreeThreadingMutex
+class ABSL_SCOPED_LOCKABLE FreeThreadingLockGuard {
+ public:
+  explicit FreeThreadingLockGuard(FreeThreadingMutex& mutex)
+      ABSL_EXCLUSIVE_LOCK_FUNCTION(mutex)
+      : mutex_(mutex) {
+    mutex_.Lock();
+  }
+  ~FreeThreadingLockGuard() ABSL_UNLOCK_FUNCTION() { mutex_.Unlock(); }
+
+  FreeThreadingLockGuard(const FreeThreadingLockGuard&) = delete;
+  FreeThreadingLockGuard& operator=(const FreeThreadingLockGuard&) = delete;
+
+ private:
+  FreeThreadingMutex& mutex_;
+};
+
+}  // namespace python
+}  // namespace protobuf
+}  // namespace google
+
+#endif  // GOOGLE_PROTOBUF_PYTHON_CPP_FREE_THREADING_MUTEX_H__