Skip to content

Commit 91e4c7e

Browse files
dmitriplotnikovcopybara-github
authored andcommitted
[Python] Fix AB-BA deadlock between Python GIL and DescriptorPool mutex
Ensure all C++ locks (Program mutex, Activation mutex, once flags, DescriptorPool mutex) precede Python GIL acquisition by releasing the GIL in PyCelExpression::Eval, PyObjectToCelValue, and PyCelEnvInternal constructor, while acquiring the GIL on callbacks into Python. PiperOrigin-RevId: 988553958
1 parent b85b57d commit 91e4c7e

9 files changed

Lines changed: 148 additions & 87 deletions

‎cel_expr_python/cel_parallel_test.py‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -170,28 +170,28 @@ def testSequentialEval(self):
170170
self._test_eval(multi_threaded=False)
171171

172172
def _test_compile(self, multi_threaded: bool):
173-
def compile_expr(n: int) -> cel.Expression:
173+
174+
def compile_and_eval(n: int) -> Any:
174175
test_case = _TEST_CASES[n % len(_TEST_CASES)]
175-
return self.env.compile(test_case.expr)
176+
expr = self.env.compile(test_case.expr)
177+
data = test_case.data(n)
178+
return expr.eval(data=data).plain_value()
176179

177180
start_time = time.perf_counter()
178181
if multi_threaded:
179182
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
180-
results = list(executor.map(compile_expr, range(_NUM_COMPILATIONS)))
183+
results = list(executor.map(compile_and_eval, range(_NUM_COMPILATIONS)))
181184
else:
182-
results = [compile_expr(n) for n in range(_NUM_COMPILATIONS)]
185+
results = [compile_and_eval(n) for n in range(_NUM_COMPILATIONS)]
183186
duration_ms = (time.perf_counter() - start_time) * 1000
184187

185188
mode = "Multi-threaded" if multi_threaded else "Sequential"
186189
logging.info("%s compilation duration: %.2f ms", mode, duration_ms)
187190

188191
self.assertLen(results, _NUM_COMPILATIONS)
189-
for i, expr in enumerate(results):
192+
for i, res in enumerate(results):
190193
test_case = _TEST_CASES[i % len(_TEST_CASES)]
191-
data = test_case.data(i)
192-
self.assertEqual(
193-
expr.eval(data=data).plain_value(), test_case.expected(i)
194-
)
194+
self.assertEqual(res, test_case.expected(i))
195195

196196
def testMultiThreadedCompilation(self):
197197
self._test_compile(multi_threaded=True)

