Skip to content

Commit 9f64325

Browse files
zahariashclaude
authored andcommitted
fix(pybind): load each import-cache entry once, under a lock
The driver loads the Python objects it needs (datetime, uuid, pandas, ...) into a process-wide import cache on first use. PythonCachedItem::operator() checked and set its loaded flag without synchronization, and addToCache appended to a shared vector, relying on the GIL. On free-threaded CPython with the GIL disabled, threads making their first calls at the same time loaded entries concurrently: the vector reallocated under another thread's push_back and corrupted the heap, which crashed later in an unrelated allocation (SIGSEGV, or "free(): invalid pointer"). Each entry now loads once. The flag is atomic, and loading and addToCache run under a mutex that waiters take with the GIL released (detached on free-threaded builds), so the thread holding it can always finish its import; the parent entry is loaded before locking. Creating the cache in the Database constructor is serialized as well. PandasAnalyzer took a copy of the cached pandas entry, so it loaded into the copy and added to the cache on every call; it now uses the cache's own entry. This does not declare free-threading support; it removes this race. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 548c183 commit 9f64325

7 files changed

Lines changed: 89 additions & 11 deletions

File tree

‎src_cpp/cached_import/py_cached_import.cpp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,13 @@ PythonCachedImport::~PythonCachedImport() {
1111
allObjects.clear();
1212
}
1313

