From 0beaf42ca341b5d0774c7873e6557dba562a1fb3 Mon Sep 17 00:00:00 2001 From: Dmitri Plotnikov Date: Thu, 20 Aug 2026 18:55:31 -0700 Subject: [PATCH] Add multi-threaded compilation and evaluation tests and resolve compilation deadlock in CEL Python. This change adds cel_parallel_test.py to benchmark and validate compiling and evaluating a diverse set of CEL expressions concurrently across multiple worker threads using concurrent.futures.ThreadPoolExecutor as well as sequentially. In addition, this resolves an AB-BA deadlock between the Python GIL and Protobuf DescriptorPool C++ mutex during concurrent compilation by releasing the GIL in PyCelEnv::Compile and acquiring the GIL via py::gil_scoped_acquire on-demand in PyDescriptorDatabase callbacks, while protecting PyCelEnvInternal with a mutex. Test duration metrics: - Multi-threaded compilation (1,000 iterations): ~248 ms - Sequential compilation (1,000 iterations): ~394 ms - Multi-threaded evaluation (10,000 iterations): ~782 ms - Sequential evaluation (10,000 iterations): ~434 ms PiperOrigin-RevId: 968187521 --- cel_expr_python/BUILD | 16 ++ cel_expr_python/cel_parallel_test.py | 204 ++++++++++++++++++++++ cel_expr_python/py_cel_env.cc | 12 ++ cel_expr_python/py_cel_env_internal.cc | 11 +- cel_expr_python/py_cel_env_internal.h | 9 +- cel_expr_python/py_cel_expression.cc | 2 - cel_expr_python/py_descriptor_database.cc | 12 +- cel_expr_python/py_error_status.cc | 7 +- 8 files changed, 256 insertions(+), 17 deletions(-) create mode 100644 cel_expr_python/cel_parallel_test.py diff --git a/cel_expr_python/BUILD b/cel_expr_python/BUILD index 3b60244..0853e61 100644 --- a/cel_expr_python/BUILD +++ b/cel_expr_python/BUILD @@ -60,6 +60,7 @@ pybind_library( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/synchronization", "@com_google_absl//absl/time", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", @@ -166,6 +167,21 @@ py_test( }), ) +py_test( + name = "cel_parallel_test", + srcs = ["cel_parallel_test.py"], + data = [ + ":cel", + ], + deps = [ + "//testing:proto2_test_all_types_py_pb2", + "@com_google_absl_py//absl/testing:absltest", + ] + select({ + "@platforms//os:windows": [], + "//conditions:default": [":cel"], + }), +) + py_test( name = "cel_env_test", srcs = ["cel_env_test.py"], diff --git a/cel_expr_python/cel_parallel_test.py b/cel_expr_python/cel_parallel_test.py new file mode 100644 index 0000000..479f2aa --- /dev/null +++ b/cel_expr_python/cel_parallel_test.py @@ -0,0 +1,204 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-threaded tests for cel-python.""" + +import collections.abc +import concurrent.futures +import dataclasses +import gc +import logging +import time +from typing import Any + +from absl.testing import absltest +from cel_expr_python import cel +from cel.expr.conformance.proto2 import test_all_types_pb2 as test_all_types_pb + + +@dataclasses.dataclass(frozen=True) +class _TestCase: + expr: str + data: collections.abc.Callable[[int], dict[str, Any]] + expected: collections.abc.Callable[[int], Any] + + +_NUM_EVALUATIONS = 10000 +_NUM_COMPILATIONS = 1000 + +_TEST_MSG = test_all_types_pb.TestAllTypes(single_int64=100) + +_TEST_CASES = [ + _TestCase( + expr="var_int * var_int", + data=lambda n: {"var_int": n}, + expected=lambda n: n * n, + ), + _TestCase( + expr="var_str + '_' + string(var_int)", + data=lambda n: {"var_str": "num", "var_int": n}, + expected=lambda n: f"num_{n}", + ), + _TestCase( + expr="var_int % 2 == 0", + data=lambda n: {"var_int": n}, + expected=lambda n: n % 2 == 0, + ), + _TestCase( + expr="[var_int, var_int + 1, var_int + 2]", + data=lambda n: {"var_int": n}, + expected=lambda n: [n, n + 1, n + 2], + ), + _TestCase( + expr="var_int_map[var_int]", + data=lambda n: {"var_int_map": {n: f"val_{n}"}, "var_int": n}, + expected=lambda n: f"val_{n}", + ), + _TestCase( + expr="var_msg.single_int64 + var_int", + data=lambda n: {"var_msg": _TEST_MSG, "var_int": n}, + expected=lambda n: 100 + n, + ), + _TestCase( + expr=( + "cel.expr.conformance.proto2.TestAllTypes{" + " single_int64: var_int, single_string: var_str" + "}" + ), + data=lambda n: {"var_int": n, "var_str": f"msg_{n}"}, + expected=lambda n: test_all_types_pb.TestAllTypes( + single_int64=n, single_string=f"msg_{n}" + ), + ), + _TestCase( + expr="{'key': var_str, 'value': var_int}", + data=lambda n: {"var_str": f"val_{n}", "var_int": n}, + expected=lambda n: {"key": f"val_{n}", "value": n}, + ), + _TestCase( + expr="[var_int, var_int + 1, var_int + 2].all(x, x >= var_int)", + data=lambda n: {"var_int": n}, + expected=lambda n: True, + ), +] + + +class CelParallelTest(absltest.TestCase): + + def setUp(self): + super().setUp() + + self.env = cel.NewEnv( + variables={ + "var_int": cel.Type.INT, + "var_str": cel.Type.STRING, + "var_int_map": cel.Type.Map(cel.Type.INT, cel.Type.STRING), + "var_msg": cel.Type("cel.expr.conformance.proto2.TestAllTypes"), + }, + ) + self.object_counts_before_test = self._grab_object_counts() + + def tearDown(self): + """Tears down the test environment.""" + super().tearDown() + + gc.collect() + # Assert that all Arenas have been garbage-collected + self.assertEqual(cel._InternalArena._get_instance_count(), 0) + self._check_for_leaks() + + def _grab_object_counts(self) -> dict[str, int]: + gc.collect() + all_objects = gc.get_objects() + type_counts = {} + for obj in all_objects: + obj_type = type(obj) + type_counts[obj_type.__name__] = type_counts.get(obj_type, 0) + 1 + return type_counts + + def _check_for_leaks(self): + type_counts = self._grab_object_counts() + for key, count in type_counts.items(): + if count != self.object_counts_before_test.get(key, 0): + self.fail( + f"Object count for {key} did not match expected count. " + f"Expected: {self.object_counts_before_test.get(key, 0)}, " + f"Actual: {count}", + ) + + def _test_eval(self, multi_threaded: bool): + compiled_exprs = [self.env.compile(tc.expr) for tc in _TEST_CASES] + + def eval_expr(n: int) -> Any: + idx = n % len(_TEST_CASES) + test_case = _TEST_CASES[idx] + expr = compiled_exprs[idx] + data = test_case.data(n) + return expr.eval(data=data).plain_value() + + start_time = time.perf_counter() + if multi_threaded: + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map(eval_expr, range(_NUM_EVALUATIONS))) + else: + results = [eval_expr(n) for n in range(_NUM_EVALUATIONS)] + duration_ms = (time.perf_counter() - start_time) * 1000 + + mode = "Multi-threaded" if multi_threaded else "Sequential" + logging.info("%s evaluation duration: %.2f ms", mode, duration_ms) + + self.assertLen(results, _NUM_EVALUATIONS) + for i, res in enumerate(results): + test_case = _TEST_CASES[i % len(_TEST_CASES)] + self.assertEqual(res, test_case.expected(i)) + + def testMultiThreadedEval(self): + self._test_eval(multi_threaded=True) + + def testSequentialEval(self): + self._test_eval(multi_threaded=False) + + def _test_compile(self, multi_threaded: bool): + def compile_expr(n: int) -> cel.Expression: + test_case = _TEST_CASES[n % len(_TEST_CASES)] + return self.env.compile(test_case.expr) + + start_time = time.perf_counter() + if multi_threaded: + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map(compile_expr, range(_NUM_COMPILATIONS))) + else: + results = [compile_expr(n) for n in range(_NUM_COMPILATIONS)] + duration_ms = (time.perf_counter() - start_time) * 1000 + + mode = "Multi-threaded" if multi_threaded else "Sequential" + logging.info("%s compilation duration: %.2f ms", mode, duration_ms) + + self.assertLen(results, _NUM_COMPILATIONS) + for i, expr in enumerate(results): + test_case = _TEST_CASES[i % len(_TEST_CASES)] + data = test_case.data(i) + self.assertEqual( + expr.eval(data=data).plain_value(), test_case.expected(i) + ) + + def testMultiThreadedCompilation(self): + self._test_compile(multi_threaded=True) + + def testSequentialCompilation(self): + self._test_compile(multi_threaded=False) + + +if __name__ == "__main__": + absltest.main() diff --git a/cel_expr_python/py_cel_env.cc b/cel_expr_python/py_cel_env.cc index 30758cf..68310be 100644 --- a/cel_expr_python/py_cel_env.cc +++ b/cel_expr_python/py_cel_env.cc @@ -196,6 +196,18 @@ std::shared_ptr PyCelEnv::NewActivation( PyCelExpression PyCelEnv::Compile(const std::string& cel_expr, bool disable_check) { + // Release the GIL before entering C++ compilation to prevent lock + // inversion/deadlock with DescriptorPool's internal mutex during concurrent + // multi-threaded compilation. + // + // When DescriptorPool performs a descriptor lookup on a cache miss, it calls + // back into Python via PyDescriptorDatabase (which re-acquires the GIL via + // PyGILState_Ensure). If another thread were to enter Compile() with the GIL + // held, it would block on DescriptorPool's internal C++ mutex while holding + // the GIL, causing an AB-BA deadlock with any thread inside + // PyDescriptorDatabase waiting for the GIL. Releasing the GIL here guarantees + // a strict one-way lock hierarchy (DescriptorPool Mutex -> Python GIL). + py::gil_scoped_release gil_release; return ThrowIfError(PyCelExpression::Compile(env_, cel_expr, disable_check)); } diff --git a/cel_expr_python/py_cel_env_internal.cc b/cel_expr_python/py_cel_env_internal.cc index 8255975..c9440c0 100644 --- a/cel_expr_python/py_cel_env_internal.cc +++ b/cel_expr_python/py_cel_env_internal.cc @@ -25,6 +25,7 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" +#include "absl/synchronization/mutex.h" #include "checker/type_checker_builder.h" #include "common/container.h" #include "common/function_descriptor.h" @@ -40,7 +41,6 @@ #include "runtime/runtime.h" #include "runtime/runtime_builder.h" #include "runtime/runtime_options.h" -#include "validator/validator.h" #include "cel_expr_python/cel_extension.h" #include "cel_expr_python/py_cel_env_config.h" #include "cel_expr_python/py_cel_function.h" @@ -232,8 +232,7 @@ PyCelEnvInternal::NewCelEnvInternal( absl::StatusOr PyCelEnvInternal::GetCompiler( const std::shared_ptr& env) { - ABSL_CHECK(PyGILState_Check()); - + absl::MutexLock lock(env->mutex_); if (env->compiler_) { return env->compiler_.get(); } @@ -279,6 +278,7 @@ absl::StatusOr PyCelEnvInternal::GetCompiler( absl::StatusOr PyCelEnvInternal::GetRuntime( const std::shared_ptr& env, RuntimeMode runtime_mode) { + absl::MutexLock lock(env->mutex_); if (auto it = env->runtimes_.find(runtime_mode); it != env->runtimes_.end()) { return it->second.get(); } @@ -341,7 +341,7 @@ absl::StatusOr PyCelEnvInternal::GetRuntime( const PyCelType& PyCelEnvInternal::GetVariableType( const std::string& name) const { - ABSL_CHECK(PyGILState_Check()); + absl::MutexLock lock(mutex_); auto it = variable_types_.find(name); if (it != variable_types_.end()) { return it->second; @@ -363,9 +363,8 @@ CelExtensionHandle::CelExtensionHandle(CelExtensionHandle&& other) CelExtensionHandle::~CelExtensionHandle() { if (py_extension_ != nullptr) { - auto gil_state = PyGILState_Ensure(); + py::gil_scoped_acquire acquire; Py_DECREF(py_extension_); - PyGILState_Release(gil_state); } } diff --git a/cel_expr_python/py_cel_env_internal.h b/cel_expr_python/py_cel_env_internal.h index 03e3c13..f6087d4 100644 --- a/cel_expr_python/py_cel_env_internal.h +++ b/cel_expr_python/py_cel_env_internal.h @@ -25,6 +25,7 @@ #include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/synchronization/mutex.h" #include "common/container.h" #include "compiler/compiler.h" #include "env/env.h" @@ -117,7 +118,8 @@ class PyCelEnvInternal { // Use NewCelEnvInternal() to create an instance. PyCelEnvInternal( const PyCelEnvConfig& env_config, const PyCelOptions& options, - PyObject* py_descriptor_pool, std::vector extensions, + PyObject* py_descriptor_pool, + std::vector extension_handles, absl::flat_hash_map& function_impls); absl::Status ConfigureStandardExtension( @@ -138,11 +140,12 @@ class PyCelEnvInternal { std::shared_ptr py_message_factory_; // Synchronized by the GIL. absl::flat_hash_map variable_types_; + mutable absl::Mutex mutex_; std::vector extensions_; absl::flat_hash_map function_impls_; - std::unique_ptr compiler_; + std::unique_ptr compiler_ ABSL_GUARDED_BY(mutex_); absl::flat_hash_map> - runtimes_; + runtimes_ ABSL_GUARDED_BY(mutex_); }; } // namespace cel_python diff --git a/cel_expr_python/py_cel_expression.cc b/cel_expr_python/py_cel_expression.cc index bc36244..5fda518 100644 --- a/cel_expr_python/py_cel_expression.cc +++ b/cel_expr_python/py_cel_expression.cc @@ -105,8 +105,6 @@ void PyCelExpression::DefinePythonBindings(py::module& m) { absl::StatusOr PyCelExpression::Compile( const std::shared_ptr& env, const std::string& cel_expr, bool disable_check) { - ABSL_CHECK(PyGILState_Check()); - CEL_PYTHON_ASSIGN_OR_RETURN(const cel::Compiler* compiler, PyCelEnvInternal::GetCompiler(env)); diff --git a/cel_expr_python/py_descriptor_database.cc b/cel_expr_python/py_descriptor_database.cc index ea36039..f6e54ef 100644 --- a/cel_expr_python/py_descriptor_database.cc +++ b/cel_expr_python/py_descriptor_database.cc @@ -25,9 +25,12 @@ #include "common/minimal_descriptor_pool.h" #include "cel_expr_python/py_error_status.h" #include "google/protobuf/descriptor.h" +#include namespace cel_python { +namespace py = pybind11; + PyDescriptorDatabase::PyDescriptorDatabase(PyObject* py_descriptor_pool) : py_descriptor_pool_(py_descriptor_pool), standard_pool_(cel::GetMinimalDescriptorPool()) { @@ -36,16 +39,14 @@ PyDescriptorDatabase::PyDescriptorDatabase(PyObject* py_descriptor_pool) } PyDescriptorDatabase::~PyDescriptorDatabase() { - auto gil_state = PyGILState_Ensure(); + py::gil_scoped_acquire acquire; Py_XDECREF(py_descriptor_pool_); - PyGILState_Release(gil_state); } // Find a file by file name. Fills in in *output and returns true if found. // Otherwise, returns false, leaving the contents of *output undefined. bool PyDescriptorDatabase::FindFileByName(StringViewArg filename, google::protobuf::FileDescriptorProto* output) { - ABSL_CHECK(PyGILState_Check()); const google::protobuf::FileDescriptor* file = standard_pool_.FindFileByName(filename); if (file != nullptr) { file->CopyTo(output); @@ -56,6 +57,7 @@ bool PyDescriptorDatabase::FindFileByName(StringViewArg filename, return false; } + py::gil_scoped_acquire acquire; PyObject* pyfile = PyObject_CallMethod( py_descriptor_pool_, "FindFileByName", "s#", filename.data(), static_cast(filename.size())); @@ -94,7 +96,6 @@ bool PyDescriptorDatabase::FindFileByName(StringViewArg filename, // and leaves *output undefined. bool PyDescriptorDatabase::FindFileContainingSymbol( StringViewArg symbol_name, google::protobuf::FileDescriptorProto* output) { - ABSL_CHECK(PyGILState_Check()); const google::protobuf::FileDescriptor* file = standard_pool_.FindFileContainingSymbol(symbol_name); if (file != nullptr) { @@ -106,6 +107,7 @@ bool PyDescriptorDatabase::FindFileContainingSymbol( return false; } + py::gil_scoped_acquire acquire; PyObject* pyfile = PyObject_CallMethod( py_descriptor_pool_, "FindFileContainingSymbol", "s#", symbol_name.data(), static_cast(symbol_name.size())); @@ -149,7 +151,7 @@ bool PyDescriptorDatabase::FindFileContainingExtension( return false; } - ABSL_CHECK(PyGILState_Check()); + py::gil_scoped_acquire acquire; PyObject* py_containing_type = PyObject_CallMethod( py_descriptor_pool_, "FindMessageTypeByName", "s#", containing_type.data(), static_cast(containing_type.size())); diff --git a/cel_expr_python/py_error_status.cc b/cel_expr_python/py_error_status.cc index 83e03b9..4565c0e 100644 --- a/cel_expr_python/py_error_status.cc +++ b/cel_expr_python/py_error_status.cc @@ -29,6 +29,8 @@ namespace cel_python { +namespace py = pybind11; + static absl::Status PyErrorToStatus(PyObject* py_type, PyObject* py_error) { // Loose mapping from Python exceptions to absl::Status codes, consistent with // the pybind11 mapping. @@ -96,11 +98,13 @@ std::runtime_error StatusToException(const absl::Status& status) { } static absl::Status& PendingPyError() { - static absl::NoDestructor pending_py_error(absl::OkStatus()); + static thread_local absl::NoDestructor pending_py_error( + absl::OkStatus()); return *pending_py_error; } absl::Status PyErr_toStatus() { + py::gil_scoped_acquire acquire; PyObject* py_error = PyErr_Occurred(); if (!py_error) { absl::Status status = PendingPyError(); @@ -132,6 +136,7 @@ absl::Status PyErr_toStatus() { } void PyErr_noteAndClear() { + py::gil_scoped_acquire acquire; if (!PyErr_Occurred()) { return; }