Visitar URL original
[Python] Fix AB-BA deadlock between Python GIL and DescriptorPool mutex · cel-expr/cel-python@37a7bca · GitHub
Skip to content

Commit 37a7bca

Browse files
dmitriplotnikovcopybara-github
dmitriplotnikov
authored and
copybara-github
committed
[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 37a7bca

11 files changed

Lines changed: 281 additions & 97 deletions

‎cel_expr_python/cel_parallel_test.py‎

Lines changed: 40 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,18 @@ def setUp(self):
106106
"var_int_map": cel.Type.Map(cel.Type.INT, cel.Type.STRING),
107107
"var_msg": cel.Type("cel.expr.conformance.proto2.TestAllTypes"),
108108
},
109+
functions=[
110+
cel.FunctionDecl(
111+
"custom_fn",
112+
[
113+
cel.Overload(
114+
"custom_fn_int",
115+
return_type=cel.Type.INT,
116+
parameters=[cel.Type.INT],
117+
)
118+
],
119+
)
120+
],
109121
)
110122
self.object_counts_before_test = self._grab_object_counts()
111123

@@ -170,28 +182,28 @@ def testSequentialEval(self):
170182
self._test_eval(multi_threaded=False)
171183

172184
def _test_compile(self, multi_threaded: bool):
173-
def compile_expr(n: int) -> cel.Expression:
185+
186+
def compile_and_eval(n: int) -> Any:
174187
test_case = _TEST_CASES[n % len(_TEST_CASES)]
175-
return self.env.compile(test_case.expr)
188+
expr = self.env.compile(test_case.expr)
189+
data = test_case.data(n)
190+
return expr.eval(data=data).plain_value()
176191

177192
start_time = time.perf_counter()
178193
if multi_threaded:
179194
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
180-
results = list(executor.map(compile_expr, range(_NUM_COMPILATIONS)))
195+
results = list(executor.map(compile_and_eval, range(_NUM_COMPILATIONS)))
181196
else:
182-
results = [compile_expr(n) for n in range(_NUM_COMPILATIONS)]
197+
results = [compile_and_eval(n) for n in range(_NUM_COMPILATIONS)]
183198
duration_ms = (time.perf_counter() - start_time) * 1000
184199

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

188203
self.assertLen(results, _NUM_COMPILATIONS)
189-
for i, expr in enumerate(results):
204+
for i, res in enumerate(results):
190205
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-
)
206+
self.assertEqual(res, test_case.expected(i))
195207

196208
def testMultiThreadedCompilation(self):
197209
self._test_compile(multi_threaded=True)
@@ -240,6 +252,25 @@ def read_values(_):
240252
for i in range(10):
241253
run_concurrent_value_test(i)
242254

255+
def testConcurrentCustomFunction(self):
256+
expr = self.env.compile("custom_fn(var_int)")
257+
258+
def eval_custom_fn(n: int):
259+
fn = cel.Function(
260+
"custom_fn",
261+
[cel.Type.INT],
262+
False,
263+
lambda x: x * 2,
264+
return_type=cel.Type.INT,
265+
)
266+
act = self.env.Activation({"var_int": n}, functions=[fn])
267+
res = expr.eval(act)
268+
self.assertEqual(res.value(), n * 2)
269+
270+
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
271+
results = list(executor.map(eval_custom_fn, range(100)))
272+
self.assertLen(results, 100)
273+
243274

244275
if __name__ == "__main__":
245276
absltest.main()

‎cel_expr_python/cel_test.py‎

Lines changed: 80 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -846,6 +846,45 @@ def testErrorHandling(self):
846846
r"Could not find file containing symbol:.* \[NOT_FOUND\]",
847847
)
848848

849+
def testDescriptorPoolMissingSerializedPb(self):
850+
bad_env = cel.NewEnv(
851+
_MissingSerializedPbPool(),
852+
variables={},
853+
options=self.options,
854+
)
855+
with self.assertRaises(Exception) as e:
856+
bad_env.compile("cel.expr.conformance.proto2.TestSomeTypes{}")
857+
self.assertIn(
858+
"Python file has no attribute 'serialized_pb'",
859+
str(e.exception),
860+
)
861+
862+
def testDescriptorPoolCorruptedSerializedPb(self):
863+
bad_env = cel.NewEnv(
864+
_CorruptedSerializedPbPool(),
865+
variables={},
866+
options=self.options,
867+
)
868+
with self.assertRaises(Exception) as e:
869+
bad_env.compile("cel.expr.conformance.proto2.TestSomeTypes{}")
870+
self.assertIn(
871+
"Failed to parse descriptor for"
872+
" cel.expr.conformance.proto2.TestSomeTypes",
873+
str(e.exception),
874+
)
875+
876+
def testProtoMessageToCelValueError(self):
877+
bad_env = cel.NewEnv(
878+
_RaisingDescriptorPool(),
879+
variables={"var_proto": cel.Type.DYN},
880+
options=self.options,
881+
)
882+
expr = bad_env.compile("var_proto", disable_check=True)
883+
msg = test_all_types_pb.TestAllTypes(single_string="Hey")
884+
res = expr.eval(data={"var_proto": msg})
885+
self.assertEqual(res.type(), cel.Type.ERROR)
886+
self.assertIn("Custom pool error", str(res.value()))
887+
849888

850889
class CompatibleNumber:
851890

@@ -872,7 +911,47 @@ def __int__(self) -> int:
872911
class _BadDescriptorPool:
873912