14+
std::unique_lock<std::mutex> PythonCachedImport::lockForLoading() {
15+
// Wait with the GIL released (detached on free-threaded builds): the holder needs the GIL to
16+
// finish its import, and a detached waiter does not stall stop-the-world pauses.
17+
py::gil_scoped_release release;
18+
return std::unique_lock<std::mutex>(loadMutex);
19+
}
20+
1421
py::handle PythonCachedImport::addToCache(py::object obj) {
1522
auto ptr = obj.ptr();
1623
allObjects.push_back(obj);

‎src_cpp/cached_import/py_cached_item.cpp‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,18 @@ namespace lbug {
77

88
py::handle PythonCachedItem::operator()() {
99
assert((bool)PyGILState_Check());
10-
// load if unloaded, return cached object if already loaded
11-
if (loaded) {
10+
if (isLoaded()) {
1211
return object;
1312
}
14-
if (parent == nullptr) {
15-
object = importCache->addToCache(std::move(py::module::import(name.c_str())));
16-
} else {
17-
object = importCache->addToCache(std::move((*parent)().attr(name.c_str())));
13+
// The parent takes the same lock, so load it before locking.
14+
py::handle parentObject = parent == nullptr ? py::handle() : (*parent)();
15+
auto lock = importCache->lockForLoading();
16+
if (!isLoaded()) {
17+
auto obj = parent == nullptr ? py::object(py::module::import(name.c_str())) :
18+
py::object(parentObject.attr(name.c_str()));
19+
object = importCache->addToCache(std::move(obj));
20+
loaded.store(true, std::memory_order_release);
1821
}
19-
loaded = true;
2022
return object;
2123
}
2224

‎src_cpp/include/cached_import/py_cached_import.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#pragma once
22

3+
#include <mutex>
34
#include <string>
45
#include <vector>
56

@@ -15,7 +16,9 @@ class PythonCachedImport {
1516
PythonCachedImport() = default;
1617
~PythonCachedImport();
1718

19+
// Call with the lock from lockForLoading() held.
1820
py::handle addToCache(py::object obj);
21+
std::unique_lock<std::mutex> lockForLoading();
1922

2023
DateTimeCachedItem datetime;
2124
DecimalCachedItem decimal;
@@ -28,6 +31,7 @@ class PythonCachedImport {
2831
UUIDCachedItem uuid;
2932

3033
private:
34+
std::mutex loadMutex;
3135
std::vector<py::object> allObjects;
3236
};
3337

‎src_cpp/include/cached_import/py_cached_item.h‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#pragma once
22

3+
#include <atomic>
34
#include <memory>
45
#include <string>
56

@@ -10,16 +11,16 @@ namespace lbug {
1011
class PythonCachedItem {
1112
public:
1213
explicit PythonCachedItem(std::string name, PythonCachedItem* parent = nullptr)
13-
: name(std::move(name)), parent(parent), loaded(false) {}
14+
: name(std::move(name)), parent(parent) {}
1415
virtual ~PythonCachedItem() = default;
1516

16-
bool isLoaded() const { return loaded; }
17+
bool isLoaded() const { return loaded.load(std::memory_order_acquire); }
1718
py::handle operator()();
1819

1920
protected:
2021
std::string name;
2122
PythonCachedItem* parent;
22-
bool loaded;
23+
std::atomic<bool> loaded{false};
2324
py::handle object;
2425
};
2526

‎src_cpp/pandas/pandas_analyzer.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ static py::object findFirstNonNull(const py::handle& row, uint64_t numRows) {
193193

194194
common::LogicalType PandasAnalyzer::innerAnalyze(py::object column, bool& canConvert) {
195195
auto numRows = py::len(column);
196-
auto pandasModule = importCache->pandas;
196+
auto& pandasModule = importCache->pandas;
197197
auto pandasSeries = pandasModule.core.series.Series();
198198

199199
if (py::isinstance(column, pandasSeries)) {

‎src_cpp/py_database.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include "include/py_database.h"
22

33
#include <memory>
4+
#include <mutex>
45

56
#include "common/exception/runtime.h"
67
#include "extension/extension.h"
@@ -65,6 +66,9 @@ PyDatabase::PyDatabase(const std::string& databasePath, uint64_t bufferPoolSize,
6566
lbug::extension::ExtensionUtils::addTableFunc<lbug::PandasScanFunction>(*state->database);
6667
state->storageDriver = std::make_unique<StorageDriver>(state->database.get());
6768
py::gil_scoped_acquire acquire;
69+
// Taken with the GIL held: nothing inside may call into Python.
70+
static std::mutex importCacheMutex;
71+
std::lock_guard<std::mutex> lock(importCacheMutex);
6872
if (lbug::importCache.get() == nullptr) {
6973
lbug::importCache = std::make_shared<lbug::PythonCachedImport>();
7074
}
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
from __future__ import annotations
2+
3+
import os
4+
import subprocess
5+
import sys
6+
import sysconfig
7+
from textwrap import dedent
8+
from typing import TYPE_CHECKING
9+
10+
import pytest
11+
12+
if TYPE_CHECKING:
13+
from pathlib import Path
14+
15+
# The driver fills its import cache on first use, and the cache cannot be reset within a
16+
# process, so every attempt runs in a fresh interpreter. Only a free-threaded build with the GIL
17+
# disabled can expose a race; on other builds this is a smoke test.
18+
FREE_THREADED = bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
19+
20+
21+
@pytest.mark.parametrize("attempt", range(10 if FREE_THREADED else 1))
22+
def test_concurrent_first_use_of_the_import_cache(
23+
attempt: int, build_dir: Path
24+
) -> None:
25+
code = dedent(f"""
26+
import sys
27+
import threading
28+
29+
sys.path.append(r"{build_dir!s}")
30+
31+
import ladybug as lb
32+
33+
db = lb.Database(":memory:")
34+
conns = [lb.Connection(db) for _ in range(16)]
35+
start = threading.Barrier(len(conns))
36+
errors = []
37+
38+
def run(conn, n):
39+
start.wait()
40+
try:
41+
assert conn.execute("RETURN $n", {{"n": n}}).get_next() == [n]
42+
except BaseException as e:
43+
errors.append(e)
44+
45+
threads = [threading.Thread(target=run, args=(c, i)) for i, c in enumerate(conns)]
46+
for t in threads:
47+
t.start()
48+
for t in threads:
49+
t.join()
50+
assert not errors, errors
51+
""")
52+
env = {**os.environ, "PYTHON_GIL": "0"} if FREE_THREADED else None
53+
result = subprocess.run(
54+
[sys.executable, "-c", code],
55+
env=env,
56+
capture_output=True,
57+
text=True,
58+
check=False,
59+
)
60+
assert result.returncode == 0, (result.returncode, result.stderr)

0 commit comments

Comments
 (0)