Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit b8bef14

Browse files
anandoleecopybara-github
authored andcommitted
Python Proto Free Threading tests/experimental
Add experimental and simple tests for free threading support on python fast cpp. PiperOrigin-RevId: 835343204
1 parent 0ce3231 commit b8bef14

File tree

4 files changed

+168
-69
lines changed

4 files changed

+168
-69
lines changed

‎python/google/protobuf/internal/thread_safe_test.py‎

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,17 @@
77

88
"""Unittest for thread safe"""
99

10+
import sys
1011
import threading
1112
import time
1213
import unittest
1314

14-
from google.protobuf import unittest_pb2
15+
from google.protobuf import descriptor_pb2
16+
from google.protobuf import descriptor_pool
17+
from google.protobuf import message_factory
18+
from google.protobuf.internal import api_implementation
1519

20+
from google.protobuf import unittest_pb2
1621

1722
class ThreadSafeTest(unittest.TestCase):
1823

@@ -35,7 +40,7 @@ def ParseMessage():
3540
field_des = unittest_pb2.TestAllTypes.DESCRIPTOR.fields_by_name[
3641
'optional_int32'
3742
]
38-
count = 5000
43+
count = 1000
3944
for x in range(0, count):
4045
# delete the _decoders because only the first time parse the field
4146
# may cause data race.
@@ -51,5 +56,57 @@ def ParseMessage():
5156
self.assertEqual(count * 2, self.success)
5257

5358

59+
class FreeThreadingTest(unittest.TestCase):
60+
61+
def RunThreads(self, thread_size, func):
62+
threads = []
63+
for i in range(0, thread_size):
64+
threads.append(threading.Thread(target=func))
65+
for thread in threads:
66+
thread.start()
67+
for thread in threads:
68+
thread.join()
69+
70+
def testDoNothing(self):
71+
thread_size = 10
72+
73+
def DoNothing():
74+
return
75+
76+
self.RunThreads(thread_size, DoNothing)
77+
78+
@unittest.skipIf(
79+
api_implementation.Type() != 'cpp',
80+
'Only cpp supports free threading for now',
81+
)
82+
def testDescriptorPoolMap(self):
83+
thread_size = 20
84+
self.success_count = 0
85+
lock = threading.Lock()
86+
87+
def CreatePool():
88+
def DoCreate():
89+
pool = descriptor_pool.DescriptorPool()
90+
file_proto = descriptor_pb2.FileDescriptorProto(name='foo')
91+
message_proto = file_proto.message_type.add(name='SomeMessage')
92+
message_proto.field.add(
93+
name='int_field',
94+
number=1,
95+
type=descriptor_pb2.FieldDescriptorProto.TYPE_INT32,
96+
label=descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL,
97+
)
98+
pool.Add(file_proto)
99+
desc = pool.FindMessageTypeByName('SomeMessage')
100+
msg = message_factory.GetMessageClass(desc)()
101+
msg.int_field = 1
102+
103+
DoCreate()
104+
with lock:
105+
self.success_count += 1
106+
107+
self.RunThreads(thread_size, CreatePool)
108+
self.assertEqual(thread_size, self.success_count)
109+
110+
54111
if __name__ == '__main__':
55112
unittest.main()

‎python/google/protobuf/pyext/descriptor.cc‎