‎cel_expr_python/py_cel_env_internal.cc‎

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -115,8 +115,13 @@ PyCelEnvInternal::PyCelEnvInternal(
115115
google::protobuf::Arena arena;
116116
for (const cel::Config::VariableConfig& variable_config :
117117
env_config_.GetConfig().GetVariableConfigs()) {
118-
auto status_or_type = cel::TypeInfoToType(variable_config.type_info,
119-
descriptor_pool_.get(), &arena);
118+
auto status_or_type = [&]() {
119+
// Release the GIL during TypeInfoToType lookups to prevent lock order
120+
// inversion with DescriptorPool's internal mutex.
121+
py::gil_scoped_release gil_release;
122+
return cel::TypeInfoToType(variable_config.type_info,
123+
descriptor_pool_.get(), &arena);
124+
}();
120125
if (status_or_type.ok()) {
121126
variable_types_[variable_config.name] =
122127
PyCelType::FromCelType(*status_or_type);
@@ -314,13 +319,6 @@ absl::StatusOr<std::unique_ptr<cel::Runtime>> PyCelEnvInternal::BuildRuntime(
314319
if (it == function_impls_.end()) {
315320
continue;
316321
}
317-
py::object py_function;
318-
if (!PyGILState_Check()) {
319-
py::gil_scoped_acquire acquire;
320-
py_function = it->second;
321-
} else {
322-
py_function = it->second;
323-
}
324322
std::vector<cel::Kind> param_kinds;
325323
param_kinds.reserve(overload_config.parameters.size());
326324
for (const cel::Config::TypeInfo& parameter :
@@ -338,10 +336,9 @@ absl::StatusOr<std::unique_ptr<cel::Runtime>> PyCelEnvInternal::BuildRuntime(
338336
cel::TypeInfoToType(overload_config.return_type,
339337
descriptor_pool_.get(), &arena));
340338
CEL_PYTHON_RETURN_IF_ERROR(builder.function_registry().Register(
341-
descriptor,
342-
std::make_unique<PyCelFunctionAdapter>(
343-
function_config.name, PyCelType::FromCelType(return_type),
344-
std::move(py_function))));
339+
descriptor, std::make_unique<PyCelFunctionAdapter>(
340+
function_config.name,
341+
PyCelType::FromCelType(return_type), it->second)));
345342
}
346343
}
347344
return std::move(builder).Build();

‎cel_expr_python/py_cel_expression.cc‎

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
#include "absl/status/status.h"
3434
#include "absl/status/statusor.h"
3535
#include "absl/strings/str_format.h"
36+
#include "absl/synchronization/mutex.h"
3637
#include "checker/validation_result.h"
3738
#include "common/ast.h"
3839
#include "common/ast_proto.h"
@@ -43,7 +44,6 @@
4344
#include "parser/parser_interface.h"
4445
#include "runtime/embedder_context.h"
4546
#include "runtime/runtime.h"
46-
#include "cel_expr_python/free_threading_mutex.h"
4747
#include "cel_expr_python/py_cel_activation.h"
4848
#include "cel_expr_python/py_cel_arena.h"
4949
#include "cel_expr_python/py_cel_env_internal.h"
@@ -129,7 +129,8 @@ absl::StatusOr<PyCelExpression> PyCelExpression::Compile(
129129
}
130130

131131
PyCelExpression::PyCelExpression(PyCelExpression&& other) noexcept {
132-
FreeThreadingLockGuard lock(other.mutex_);
132+
// Lock other.mutex_ to safely move the Program and expression state.
133+
absl::MutexLock lock(other.mutex_);
133134
expr_ = std::move(other.expr_);
134135
env_ = std::move(other.env_);
135136
cel_program_ = std::move(other.cel_program_);
@@ -155,7 +156,10 @@ PyCelType PyCelExpression::GetReturnType() {
155156
}
156157

157158
absl::StatusOr<const cel::Program*> PyCelExpression::GetProgram() {
158-
FreeThreadingLockGuard lock(mutex_);
159+
// Lock is needed in both GIL-enabled and free-threaded builds because Eval()
160+
// releases the Python GIL before calling GetProgram(), allowing multiple
161+
// threads to concurrently initialize cel_program_.
162+
absl::MutexLock lock(mutex_);
159163
if (cel_program_) {
160164
return cel_program_.get();
161165
}
@@ -179,17 +183,26 @@ absl::StatusOr<const cel::Program*> PyCelExpression::GetProgram() {
179183
absl::StatusOr<PyCelValue> PyCelExpression::Eval(
180184
const PyCelActivation& activation) {
181185
ABSL_CHECK(PyGILState_Check());
182-
CEL_PYTHON_ASSIGN_OR_RETURN(const cel::Program* program, GetProgram());
183186
std::shared_ptr<PyCelArena> arena = activation.GetArena();
184187
std::shared_ptr<PyCelEnvInternal> env = activation.GetEnv();
185-
cel::EmbedderContext embedder_context = cel::EmbedderContext::From(&env);
186-
cel::EvaluateOptions options;
187-
options.message_factory = env->GetMessageFactory();
188-
options.embedder_context = &embedder_context;
189-
CEL_PYTHON_ASSIGN_OR_RETURN(
190-
cel::Value result,
191-
program->Evaluate(arena->GetArena(), *activation.GetActivation(),
192-
std::move(options)));
188+
cel::Value result;
189+
{
190+
// Release the GIL before entering C++ program creation and evaluation to
191+
// prevent lock inversion/deadlock with DescriptorPool's internal mutex
192+
// (C++ Locks -> DescriptorPool Mutex -> Python GIL). Callbacks into Python
193+
// (such as PyCelValueProvider::Provide and PyCelFunctionAdapter::Invoke)
194+
// re-acquire the GIL on demand.
195+
py::gil_scoped_release gil_release;
196+
CEL_PYTHON_ASSIGN_OR_RETURN(const cel::Program* program, GetProgram());
197+
cel::EmbedderContext embedder_context = cel::EmbedderContext::From(&env);
198+
cel::EvaluateOptions options;
199+
options.message_factory = env->GetMessageFactory();
200+
options.embedder_context = &embedder_context;
201+
CEL_PYTHON_ASSIGN_OR_RETURN(
202+
result,
203+
program->Evaluate(arena->GetArena(), *activation.GetActivation(),
204+
std::move(options)));
205+
}
193206
return PyCelValue(result, arena, std::move(env));
194207
}
195208

‎cel_expr_python/py_cel_expression.h‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@
2424
#include "cel/expr/syntax.pb.h"
2525
#include "absl/base/thread_annotations.h"
2626
#include "absl/status/statusor.h"
27+
#include "absl/synchronization/mutex.h"
2728
#include "runtime/runtime.h"
28-
#include "cel_expr_python/free_threading_mutex.h"
2929
#include "cel_expr_python/py_cel_activation.h"
3030
#include "cel_expr_python/py_cel_type.h"
3131
#include "cel_expr_python/py_cel_value.h"
@@ -70,7 +70,13 @@ class PyCelExpression {
7070
std::variant<cel::expr::ParsedExpr, cel::expr::CheckedExpr>
7171
expr_;
7272
std::shared_ptr<PyCelEnvInternal> env_;
73-
mutable FreeThreadingMutex mutex_;
73+
// mutex_ protects lazy initialization of cel_program_ in GetProgram().
74+
// Unlike FreeThreadingMutex (which compiles to a no-op in GIL builds), a real
75+
// absl::Mutex is required here because Eval() releases the Python GIL before
76+
// calling GetProgram(), allowing multiple threads in both GIL-enabled and
77+
// free-threaded Python builds to concurrently access and initialize
78+
// cel_program_.
79+
mutable absl::Mutex mutex_;
7480
std::unique_ptr<cel::Program> cel_program_ ABSL_GUARDED_BY(mutex_);
7581
};
7682

‎cel_expr_python/py_cel_function.cc‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -75,15 +75,33 @@ PyCelFunction::PyCelFunction(std::string function_name,
7575

7676
PyCelFunctionAdapter::PyCelFunctionAdapter(std::string function_name,
7777
PyCelType return_type,
78-
py::object py_function)
78+
const py::object& py_function)
7979
: function_name_(std::move(function_name)),
8080
return_type_(std::move(return_type)),
81-
py_function_(std::move(py_function)) {}
81+
py_function_(py_function.ptr()) {
82+
if (!PyGILState_Check()) {
83+
py::gil_scoped_acquire acquire;
84+
Py_XINCREF(py_function_);
85+
} else {
86+
Py_XINCREF(py_function_);
87+
}
88+
}
89+
90+
PyCelFunctionAdapter::~PyCelFunctionAdapter() {
91+
if (py_function_ != nullptr) {
92+
if (!PyGILState_Check()) {
93+
py::gil_scoped_acquire acquire;
94+
Py_XDECREF(py_function_);
95+
} else {
96+
Py_XDECREF(py_function_);
97+
}
98+
}
99+
}
82100

83101
absl::StatusOr<cel::Value> PyCelFunctionAdapter::Invoke(
84102
absl::Span<const cel::Value> args,
85103
const cel::Function::InvokeContext& context) const {
86-
ABSL_CHECK(PyGILState_Check());
104+
py::gil_scoped_acquire acquire;
87105

88106
std::shared_ptr<PyCelEnvInternal> env = GetEnvFromContext(context);
89107
CEL_PYTHON_ASSIGN_OR_RETURN(auto py_arena,
@@ -94,7 +112,7 @@ absl::StatusOr<cel::Value> PyCelFunctionAdapter::Invoke(
94112
CelValueToPyObject(args[i], env, py_arena,
95113
/*plain_value=*/true));
96114
}
97-
PyObject* result = PyObject_CallObject(py_function_.ptr(), py_args);
115+
PyObject* result = PyObject_CallObject(py_function_, py_args);
98116
Py_DECREF(py_args);
99117
absl::Status status = PyErr_toStatus();
100118
if (!status.ok()) {
@@ -105,9 +123,8 @@ absl::StatusOr<cel::Value> PyCelFunctionAdapter::Invoke(
105123
absl::StatusOr<cel::Value> cel_result = PyObjectToCelValue(
106124
result, return_type_,
107125
[this]() {
108-
return absl::StrFormat(
109-
"Python function '%s'",
110-
PyUnicode_AsUTF8(PyObject_Repr(py_function_.ptr())));
126+
return absl::StrFormat("Python function '%s'",
127+
PyUnicode_AsUTF8(PyObject_Repr(py_function_)));
111128
},
112129
env, context.arena());
113130
Py_XDECREF(result);

‎cel_expr_python/py_cel_function.h‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ class PyCelFunction {
4545
std::string function_name() const { return function_name_; }
4646
const std::vector<PyCelType>& parameters() const { return parameters_; }
4747
bool is_member() const { return is_member_; }
48-
py::object impl() const { return impl_; }
48+
const py::object& impl() const { return impl_; }
4949
const PyCelType& return_type() const { return return_type_; }
5050

5151
private:
@@ -61,7 +61,11 @@ class PyCelFunction {
6161
class PyCelFunctionAdapter : public cel::Function {
6262
public:
6363
PyCelFunctionAdapter(std::string function_name, PyCelType return_type,
64-
py::object py_function);
64+
const py::object& py_function);
65+
~PyCelFunctionAdapter() override;
66+
67+
PyCelFunctionAdapter(const PyCelFunctionAdapter&) = delete;
68+
PyCelFunctionAdapter& operator=(const PyCelFunctionAdapter&) = delete;
6569

6670
absl::StatusOr<cel::Value> Invoke(
6771
absl::Span<const cel::Value> args,
@@ -70,7 +74,7 @@ class PyCelFunctionAdapter : public cel::Function {
7074
private:
7175
std::string function_name_;
7276
PyCelType return_type_;
73-
py::object py_function_;
77+
PyObject* py_function_;
7478
};
7579

7680
} // namespace cel_python

‎cel_expr_python/py_cel_overload.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ class PyCelOverload {
4343
PyCelType return_type() const { return return_type_; }
4444
const std::vector<PyCelType>& parameters() const { return parameters_; }
4545
bool is_member() const { return is_member_; }
46-
py::object py_function() const { return py_function_; }
46+
const py::object& py_function() const { return py_function_; }
4747

4848
cel::Config::FunctionOverloadConfig ToFunctionOverloadConfig() const;
4949

0 commit comments

Comments
 (0)