Skip to content

Commit 3c14af7

Browse files
zahariashclaude
andcommitted
Merge origin/main into fix/nonblocking-set-query-timeout
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2 parents faf7159 + c852805 commit 3c14af7

16 files changed

Lines changed: 686 additions & 57 deletions

‎.github/workflows/ci.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ jobs:
113113
- name: Check formatting (black)
114114
working-directory: ladybug/tools/python_api
115115
run: |
116-
uv pip install black==26.3.0
116+
uv pip install black==26.10.0
117117
.venv/bin/black --check src_py test
118118
119119
- name: Run ruff check
@@ -184,7 +184,7 @@ jobs:
184184
- name: Check formatting (black)
185185
working-directory: ladybug/tools/python_api
186186
run: |
187-
uv pip install black==26.3.0
187+
uv pip install black==26.10.0
188188
.venv/bin/black --check src_py test
189189
190190
- name: Run ruff check

‎pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@ dev = [
2424
"pytest",
2525
"pytest-asyncio~=1.0",
2626
"pytest-xdist~=3.0",
27-
"setuptools~=80.9",
28-
"black==26.3.0",
27+
"setuptools>=80.9",
28+
"black>=26.3.1",
2929
"ruff==0.15.19",
3030
"mypy==1.16.0",
3131
"torch>=2.5.0",

‎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
}

‎src_py/async_connection.py‎

Lines changed: 84 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import asyncio
44
import threading
55
import warnings
6+
import weakref
67
from concurrent.futures import ThreadPoolExecutor
78
from typing import TYPE_CHECKING, Any
89

@@ -51,6 +52,14 @@ def __init__(
5152
self.connections = [Connection(database) for _ in range(max_concurrent_queries)]
5253
self.connections_counter = [0 for _ in range(max_concurrent_queries)]
5354
self.lock = threading.Lock()
55+
# Calls on one connection run one at a time (the engine serializes them
56+
# anyway), so the call holding a connection's lock is the one an
57+
# interrupt would hit. Queries run directly on a connection obtained
58+
# from acquire_connection() are not tracked.
59+
self._connection_locks: weakref.WeakKeyDictionary[
60+
Connection, threading.Lock
61+
] = weakref.WeakKeyDictionary()
62+
self._running: dict[Connection, _Call] = {}
5463

5564
for conn in self.connections:
5665
conn.init_connection()
@@ -114,10 +123,9 @@ def release_connection(self, conn: Connection) -> None:
114123
115124
116125
"""
117-
for i, existing_conn in enumerate(self.connections):
118-
if existing_conn == conn:
119-
self.__decrement_connection_counter(i)
120-
break
126+
conn_index = self.__index_of(conn)
127+
if conn_index is not None:
128+
self.__decrement_connection_counter(conn_index)
121129

122130
def set_query_timeout(self, timeout_in_ms: int) -> None:
123131
"""
@@ -153,28 +161,22 @@ async def execute(
153161
QueryResult
154162
Query result.
155163
164+
Raises
165+
------
166+
asyncio.CancelledError
167+
If the awaiting task is cancelled. A running query is interrupted; a
168+
query still waiting for a free connection is not started.
169+
156170
"""
157-
loop = asyncio.get_running_loop()
158-
# If the query is a prepared statement, use the connection associated with it
159171
if isinstance(query, PreparedStatement):
160172
conn = query._connection
161-
for i, existing_conn in enumerate(self.connections):
162-
if existing_conn == conn:
163-
conn_index = i
164-
with self.lock:
165-
self.connections_counter[conn_index] += 1
166-
break
173+
conn_index = self.__index_of(conn)
174+
if conn_index is not None:
175+
with self.lock:
176+
self.connections_counter[conn_index] += 1
167177
else:
168178
conn, conn_index = self.__get_connection_with_least_queries()
169-
170-
try:
171-
return await loop.run_in_executor(
172-
self.executor, conn.execute, query, parameters
173-
)
174-
except asyncio.CancelledError:
175-
conn.interrupt()
176-
finally:
177-
self.__decrement_connection_counter(conn_index)
179+
return await self.__submit(conn, conn_index, conn.execute, query, parameters)
178180

179181
async def _prepare(
180182
self, query: str, parameters: dict[str, Any] | None = None
@@ -183,16 +185,64 @@ async def _prepare(
183185
The only parameters supported during prepare are dataframes.
184186
Any remaining parameters will be ignored and should be passed to execute().
185187
""" # noqa: D401
186-
loop = asyncio.get_running_loop()
187188
conn, conn_index = self.__get_connection_with_least_queries()
189+
return await self.__submit(conn, conn_index, conn._prepare, query, parameters)
188190

191+
def __index_of(self, conn: Connection) -> int | None:
192+
for i, existing_conn in enumerate(self.connections):
193+
if existing_conn is conn:
194+
return i
195+
return None
196+
197+
async def __submit(
198+
self, conn: Connection, conn_index: int | None, fn: Any, *args: Any
199+
) -> Any:
200+
call = _Call()
189201
try:
190-
preparedStatement = await loop.run_in_executor(
191-
self.executor, conn.prepare, query, parameters
202+
future = self.executor.submit(self.__run, conn, call, fn, *args)
203+
except RuntimeError:
204+
if conn_index is not None:
205+
self.__decrement_connection_counter(conn_index)
206+
raise
207+
if conn_index is not None:
208+
# Release the slot when the work has actually finished (or was
209+
# cancelled before it started), not when the awaiting task exits.
210+
future.add_done_callback(
211+
lambda _: self.__decrement_connection_counter(conn_index)
192212
)
193-
return preparedStatement
194-
finally:
195-
self.__decrement_connection_counter(conn_index)
213+
try:
214+
return await asyncio.wrap_future(future)
215+
except asyncio.CancelledError:
216+
self.__cancel(conn, call)
217+
raise
218+
219+
def __run(self, conn: Connection, call: _Call, fn: Any, *args: Any) -> Any:
220+
with self.lock:
221+
connection_lock = self._connection_locks.setdefault(conn, threading.Lock())
222+
with connection_lock:
223+
with self.lock:
224+
if call.cancelled:
225+
return None
226+
self._running[conn] = call
227+
try:
228+
return fn(*args)
229+
finally:
230+
with self.lock:
231+
del self._running[conn]
232+
233+
def __cancel(self, conn: Connection, call: _Call) -> None:
234+
# Interrupting acts on whatever the connection is running, so only
235+
# interrupt while this call is the one running. A call still waiting
236+
# for its connection is skipped when its turn comes.
237+
with self.lock:
238+
call.cancelled = True
239+
running = self._running.get(conn) is call
240+
if running:
241+
conn.interrupt()
242+
if running:
243+
# An interrupt that arrives while the query is still compiling is
244+
# cleared when execution starts, so repeat it until the call returns.
245+
asyncio.get_running_loop().call_later(0.05, self.__cancel, conn, call)
196246

197247
async def prepare(
198248
self, query: str, parameters: dict[str, Any] | None = None
@@ -225,10 +275,17 @@ def close(self) -> None:
225275
"""
226276
Close all connections and shutdown the thread pool.
227277
278+
Queries still running are interrupted and fail with "Interrupted.".
279+
228280
Note: Call to this method is optional. The connections and thread pool
229281
will be closed automatically when the instance is garbage collected.
230282
"""
231283
for conn in self.connections:
232284
conn.close()
233285

234286
self.executor.shutdown(wait=True)
287+
288+
289+
class _Call:
290+
def __init__(self) -> None:
291+
self.cancelled = False

0 commit comments

Comments
 (0)