Lines changed: 1 addition & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,13 @@
2525
#include "absl/container/flat_hash_map.h"
2626
#include "absl/log/absl_check.h"
2727
#include "absl/strings/string_view.h"
28-
#ifdef Py_GIL_DISABLED
29-
// Only include mutex for free-threaded builds
30-
#include "absl/synchronization/mutex.h"
31-
#endif
3228
#include "google/protobuf/descriptor.h"
3329
#include "google/protobuf/dynamic_message.h"
3430
#include "google/protobuf/internal_feature_helper.h"
3531
#include "google/protobuf/io/coded_stream.h"
3632
#include "google/protobuf/pyext/descriptor_containers.h"
3733
#include "google/protobuf/pyext/descriptor_pool.h"
34+
#include "google/protobuf/pyext/free_threading_mutex.h"
3835
#include "google/protobuf/pyext/message.h"
3936
#include "google/protobuf/pyext/message_factory.h"
4037
#include "google/protobuf/pyext/scoped_pyobject_ptr.h"
@@ -78,51 +75,6 @@ namespace google {
7875
namespace protobuf {
7976
namespace python {
8077

81-
// Zero-cost mutex wrapper that compiles away to nothing in GIL-enabled builds.
82-
// Similar to nanobind's ft_mutex pattern.
83-
class ABSL_LOCKABLE ABSL_ATTRIBUTE_WARN_UNUSED FreeThreadingMutex {
84-
public:
85-
FreeThreadingMutex() = default;
86-
explicit constexpr FreeThreadingMutex(absl::ConstInitType)
87-
#ifdef Py_GIL_DISABLED
88-
: mutex_(absl::kConstInit)
89-
#endif
90-
{
91-
}
92-
FreeThreadingMutex(const FreeThreadingMutex&) = delete;
93-
FreeThreadingMutex& operator=(const FreeThreadingMutex&) = delete;
94-
95-
#ifndef Py_GIL_DISABLED
96-
// GIL-enabled build: no-op mutex (zero cost)
97-
void Lock() {}
98-
void Unlock() {}
99-
#else
100-
// Free-threaded build: real mutex
101-
void Lock() ABSL_EXCLUSIVE_LOCK_FUNCTION() { mutex_.Lock(); }
102-
void Unlock() ABSL_UNLOCK_FUNCTION() { mutex_.Unlock(); }
103-
104-
private:
105-
absl::Mutex mutex_;
106-
#endif
107-
};
108-
109-
// RAII lock guard for FreeThreadingMutex
110-
class ABSL_SCOPED_LOCKABLE FreeThreadingLockGuard {
111-
public:
112-
explicit FreeThreadingLockGuard(FreeThreadingMutex& mutex)
113-
ABSL_EXCLUSIVE_LOCK_FUNCTION(mutex)
114-
: mutex_(mutex) {
115-
mutex_.Lock();
116-
}
117-
~FreeThreadingLockGuard() ABSL_UNLOCK_FUNCTION() { mutex_.Unlock(); }
118-
119-
FreeThreadingLockGuard(const FreeThreadingLockGuard&) = delete;
120-
FreeThreadingLockGuard& operator=(const FreeThreadingLockGuard&) = delete;
121-
122-
private:
123-
FreeThreadingMutex& mutex_;
124-
};
125-
12678
// Mutex to protect interned_descriptors from concurrent access in
12779
// free-threading Python builds. Zero-cost in GIL-enabled builds.
12880
// NOTE: Free-threading support is still experimental.

‎python/google/protobuf/pyext/descriptor_pool.cc‎

Lines changed: 35 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
#include <utility>
1212
#include <vector>
1313

14+
#include "absl/base/const_init.h"
15+
1416
#define PY_SSIZE_T_CLEAN
1517
#include <Python.h>
1618

@@ -22,6 +24,7 @@
2224
#include "google/protobuf/pyext/descriptor.h"
2325
#include "google/protobuf/pyext/descriptor_database.h"
2426
#include "google/protobuf/pyext/descriptor_pool.h"
27+
#include "google/protobuf/pyext/free_threading_mutex.h"
2528
#include "google/protobuf/pyext/message.h"
2629
#include "google/protobuf/pyext/message_factory.h"
2730
#include "google/protobuf/pyext/scoped_pyobject_ptr.h"
@@ -46,6 +49,8 @@ namespace python {
4649
static absl::flat_hash_map<const DescriptorPool*, PyDescriptorPool*>*
4750
descriptor_pool_map;
4851

52+
static FreeThreadingMutex descriptor_pool_map_mutex(absl::kConstInit);
53+
4954
namespace cdescriptor_pool {
5055

5156
// Collects errors that occur during proto file building to allow them to be
@@ -127,11 +132,14 @@ static PyDescriptorPool* PyDescriptorPool_NewWithUnderlay(
127132
cpool->is_mutable = true;
128133
cpool->underlay = underlay;
129134

130-
if (!descriptor_pool_map->insert(
131-
std::make_pair(cpool->pool, cpool)).second) {
132-
// Should never happen -- would indicate an internal error / bug.
133-
PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
134-
return nullptr;
135+
{
136+
FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
137+
if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool))
138+
.second) {
139+
// Should never happen -- would indicate an internal error / bug.
140+
PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
141+
return nullptr;
142+
}
135143
}
136144

137145
return cpool;
@@ -162,10 +170,14 @@ static PyDescriptorPool* PyDescriptorPool_NewWithDatabase(
162170
cpool->pool = pool;
163171
cpool->is_owned = true;
164172

165-
if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool)).second) {
166-
// Should never happen -- would indicate an internal error / bug.
167-
PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
168-
return nullptr;
173+
{
174+
FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
175+
if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool))
176+
.second) {
177+
// Should never happen -- would indicate an internal error / bug.
178+
PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
179+
return nullptr;
180+
}
169181
}
170182

