Skip to content

Commit 2acc78b

Browse files
committed
Close pending statements on connection close
This is a version 3 of this PR that attempts to safely close pending statements when the connection is closed. In it all tracking and locking logic is moved from C++ to Java: - `Connection`, `Statement` and `ResultSet` instances use their own locks; the check that corresponding reference is still alive is performed before every native call after obtaining the lock. - `Statement` lock is held during the query execution, note, this lock is NOT requied to call `statement#cancel()` because this operation is implemented on a `Connection` level. - When a `Connection` is being closed, pending query is cancelled first, and then all active statements are closed in a reverse creation order. Note, thread safety for Appender and Arrow interfaces is going to be addressed in subsequent PRs. Testing: new tests added for various sequential and concurrent closure scenarios. Fixes: #101
1 parent e4f5592 commit 2acc78b

14 files changed

+821
-406
lines changed

CMakeLists.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -439,7 +439,7 @@ set(DUCKDB_SRC_FILES
439439
src/duckdb/extension/json/json_serializer.cpp
440440
src/duckdb/ub_extension_json_json_functions.cpp)
441441

442-
set(JEMACLLOC_SRC_FILES
442+
set(JEMALLOC_SRC_FILES
443443
src/duckdb/extension/jemalloc/jemalloc_extension.cpp
444444
src/duckdb/extension/jemalloc/jemalloc/src/jemalloc.c
445445
src/duckdb/extension/jemalloc/jemalloc/src/arena.c
@@ -553,7 +553,7 @@ add_jar(duckdb_jdbc_tests ${JAVA_TEST_FILES} INCLUDE_JARS duckdb_jdbc)
553553
if(MSVC)
554554
list(APPEND DUCKDB_SRC_FILES duckdb_java.def)
555555
else()
556-
list(APPEND DUCKDB_SRC_FILES ${JEMACLLOC_SRC_FILES})
556+
list(APPEND DUCKDB_SRC_FILES ${JEMALLOC_SRC_FILES})
557557
endif()
558558

559559
add_library(duckdb_java SHARED

CMakeLists.txt.in

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ set(DUCKDB_DEFINITIONS
4646
set(DUCKDB_SRC_FILES
4747
${SOURCES})
4848

49-
set(JEMACLLOC_SRC_FILES
49+
set(JEMALLOC_SRC_FILES
5050
${JEMALLOC_SOURCES})
5151

5252

@@ -95,7 +95,7 @@ add_jar(duckdb_jdbc_tests ${JAVA_TEST_FILES} INCLUDE_JARS duckdb_jdbc)
9595
if(MSVC)
9696
list(APPEND DUCKDB_SRC_FILES duckdb_java.def)
9797
else()
98-
list(APPEND DUCKDB_SRC_FILES ${JEMACLLOC_SRC_FILES})
98+
list(APPEND DUCKDB_SRC_FILES ${JEMALLOC_SRC_FILES})
9999
endif()
100100

101101
add_library(duckdb_java SHARED

src/jni/duckdb_java.cpp

Lines changed: 2 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
#include "duckdb/main/extension_util.hpp"
1414
#include "duckdb/parser/parsed_data/create_type_info.hpp"
1515
#include "functions.hpp"
16+
#include "holders.hpp"
1617
#include "refs.hpp"
1718
#include "types.hpp"
1819
#include "util.hpp"
@@ -59,40 +60,6 @@ JNIEXPORT void JNICALL JNI_OnUnload(JavaVM *vm, void *reserved) {
5960
delete_global_refs(env);
6061
}
6162

62-
/**
63-
* Associates a duckdb::Connection with a duckdb::DuckDB. The DB may be shared amongst many ConnectionHolders, but the
64-
* Connection is unique to this holder. Every Java DuckDBConnection has exactly 1 of these holders, and they are never
65-
* shared. The holder is freed when the DuckDBConnection is closed. When the last holder sharing a DuckDB is freed, the
66-
* DuckDB is released as well.
67-
*/
68-
struct ConnectionHolder {
69-
const duckdb::shared_ptr<duckdb::DuckDB> db;
70-
const duckdb::unique_ptr<duckdb::Connection> connection;
71-
72-
ConnectionHolder(duckdb::shared_ptr<duckdb::DuckDB> _db)
73-
: db(_db), connection(make_uniq<duckdb::Connection>(*_db)) {
74-
}
75-
};
76-
77-
/**
78-
* Throws a SQLException and returns nullptr if a valid Connection can't be retrieved from the buffer.
79-
*/
80-
static Connection *get_connection(JNIEnv *env, jobject conn_ref_buf) {
81-
if (!conn_ref_buf) {
82-
throw ConnectionException("Invalid connection");
83-
}
84-
auto conn_holder = (ConnectionHolder *)env->GetDirectBufferAddress(conn_ref_buf);
85-
if (!conn_holder) {
86-
throw ConnectionException("Invalid connection");
87-
}
88-
auto conn_ref = conn_holder->connection.get();
89-
if (!conn_ref || !conn_ref->context) {
90-
throw ConnectionException("Invalid connection");
91-
}
92-
93-
return conn_ref;
94-
}
95-
9663
//! The database instance cache, used so that multiple connections to the same file point to the same database object
9764
duckdb::DBInstanceCache instance_cache;
9865

@@ -189,10 +156,6 @@ void _duckdb_jdbc_disconnect(JNIEnv *env, jclass, jobject conn_ref_buf) {
189156
}
190157
}
191158

192-
struct StatementHolder {
193-
duckdb::unique_ptr<PreparedStatement> stmt;
194-
};
195-
196159
#include "utf8proc_wrapper.hpp"
197160

198161
jobject _duckdb_jdbc_prepare(JNIEnv *env, jclass, jobject conn_ref_buf, jbyteArray query_j) {
@@ -233,11 +196,6 @@ jobject _duckdb_jdbc_prepare(JNIEnv *env, jclass, jobject conn_ref_buf, jbyteArr
233196
return env->NewDirectByteBuffer(stmt_ref, 0);
234197
}
235198

236-
struct ResultHolder {
237-
duckdb::unique_ptr<QueryResult> res;
238-
duckdb::unique_ptr<DataChunk> chunk;
239-
};
240-
241199
Value ToValue(JNIEnv *env, jobject param, duckdb::shared_ptr<ClientContext> context) {
242200
param = env->CallStaticObjectMethod(J_Timestamp, J_Timestamp_valueOf, param);
243201

@@ -930,6 +888,7 @@ static ProfilerPrintFormat GetProfilerPrintFormat(JNIEnv *env, jobject format) {
930888
if (env->IsSameObject(format, J_ProfilerPrintFormat_GRAPHVIZ)) {
931889
return ProfilerPrintFormat::GRAPHVIZ;
932890
}
891+
throw InvalidInputException("Invalid profiling format");
933892
}
934893

935894
jstring _duckdb_jdbc_get_profiling_information(JNIEnv *env, jclass, jobject conn_ref_buf, jobject j_format) {

src/jni/functions.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,5 +404,6 @@ JNIEXPORT jstring JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1get_1profil
404404
duckdb::ErrorData error(e);
405405
ThrowJNI(env, error.Message().c_str());
406406

407+
return nullptr;
407408
}
408409
}

src/jni/holders.hpp

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
#pragma once
2+
3+
#include "duckdb.hpp"
4+
5+
#include <jni.h>
6+
7+
/**
8+
* Associates a duckdb::Connection with a duckdb::DuckDB. The DB may be shared amongst many ConnectionHolders, but the
9+
* Connection is unique to this holder. Every Java DuckDBConnection has exactly 1 of these holders, and they are never
10+
* shared. The holder is freed when the DuckDBConnection is closed. When the last holder sharing a DuckDB is freed, the
11+
* DuckDB is released as well.
12+
*/
13+
struct ConnectionHolder {
14+
const duckdb::shared_ptr<duckdb::DuckDB> db;
15+
const duckdb::unique_ptr<duckdb::Connection> connection;
16+
17+
ConnectionHolder(duckdb::shared_ptr<duckdb::DuckDB> _db)
18+
: db(_db), connection(duckdb::make_uniq<duckdb::Connection>(*_db)) {
19+
}
20+
};
21+
22+
struct StatementHolder {
23+
duckdb::unique_ptr<duckdb::PreparedStatement> stmt;
24+
};
25+
26+
struct ResultHolder {
27+
duckdb::unique_ptr<duckdb::QueryResult> res;
28+
duckdb::unique_ptr<duckdb::DataChunk> chunk;
29+
};
30+
31+
/**
32+
* Throws a SQLException and returns nullptr if a valid Connection can't be retrieved from the buffer.
33+
*/
34+
inline duckdb::Connection *get_connection(JNIEnv *env, jobject conn_ref_buf) {
35+
if (!conn_ref_buf) {
36+
throw duckdb::ConnectionException("Invalid connection");
37+
}
38+
auto conn_holder = (ConnectionHolder *)env->GetDirectBufferAddress(conn_ref_buf);
39+
if (!conn_holder) {
40+
throw duckdb::ConnectionException("Invalid connection");
41+
}
42+
auto conn_ref = conn_holder->connection.get();
43+
if (!conn_ref || !conn_ref->context) {
44+
throw duckdb::ConnectionException("Invalid connection");
45+
}
46+
47+
return conn_ref;
48+
}

src/main/java/org/duckdb/DuckDBAppender.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ public DuckDBAppender(DuckDBConnection con, String schemaName, String tableName)
1616
throw new SQLException("Invalid connection");
1717
}
1818
appender_ref = DuckDBNative.duckdb_jdbc_create_appender(
19-
con.conn_ref, schemaName.getBytes(StandardCharsets.UTF_8), tableName.getBytes(StandardCharsets.UTF_8));
19+
con.connRef, schemaName.getBytes(StandardCharsets.UTF_8), tableName.getBytes(StandardCharsets.UTF_8));
2020
}
2121

2222
public void beginRow() throws SQLException {

0 commit comments

Comments
 (0)