874913
def FindFileContainingSymbol(self, symbol_name: str): # pylint: disable=invalid-name
875-
raise LookupError("Could not find file containing symbol: %s" % symbol_name)
914+
if symbol_name.startswith("cel.expr.conformance"):
915+
raise LookupError(
916+
"Could not find file containing symbol: %s" % symbol_name
917+
)
918+
raise KeyError(symbol_name)
919+
920+
921+
class _MissingSerializedPbPool:
922+
923+
def FindFileByName( # pylint: disable=invalid-name,unused-argument
924+
self, filename: str
925+
):
926+
raise KeyError(filename)
927+
928+
def FindFileContainingSymbol( # pylint: disable=invalid-name,unused-argument
929+
self, symbol_name: str
930+
):
931+
if symbol_name.startswith("cel.expr.conformance"):
932+
return object()
933+
raise KeyError(symbol_name)
934+
935+
936+
class _CorruptedSerializedPbPool:
937+
938+
class _FakeFile:
939+
serialized_pb = b"corrupted proto descriptor bytes"
940+
941+
def FindFileContainingSymbol( # pylint: disable=invalid-name,unused-argument
942+
self, symbol_name: str
943+
):
944+
if symbol_name.startswith("cel.expr.conformance"):
945+
return self._FakeFile()
946+
raise KeyError(symbol_name)
947+
948+
949+
class _RaisingDescriptorPool:
950+
951+
def FindFileContainingSymbol(self, symbol_name: str): # pylint: disable=invalid-name
952+
if symbol_name.startswith("cel.expr.conformance"):
953+
raise RuntimeError("Custom pool error: %s" % symbol_name)
954+
raise KeyError(symbol_name)
876955

877956

878957
class CelWithoutProtoSupportTest(absltest.TestCase):

‎cel_expr_python/py_cel_env_internal.cc‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -115,11 +115,16 @@ 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);
120-
if (status_or_type.ok()) {
121-
variable_types_[variable_config.name] =
122-
PyCelType::FromCelType(*status_or_type);
118+
absl::StatusOr<cel::Type> type;
119+
{
120+
// Release the GIL during TypeInfoToType lookups to prevent lock order
121+
// inversion with DescriptorPool's internal mutex.
122+
py::gil_scoped_release gil_release;
123+
type = cel::TypeInfoToType(variable_config.type_info,
124+
descriptor_pool_.get(), &arena);
125+
}
126+
if (type.ok()) {
127+
variable_types_[variable_config.name] = PyCelType::FromCelType(*type);
123128
}
124129
}
125130
}
@@ -132,7 +137,7 @@ PyCelEnvInternal::NewCelEnvInternal(
132137
PyObject* py_descriptor_pool,
133138
const std::unordered_map<std::string, PyCelType>& variable_types,
134139
const std::vector<PyObject*>& extensions,
135-
cel::ExpressionContainer container,
140+
const cel::ExpressionContainer& container,
136141
const std::vector<std::shared_ptr<PyCelFunctionDecl>>& functions,
137142
const std::unordered_map<std::string, py::object>& function_impls) {
138143
cel::Config config = env_config.GetConfig();
@@ -310,17 +315,6 @@ absl::StatusOr<std::unique_ptr<cel::Runtime>> PyCelEnvInternal::BuildRuntime(
310315
GetEnvConfig().GetConfig().GetFunctionConfigs()) {
311316
for (const cel::Config::FunctionOverloadConfig& overload_config :
312317
function_config.overload_configs) {
313-
auto it = function_impls_.find(overload_config.overload_id);
314-
if (it == function_impls_.end()) {
315-
continue;
316-
}
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-
}
324318
std::vector<cel::Kind> param_kinds;
325319
param_kinds.reserve(overload_config.parameters.size());
326320
for (const cel::Config::TypeInfo& parameter :
@@ -333,15 +327,20 @@ absl::StatusOr<std::unique_ptr<cel::Runtime>> PyCelEnvInternal::BuildRuntime(
333327
cel::FunctionDescriptor descriptor(
334328
function_config.name, overload_config.is_member_function, param_kinds,
335329
kFunctionDescriptorOptions);
330+
auto it = function_impls_.find(overload_config.overload_id);
331+
if (it == function_impls_.end()) {
332+
CEL_PYTHON_RETURN_IF_ERROR(
333+
builder.function_registry().RegisterLazyFunction(descriptor));
334+
continue;
335+
}
336336
CEL_PYTHON_ASSIGN_OR_RETURN(
337337
cel::Type return_type,
338338
cel::TypeInfoToType(overload_config.return_type,
339339
descriptor_pool_.get(), &arena));
340340
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))));
341+
descriptor, std::make_unique<PyCelFunctionAdapter>(
342+
function_config.name,
343+
PyCelType::FromCelType(return_type), it->second)));
345344
}
346345
}
347346
return std::move(builder).Build();

‎cel_expr_python/py_cel_env_internal.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ class PyCelEnvInternal {
7676
PyObject* py_descriptor_pool,
7777
const std::unordered_map<std::string, PyCelType>& variable_types,
7878
const std::vector<PyObject*>& extensions,
79-
cel::ExpressionContainer container,
79+
const cel::ExpressionContainer& container,
8080
const std::vector<std::shared_ptr<PyCelFunctionDecl>>& functions,
8181
const std::unordered_map<std::string, py::object>& function_impls);
8282

‎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

0 commit comments

Comments
 (0)