171183
return cpool;
@@ -191,7 +203,10 @@ static PyObject* New(PyTypeObject* type,
191203

192204
static void Dealloc(PyObject* pself) {
193205
PyDescriptorPool* self = reinterpret_cast<PyDescriptorPool*>(pself);
194-
descriptor_pool_map->erase(self->pool);
206+
{
207+
FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
208+
descriptor_pool_map->erase(self->pool);
209+
}
195210
Py_CLEAR(self->py_message_factory);
196211
for (auto it = self->descriptor_options->begin();
197212
it != self->descriptor_options->end(); ++it) {
@@ -692,9 +707,7 @@ bool InitDescriptorPool() {
692707

693708
// Register this pool to be found for C++-generated descriptors.
694709
descriptor_pool_map->insert(
695-
std::make_pair(DescriptorPool::generated_pool(),
696-
python_generated_pool));
697-
710+
std::make_pair(DescriptorPool::generated_pool(), python_generated_pool));
698711
return true;
699712
}
700713

@@ -712,6 +725,7 @@ PyDescriptorPool* GetDescriptorPool_FromPool(const DescriptorPool* pool) {
712725
pool == DescriptorPool::generated_pool()) {
713726
return python_generated_pool;
714727
}
728+
FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
715729
auto it = descriptor_pool_map->find(pool);
716730
if (it == descriptor_pool_map->end()) {
717731
PyErr_SetString(PyExc_KeyError, "Unknown descriptor pool");
@@ -737,11 +751,14 @@ PyObject* PyDescriptorPool_FromPool(const DescriptorPool* pool) {
737751
cpool->is_owned = false;
738752
cpool->is_mutable = false;
739753
cpool->underlay = nullptr;
740-
741-
if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool)).second) {
742-
// Should never happen -- We already checked the existence above.
743-
PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
744-
return nullptr;
754+
{
755+
FreeThreadingLockGuard lock(descriptor_pool_map_mutex);
756+
if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool))
757+
.second) {
758+
// Should never happen -- We already checked the existence above.
759+
PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
760+
return nullptr;
761+
}
745762
}
746763

747764
return reinterpret_cast<PyObject*>(cpool);
Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
// Protocol Buffers - Google's data interchange format
2+
// Copyright 2025 Google Inc. All rights reserved.
3+
//
4+
// Use of this source code is governed by a BSD-style
5+
// license that can be found in the LICENSE file or at
6+
// https://developers.google.com/open-source/licenses/bsd
7+
8+
#ifndef GOOGLE_PROTOBUF_PYTHON_CPP_FREE_THREADING_MUTEX_H__
9+
#define GOOGLE_PROTOBUF_PYTHON_CPP_FREE_THREADING_MUTEX_H__
10+
11+
#include "absl/base/attributes.h"
12+
#include "absl/base/const_init.h"
13+
#include "absl/base/thread_annotations.h"
14+
#ifdef Py_GIL_DISABLED
15+
// Only include mutex for free-threaded builds
16+
#include "absl/synchronization/mutex.h"
17+
#endif
18+
19+
namespace google {
20+
namespace protobuf {
21+
namespace python {
22+
23+
// Zero-cost mutex wrapper that compiles away to nothing in GIL-enabled builds.
24+
// Similar to nanobind's ft_mutex pattern.
25+
// NOTE: Protobuf Free-threading support is still experimental.
26+
class ABSL_LOCKABLE ABSL_ATTRIBUTE_WARN_UNUSED FreeThreadingMutex {
27+
public:
28+
FreeThreadingMutex() = default;
29+
explicit constexpr FreeThreadingMutex(absl::ConstInitType)
30+
#ifdef Py_GIL_DISABLED
31+
: mutex_(absl::kConstInit)
32+
#endif
33+
{
34+
}
35+
FreeThreadingMutex(const FreeThreadingMutex&) = delete;
36+
FreeThreadingMutex& operator=(const FreeThreadingMutex&) = delete;
37+
38+
#ifndef Py_GIL_DISABLED
39+
// GIL-enabled build: no-op mutex (zero cost)
40+
void Lock() {}
41+
void Unlock() {}
42+
#else
43+
// Free-threaded build: real mutex
44+
void Lock() ABSL_EXCLUSIVE_LOCK_FUNCTION() { mutex_.Lock(); }
45+
void Unlock() ABSL_UNLOCK_FUNCTION() { mutex_.Unlock(); }
46+
47+
private:
48+
absl::Mutex mutex_;
49+
#endif
50+
};
51+
52+
// RAII lock guard for FreeThreadingMutex
53+
class ABSL_SCOPED_LOCKABLE FreeThreadingLockGuard {
54+
public:
55+
explicit FreeThreadingLockGuard(FreeThreadingMutex& mutex)
56+
ABSL_EXCLUSIVE_LOCK_FUNCTION(mutex)
57+
: mutex_(mutex) {
58+
mutex_.Lock();
59+
}
60+
~FreeThreadingLockGuard() ABSL_UNLOCK_FUNCTION() { mutex_.Unlock(); }
61+
62+
FreeThreadingLockGuard(const FreeThreadingLockGuard&) = delete;
63+
FreeThreadingLockGuard& operator=(const FreeThreadingLockGuard&) = delete;
64+
65+
private:
66+
FreeThreadingMutex& mutex_;
67+
};
68+
69+
} // namespace python
70+
} // namespace protobuf
71+
} // namespace google
72+
73+
#endif // GOOGLE_PROTOBUF_PYTHON_CPP_FREE_THREADING_MUTEX_H__

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /