opensource/WrenAI
3
0
Fork 0
mirror of https://github.com/Canner/WrenAI.git synced 2026-09-01 04:57:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

perf(core-py): remove per-context call lock for same-context concurrency (#2666)

Browse Source
This commit is contained in:
Peter
2026-08-17 15:41:18 +08:00
committed by GitHub
parent f29b5b0a01
commit 28920e7e26
4 changed files with 322 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/wren-core-py/README.md
+22 -2
View File
@@ -63,12 +63,32 @@ Visibility contract:
- Brand-new *top-level* catalogs are the exception — they must exist before
MDL construction, `load_mdl`, or a transform, each of which snapshots the
top-level catalog list.
- `load_mdl` must not overlap other calls on the same context; overlapping
calls raise `RuntimeError`.
- For `load_mdl`'s overlap rule, see the Concurrency section below.
For complete runnable examples (fixture files, matching manifests, decoding
the returned bytes), see `tests/test_physical_tables.py`.
### Concurrency
Calls on one `SessionContext` run in parallel. Each `transform_sql` works
on a private top-level catalog snapshot and analyzer state is
per-invocation, so supported concurrent calls never observe each other's
intermediate state. The contract:
- Concurrent execution is supported for the read-only inputs accepted by
`transform_sql` and `query`, and for the registration APIs. `dry_run` is
concurrency-safe for statements that `EXPLAIN` only plans. An
`ANALYZE`-prefixed input becomes `EXPLAIN ANALYZE` and executes; like a
state-mutating statement accepted by `query()`, it is outside the
concurrency contract. Function lookup methods are read-only and
concurrency-safe.
- `register_parquet` / `register_csv` are safe under distinct table names;
registering the same name concurrently is unsupported.
- `list_tables` is a best-effort enumeration: registrations that land
mid-call may or may not appear, but the result is always well-formed.
- `load_mdl` must not overlap other calls on the same context; overlapping
calls raise `RuntimeError`.
## Developer Guide
### Environment Setup
core/wren-core-py/src/context.rs
+14 -40
View File
@@ -72,6 +72,14 @@ fn shared_runtime() -> PyResult<Arc<Runtime>> {
}
/// The Python wrapper for the Wren Core session context.
///
/// Calls on one context can run concurrently. Each `transform_sql` applies
/// the MDL onto a private top-level catalog snapshot (see
/// `clone_catalog_list` in wren-core's `mdl::context`), and analyzer rules
/// keep their mutable state per invocation. See the Concurrency section in
/// `README.md` for the supported operation contract. `load_mdl` takes
/// `&mut self`, so PyO3's exclusive borrow rejects overlapping calls on the
/// same context with a `RuntimeError`.
#[pyclass(name = "SessionContext")]
pub struct PySessionContext {
/// Base context — physical tables are registered here.
@@ -81,16 +89,6 @@ pub struct PySessionContext {
exec_ctx: wren_core::SessionContext,
mdl: Arc<AnalyzedWrenMDL>,
properties: Arc<HashMap<String, Option<String>>>,
/// Serializes engine calls on this context. `transform_sql` re-applies
/// the MDL catalog onto the context's shared catalog list on every call
/// (`apply_wren_on_ctx` -> `register_table_with_mdl`), while `query`,
/// `dry_run`, and `list_tables` read that same shared list, so two
/// concurrent calls on the same context can observe each other's
/// half-registered catalogs. The GIL used to provide this serialization
/// implicitly; now that the GIL is released around blocking work, keep
/// the same per-context ordering with a lock. Calls on *different*
/// contexts run in parallel.
call_lock: Arc<Mutex<()>>,
}
impl Hash for PySessionContext {
@@ -108,7 +106,6 @@ impl Default for PySessionContext {
exec_ctx: ctx,
mdl: Arc::new(AnalyzedWrenMDL::default()),
properties: Arc::new(HashMap::new()),
call_lock: Arc::new(Mutex::new(())),
}
}
}
@@ -148,7 +145,6 @@ impl PySessionContext {
exec_ctx: ctx,
mdl: Arc::new(AnalyzedWrenMDL::default()),
properties: Arc::new(HashMap::new()),
call_lock: Arc::new(Mutex::new(())),
});
};
@@ -245,7 +241,6 @@ impl PySessionContext {
exec_ctx,
mdl: analyzed_mdl,
properties: properties_ref,
call_lock: Arc::new(Mutex::new(())),
})
}
Err(e) => Err(CoreError::new(
@@ -268,7 +263,6 @@ impl PySessionContext {
let sql = sql.to_owned();
let rt = shared_runtime()?;
py.detach(|| {
let _guard = self.lock_calls();
rt.block_on(mdl::transform_sql_with_ctx(
&self.ctx,
Arc::clone(&self.mdl),
@@ -289,10 +283,7 @@ impl PySessionContext {
) -> PyResult<Vec<PyRemoteFunction>> {
let rt = shared_runtime()?;
let registered_functions: Vec<PyRemoteFunction> = py
.detach(|| {
let _guard = self.lock_calls();
rt.block_on(Self::get_registered_functions(&self.exec_ctx))
})
.detach(|| rt.block_on(Self::get_registered_functions(&self.exec_ctx)))
.map_err(CoreError::from)?
.into_iter()
.map(|f| f.into())
@@ -309,7 +300,6 @@ impl PySessionContext {
let rt = shared_runtime()?;
let functions = py
.detach(|| {
let _guard = self.lock_calls();
rt.block_on(Self::get_registered_function(
&function_name,
&self.exec_ctx,
@@ -392,7 +382,6 @@ impl PySessionContext {
let sql = sql.to_owned();
let rt = shared_runtime()?;
py.detach(|| {
let _guard = self.lock_calls();
let (batches, schema) = rt
.block_on(async {
let df = self.exec_ctx.sql(&sql).await?;
@@ -446,7 +435,6 @@ impl PySessionContext {
let (name, path) = (name.to_owned(), path.to_owned());
let rt = shared_runtime()?;
py.detach(|| {
let _guard = self.lock_calls();
rt.block_on(self.base_ctx.register_parquet(
&name,
&path,
@@ -467,7 +455,6 @@ impl PySessionContext {
let (name, path) = (name.to_owned(), path.to_owned());
let rt = shared_runtime()?;
py.detach(|| {
let _guard = self.lock_calls();
rt.block_on(self.base_ctx.register_csv(
&name,
&path,
@@ -479,12 +466,11 @@ impl PySessionContext {
}
/// List registered table names in the execution context.
///
/// This is a best-effort enumeration: registrations that land during
/// traversal may or may not appear, but the returned list is well-formed.
pub fn list_tables(&self, py: Python<'_>) -> PyResult<Vec<String>> {
// No block_on here, but the traversal reads the shared catalog
// list, so it must be ordered against catalog re-registration in
// transform/query/load_mdl via the same per-context lock.
py.detach(|| {
let _guard = self.lock_calls();
let catalog_names = self.exec_ctx.catalog_names();
let mut tables = Vec::new();
for catalog_name in &catalog_names {
@@ -500,14 +486,14 @@ impl PySessionContext {
})
}
/// Dry-run SQL (EXPLAIN) to validate without executing.
/// Explain SQL to validate its plan. An `ANALYZE`-prefixed input becomes
/// `EXPLAIN ANALYZE` and executes the statement.
#[pyo3(signature = (sql))]
pub fn dry_run(&self, py: Python<'_>, sql: &str) -> PyResult<String> {
let sql = sql.to_owned();
let rt = shared_runtime()?;
let result = py
.detach(|| {
let _guard = self.lock_calls();
rt.block_on(async {
let df = self.exec_ctx.sql(&format!("EXPLAIN {sql}")).await?;
df.collect().await
@@ -535,7 +521,6 @@ impl PySessionContext {
// resolve table references to real data during LocalRuntime execution.
let rt = shared_runtime()?;
let register_tables = py.detach(|| {
let _guard = self.lock_calls();
rt.block_on(async {
let mut tables = HashMap::new();
for catalog_name in self.base_ctx.catalog_names() {
@@ -566,7 +551,6 @@ impl PySessionContext {
);
let (unparser_ctx, exec_ctx) = py.detach(|| {
let _guard = self.lock_calls();
let unparser_ctx = rt
.block_on(apply_wren_on_ctx(
&self.base_ctx,
@@ -595,16 +579,6 @@ impl PySessionContext {
}
impl PySessionContext {
/// Take the per-context call lock. Must only be called from inside a
/// `py.detach(..)` closure (i.e. with the GIL released): blocking on
/// this lock while holding the GIL could deadlock against a thread
/// that holds the lock and is about to re-acquire the GIL.
fn lock_calls(&self) -> std::sync::MutexGuard<'_, ()> {
self.call_lock
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
fn register_remote_function(
ctx: &wren_core::SessionContext,
mut remote_function: RemoteFunction,
core/wren-core-py/tests/test_modeling_core.py
+4 -3
View File
@@ -618,9 +618,10 @@ def test_concurrent_calls_from_threads():
results = list(ex.map(worker, contexts))
assert all(r == expected for batch in results for r in batch)
# mode 2: one shared context, mixed methods. This is the mode that
# catches the same-context catalog race if the per-context call lock
# is ever removed while the GIL is released.
# mode 2: one shared context, mixed methods. Same-context concurrency
# relies on private catalog snapshots and per-invocation analyzer
# state; this hammer pins that contract (the register-heavy variant
# lives in test_same_context_concurrency.py).
shared = SessionContext(manifest_str, None)
expected_functions = len(shared.get_available_functions())
barrier = threading.Barrier(4)
core/wren-core-py/tests/test_same_context_concurrency.py
+282
View File
@@ -0,0 +1,282 @@
"""Same-context concurrency: mixed engine calls on one shared SessionContext.
One context, eight barrier-started threads first contend on one RLAC-protected
semantic query, exercising the execution context's long-lived
``ModelAnalyzeRule`` and its nested cycle-detection path. They then cycle
through a named operation schedule — ``transform_sql``, semantic and physical
``query``, ``dry_run``, singular function lookup, and ``list_tables`` — while
two threads also register Parquet/CSV files under distinct late-table names.
Every thread asserts content, so a race that corrupts catalog state or
analyzer state fails loudly rather than silently.
The hammer runs in a spawned child process because a native deadlock inside
a GIL-released ``block_on`` section cannot be interrupted from Python: the
parent enforces a wall-clock deadline and terminates the child on breach,
so a deadlock regression fails the suite instead of hanging it. Barrier
timeouts remain as a secondary guard for start-line stragglers.
Deliberately NOT tested: registering the same table name concurrently
(documented-unsupported) and creating brand-new top-level catalogs after
context derivation (outside the snapshot visibility contract).
"""
import base64
import faulthandler
import io
import json
import signal
import subprocess
import sys
import threading
import traceback
from pathlib import Path
import pyarrow as pa
import pyarrow.parquet as pq
from pyarrow import ipc
from wren_core import SessionContext
N_THREADS = 8
RLAC_ITERS = 20
ITERS = 30
BARRIER_TIMEOUT = 30
CHILD_DEADLINE = 300
CUSTOMER_KEYS = [1, 2, 3]
def _ipc_to_pydict(ipc_bytes):
return ipc.open_stream(io.BytesIO(bytes(ipc_bytes))).read_all().to_pydict()
def _customer_manifest_b64():
# Same shape as test_physical_tables: the three-part tableReference must
# match the default catalog/schema the register APIs write into.
manifest = {
"catalog": "my_catalog",
"schema": "my_schema",
"dataSource": "datafusion",
"models": [
{
"name": "customer",
"tableReference": {
"catalog": "datafusion",
"schema": "public",
"table": "customer",
},
"columns": [
{"name": "c_custkey", "type": "integer"},
{"name": "c_name", "type": "varchar"},
],
"rowLevelAccessControls": [
{
"name": "by_allowed",
"requiredProperties": [
{"name": "session_user", "required": True}
],
"condition": (
"c_custkey IN (SELECT allowed_id FROM allowed "
"WHERE allowed_user = @session_user)"
),
}
],
"primaryKey": "c_custkey",
},
{
"name": "allowed",
"tableReference": {
"catalog": "datafusion",
"schema": "public",
"table": "allowed",
},
"columns": [
{"name": "allowed_id", "type": "integer"},
{"name": "allowed_user", "type": "varchar"},
],
"primaryKey": "allowed_id",
},
],
}
return base64.b64encode(json.dumps(manifest).encode("utf-8")).decode("utf-8")
def _hammer(tmp_dir):
"""Child-process body. Raises on any correctness violation."""
tmp = Path(tmp_dir)
customer_path = tmp / "customer.parquet"
pq.write_table(
pa.table(
{
"c_custkey": pa.array(CUSTOMER_KEYS, type=pa.int32()),
"c_name": ["a", "b", "c"],
}
),
customer_path,
)
allowed_path = tmp / "allowed.parquet"
pq.write_table(
pa.table(
{
"allowed_id": pa.array([1, 3, 2], type=pa.int32()),
"allowed_user": ["alice", "alice", "bob"],
}
),
allowed_path,
)
# Shared files registered under many distinct late-table names; both
# formats carry the same rows for post-join content checks.
late_parquet = tmp / "late.parquet"
pq.write_table(pa.table({"v": pa.array([10, 20], type=pa.int64())}), late_parquet)
late_csv = tmp / "late.csv"
late_csv.write_text("v\n10\n20\n")
manifest = _customer_manifest_b64()
properties = frozenset({("session_user", "'alice'")})
ctx = SessionContext(manifest, None, properties)
ctx.register_parquet("customer", str(customer_path))
ctx.register_parquet("allowed", str(allowed_path))
# Physical-only sentinel, deliberately absent from the MDL: its name
# can only come from the physical catalog, so asserting it mid-race
# pins physical enumeration ("customer" would be satisfied by the MDL
# model alone — list_tables flattens names across catalogs).
ctx.register_parquet("physical_sentinel", str(late_parquet))
ctx.load_mdl(manifest)
semantic_sql = (
"SELECT c_custkey FROM my_catalog.my_schema.customer ORDER BY c_custkey"
)
physical_sql = "SELECT c_custkey FROM datafusion.public.customer ORDER BY c_custkey"
expected_semantic_rows = {"c_custkey": [1, 3]}
expected_physical_rows = {"c_custkey": CUSTOMER_KEYS}
expected_transform = ctx.transform_sql(semantic_sql)
expected_dry = ctx.dry_run(semantic_sql)
fn_name = ctx.get_available_functions()[0].name
assert _ipc_to_pydict(ctx.query(semantic_sql)) == expected_semantic_rows
assert _ipc_to_pydict(ctx.query(physical_sql)) == expected_physical_rows
def op_transform():
assert ctx.transform_sql(semantic_sql) == expected_transform
def op_semantic_query():
# The customer -> allowed RLAC subquery keeps the cycle stack live
# across nested analysis on exec_ctx's shared ModelAnalyzeRule.
assert _ipc_to_pydict(ctx.query(semantic_sql)) == expected_semantic_rows
def op_physical_query():
assert _ipc_to_pydict(ctx.query(physical_sql)) == expected_physical_rows
def op_dry_run():
assert ctx.dry_run(semantic_sql) == expected_dry
def op_function_lookup():
functions = ctx.get_available_function(fn_name)
assert functions
assert all(function.name == fn_name for function in functions)
def op_list_tables():
# Best-effort enumeration for late registrations, but a physical
# table that existed before the race must stay visible in every
# snapshot; late-table content is asserted post-join.
assert "physical_sentinel" in ctx.list_tables()
# Every thread cycles through the full schedule from a tid-dependent
# offset, so each operation runs concurrently with every other.
schedule = (
op_transform,
op_semantic_query,
op_physical_query,
op_dry_run,
op_function_lookup,
op_list_tables,
)
rlac_barrier = threading.Barrier(N_THREADS)
errors = []
def worker(tid):
try:
# Reuse one barrier per iteration so all threads contend on the
# same exec_ctx/model key while nested RLAC analysis is active.
for _ in range(RLAC_ITERS):
rlac_barrier.wait(timeout=BARRIER_TIMEOUT)
op_semantic_query()
for i in range(ITERS):
schedule[(tid + i) % len(schedule)]()
if tid < 2:
# Late registration under distinct names, alternating
# file formats; same-name concurrent registration is
# documented-unsupported and deliberately not exercised.
if i % 2 == 0:
ctx.register_parquet(f"late_{tid}_{i}", str(late_parquet))
else:
ctx.register_csv(f"late_{tid}_{i}", str(late_csv))
except BaseException:
rlac_barrier.abort()
errors.append(f"thread {tid}:\n{traceback.format_exc()}")
threads = [threading.Thread(target=worker, args=(tid,)) for tid in range(N_THREADS)]
for t in threads:
t.start()
for t in threads:
t.join()
if errors:
raise AssertionError("\n".join(errors))
# Every distinct-name registration must land. One table per format and
# registering thread also proves that each provider returns the fixture.
tables = ctx.list_tables()
for tid in range(2):
for i in range(ITERS):
name = f"late_{tid}_{i}"
assert name in tables
for i in (0, 1):
name = f"late_{tid}_{i}"
got = _ipc_to_pydict(
ctx.query(f'SELECT v FROM datafusion.public."{name}" ORDER BY v')
)
assert got == {"v": [10, 20]}
def _hammer_child(tmp_dir):
faulthandler.enable()
try:
_hammer(tmp_dir)
except BaseException:
traceback.print_exc()
sys.exit(1)
sys.exit(0)
def test_same_context_concurrent_engine_calls(tmp_path):
child = subprocess.Popen(
[sys.executable, __file__, "--hammer-child", str(tmp_path)]
)
try:
child.wait(timeout=CHILD_DEADLINE)
except subprocess.TimeoutExpired:
child.terminate()
try:
child.wait(timeout=10)
except subprocess.TimeoutExpired:
child.kill()
child.wait()
raise AssertionError(
f"hammer exceeded {CHILD_DEADLINE}s deadline — possible native deadlock"
)
if child.returncode < 0:
signum = -child.returncode
try:
signal_name = signal.Signals(signum).name
except ValueError:
signal_name = f"signal {signum}"
raise AssertionError(
f"hammer child terminated by {signal_name} (exit code {child.returncode})"
)
assert child.returncode == 0, (
f"hammer child failed with exit code {child.returncode}"
)
if __name__ == "__main__" and sys.argv[1:2] == ["--hammer-child"]:
_hammer_child(sys.argv[2])