From f19a1fbfa27800ca5fdff097282594eb42d9af6c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 6 Aug 2026 08:50:17 -0600 Subject: [PATCH] fix: serialise native handle lifetime against concurrent close SessionContext and DataFrame each held a raw `long nativeHandle` guarded by `if (nativeHandle == 0) throw`. The check and the JNI call were separate operations, so a concurrent close() could free the native Box between them and leave the other thread dereferencing freed memory. Introduce a package-private NativeHandle that owns the pointer. Every call pins the handle for its duration (acquire/release); close() and the consuming operations claim it, waiting for in-flight calls to drain before handing the pointer back to Rust to be freed. acquire() never blocks, which is what lets the two-handle set operations and joins pin both DataFrames without a lock-ordering hazard. A read/write lock would not: its readers queue behind a waiting writer, so opposing pin orders with closes interleaved could deadlock. claim() is the only blocking operation and always makes progress, since a thread holding a pin is inside a native call and never claims. No native change is needed. Every non-consuming JNI entry point already takes a shared reference, and DataFusion's SessionContext and DataFrame are Send + Sync, so only the handle's lifetime had to be serialised. Closes #40 --- README.md | 5 +- .../java/org/apache/datafusion/DataFrame.java | 449 +++++++++------- .../org/apache/datafusion/NativeHandle.java | 172 +++++++ .../org/apache/datafusion/SessionContext.java | 480 ++++++++++-------- .../datafusion/DataFrameConcurrencyTest.java | 258 ++++++++++ .../apache/datafusion/NativeHandleTest.java | 237 +++++++++ .../SessionContextConcurrencyTest.java | 196 +++++++ docs/source/user-guide/quickstart.md | 5 +- docs/source/user-guide/sessioncontext.md | 20 +- 9 files changed, 1417 insertions(+), 405 deletions(-) create mode 100644 core/src/main/java/org/apache/datafusion/NativeHandle.java create mode 100644 core/src/test/java/org/apache/datafusion/DataFrameConcurrencyTest.java create mode 100644 core/src/test/java/org/apache/datafusion/NativeHandleTest.java create mode 100644 core/src/test/java/org/apache/datafusion/SessionContextConcurrencyTest.java diff --git a/README.md b/README.md index c14eb8c..dec371e 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,10 @@ try (var allocator = new RootAllocator(); } ``` -`SessionContext` and `DataFrame` are `AutoCloseable` and not thread-safe. +`SessionContext` and `DataFrame` are `AutoCloseable` and safe to share +across threads: `close()` waits for calls already in flight to return +before releasing the native object, and a call that loses the race to a +`close()` throws `IllegalStateException`. ## Documentation diff --git a/core/src/main/java/org/apache/datafusion/DataFrame.java b/core/src/main/java/org/apache/datafusion/DataFrame.java index d4e0226..330cbd7 100644 --- a/core/src/main/java/org/apache/datafusion/DataFrame.java +++ b/core/src/main/java/org/apache/datafusion/DataFrame.java @@ -22,6 +22,7 @@ import java.io.ByteArrayInputStream; import java.io.IOException; import java.nio.channels.Channels; +import java.util.function.LongBinaryOperator; import org.apache.arrow.c.ArrowArrayStream; import org.apache.arrow.c.Data; @@ -37,23 +38,32 @@ * {@link #collect} (materializes every batch on the native heap before returning) or {@link * #executeStream} (yields one batch at a time as Java drains the reader). * - *

Instances are not thread-safe and must be closed. Both {@link #collect} and - * {@link #executeStream} consume the DataFrame: a successfully consumed DataFrame cannot be - * consumed again by either method (or by other executors such as {@link #count}), and {@link - * #close()} on an already-consumed instance is a no-op. + *

Instances must be closed. Both {@link #collect} and {@link #executeStream} consume the + * DataFrame: a successfully consumed DataFrame cannot be consumed again by either method (or by + * other executors such as {@link #count}), and {@link #close()} on an already-consumed instance is + * a no-op. + * + *

Instances are safe to share between threads. This class pins its native handle for the + * duration of every call, so a {@link #close} or a consuming operation racing with work on another + * thread cannot free the plan out from under it; the call that loses such a race throws {@link + * IllegalStateException}. When several threads race to consume the same DataFrame, exactly one + * succeeds. {@link #close} and the consuming operations block until calls already in flight have + * returned. */ public final class DataFrame implements AutoCloseable { static { NativeLibraryLoader.loadLibrary(); } - private long nativeHandle; + private static final String CLOSED = "DataFrame is closed or already collected"; + + private final NativeHandle handle; DataFrame(long nativeHandle) { if (nativeHandle == 0) { throw new IllegalArgumentException("DataFrame native handle is null"); } - this.nativeHandle = nativeHandle; + this.handle = new NativeHandle(nativeHandle, CLOSED); } /** @@ -68,14 +78,9 @@ public final class DataFrame implements AutoCloseable { * {@link #executeStream(BufferAllocator)} for analytics-scale queries. */ public ArrowReader collect(BufferAllocator allocator) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator); - long handle = nativeHandle; - nativeHandle = 0; try { - collectDataFrame(handle, stream.memoryAddress()); + collectDataFrame(handle.claim(), stream.memoryAddress()); return Data.importArrayStream(allocator, stream); } catch (Throwable e) { stream.close(); @@ -98,14 +103,9 @@ public ArrowReader collect(BufferAllocator allocator) { * use this method. */ public ArrowReader executeStream(BufferAllocator allocator) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator); - long handle = nativeHandle; - nativeHandle = 0; try { - executeStreamDataFrame(handle, stream.memoryAddress()); + executeStreamDataFrame(handle.claim(), stream.memoryAddress()); return Data.importArrayStream(allocator, stream); } catch (Throwable e) { stream.close(); @@ -121,10 +121,13 @@ public ArrowReader executeStream(BufferAllocator allocator) { * schema carries no buffer data. */ public Schema schema() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + byte[] ipcBytes; + long h = handle.acquire(); + try { + ipcBytes = schemaIpc(h); + } finally { + handle.release(); } - byte[] ipcBytes = schemaIpc(nativeHandle); try { return MessageSerializer.deserializeSchema( new ReadChannel(Channels.newChannel(new ByteArrayInputStream(ipcBytes)))); @@ -143,10 +146,12 @@ public Schema schema() { * #show()} or {@link #collect(BufferAllocator)}. */ public DataFrame explain(boolean verbose, boolean analyze) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(explainPlan(h, verbose, analyze)); + } finally { + handle.release(); } - return new DataFrame(explainPlan(nativeHandle, verbose, analyze)); } /** @@ -160,10 +165,12 @@ public DataFrame explain(boolean verbose, boolean analyze) { * @throws RuntimeException if execution fails. */ public DataFrame cache() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(cachePlan(h)); + } finally { + handle.release(); } - return new DataFrame(cachePlan(nativeHandle)); } /** @@ -178,34 +185,42 @@ public DataFrame cache() { * @throws RuntimeException if execution fails. */ public DataFrame describe() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(describePlan(h)); + } finally { + handle.release(); } - return new DataFrame(describePlan(nativeHandle)); } /** Execute the plan and return the number of rows. */ public long count() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return countRows(h); + } finally { + handle.release(); } - return countRows(nativeHandle); } /** Execute the plan and print formatted batches to native stdout. */ public void show() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + showDataFrame(h); + } finally { + handle.release(); } - showDataFrame(nativeHandle); } /** Execute the plan and print the first {@code limit} rows to native stdout. */ public void show(int limit) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + showDataFrameWithLimit(h, limit); + } finally { + handle.release(); } - showDataFrameWithLimit(nativeHandle, limit); } /** @@ -213,10 +228,12 @@ public void show(int limit) { * closed independently. */ public DataFrame select(String... columnNames) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(selectColumns(h, columnNames)); + } finally { + handle.release(); } - return new DataFrame(selectColumns(nativeHandle, columnNames)); } /** @@ -224,10 +241,12 @@ public DataFrame select(String... columnNames) { * DataFrame's own schema. The receiver remains usable and must still be closed independently. */ public DataFrame filter(String predicate) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(filterRows(h, predicate)); + } finally { + handle.release(); } - return new DataFrame(filterRows(nativeHandle, predicate)); } /** @@ -249,10 +268,12 @@ public DataFrame limit(int skip, int fetch) { if (fetch < 0) { throw new IllegalArgumentException("fetch must be non-negative, was " + fetch); } - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(limitRows(h, skip, fetch)); + } finally { + handle.release(); } - return new DataFrame(limitRows(nativeHandle, skip, fetch)); } /** @@ -260,10 +281,12 @@ public DataFrame limit(int skip, int fetch) { * independently. */ public DataFrame distinct() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(distinctRows(h)); + } finally { + handle.release(); } - return new DataFrame(distinctRows(nativeHandle)); } /** @@ -271,18 +294,22 @@ public DataFrame distinct() { * and must still be closed independently. */ public DataFrame dropColumns(String... columnNames) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(dropColumns(h, columnNames)); + } finally { + handle.release(); } - return new DataFrame(dropColumns(nativeHandle, columnNames)); } /** Rename a column. The receiver remains usable and must still be closed independently. */ public DataFrame withColumnRenamed(String oldName, String newName) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(renameColumn(h, oldName, newName)); + } finally { + handle.release(); } - return new DataFrame(renameColumn(nativeHandle, oldName, newName)); } /** @@ -294,16 +321,18 @@ public DataFrame withColumnRenamed(String oldName, String newName) { * @throws IllegalArgumentException if {@code name} or {@code expr} is {@code null}. */ public DataFrame withColumn(String name, String expr) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (name == null) { - throw new IllegalArgumentException("withColumn name must be non-null"); - } - if (expr == null) { - throw new IllegalArgumentException("withColumn expr must be non-null"); + long h = handle.acquire(); + try { + if (name == null) { + throw new IllegalArgumentException("withColumn name must be non-null"); + } + if (expr == null) { + throw new IllegalArgumentException("withColumn expr must be non-null"); + } + return new DataFrame(withColumnExpr(h, name, expr)); + } finally { + handle.release(); } - return new DataFrame(withColumnExpr(nativeHandle, name, expr)); } /** @@ -322,16 +351,18 @@ public DataFrame unnestColumns(String... columns) { * @throws IllegalArgumentException if {@code options} or {@code columns} is {@code null}. */ public DataFrame unnestColumns(UnnestOptions options, String... columns) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (options == null) { - throw new IllegalArgumentException("unnestColumns options must be non-null"); - } - if (columns == null) { - throw new IllegalArgumentException("unnestColumns columns must be non-null"); + long h = handle.acquire(); + try { + if (options == null) { + throw new IllegalArgumentException("unnestColumns options must be non-null"); + } + if (columns == null) { + throw new IllegalArgumentException("unnestColumns columns must be non-null"); + } + return new DataFrame(unnestColumns(h, columns, options.preserveNulls())); + } finally { + handle.release(); } - return new DataFrame(unnestColumns(nativeHandle, columns, options.preserveNulls())); } // -- Set operations ------------------------------------------------------ @@ -365,7 +396,7 @@ public DataFrame unnestColumns(UnnestOptions options, String... columns) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame union(DataFrame other) { - return new DataFrame(unionRows(nativeHandle, otherHandle("union", other))); + return combine("union", other, DataFrame::unionRows); } /** @@ -377,7 +408,7 @@ public DataFrame union(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame unionDistinct(DataFrame other) { - return new DataFrame(unionDistinctRows(nativeHandle, otherHandle("unionDistinct", other))); + return combine("unionDistinct", other, DataFrame::unionDistinctRows); } /** @@ -388,7 +419,7 @@ public DataFrame unionDistinct(DataFrame other) { * @throws RuntimeException if column types disagree on a shared name. */ public DataFrame unionByName(DataFrame other) { - return new DataFrame(unionByNameRows(nativeHandle, otherHandle("unionByName", other))); + return combine("unionByName", other, DataFrame::unionByNameRows); } /** @@ -399,8 +430,7 @@ public DataFrame unionByName(DataFrame other) { * @throws RuntimeException if column types disagree on a shared name. */ public DataFrame unionByNameDistinct(DataFrame other) { - return new DataFrame( - unionByNameDistinctRows(nativeHandle, otherHandle("unionByNameDistinct", other))); + return combine("unionByNameDistinct", other, DataFrame::unionByNameDistinctRows); } /** @@ -420,7 +450,7 @@ public DataFrame unionByNameDistinct(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame intersect(DataFrame other) { - return new DataFrame(intersectRows(nativeHandle, otherHandle("intersect", other))); + return combine("intersect", other, DataFrame::intersectRows); } /** @@ -431,8 +461,7 @@ public DataFrame intersect(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame intersectDistinct(DataFrame other) { - return new DataFrame( - intersectDistinctRows(nativeHandle, otherHandle("intersectDistinct", other))); + return combine("intersectDistinct", other, DataFrame::intersectDistinctRows); } /** @@ -452,7 +481,7 @@ public DataFrame intersectDistinct(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame except(DataFrame other) { - return new DataFrame(exceptRows(nativeHandle, otherHandle("except", other))); + return combine("except", other, DataFrame::exceptRows); } /** @@ -463,24 +492,44 @@ public DataFrame except(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame exceptDistinct(DataFrame other) { - return new DataFrame(exceptDistinctRows(nativeHandle, otherHandle("exceptDistinct", other))); + return combine("exceptDistinct", other, DataFrame::exceptDistinctRows); } /** - * Validate the receiver and the other DataFrame and return {@code other.nativeHandle}. Common - * preamble for the eight set-operation methods so the validation logic stays in one place. + * Validate the receiver and {@code other}, pin both handles, and apply a two-handle native set + * operation. Common body for the eight set-operation methods so the validation and pinning stay + * in one place. + * + *

Pinning two handles cannot deadlock: {@link NativeHandle#acquire()} never blocks, so a + * thread holding the receiver's pin is never waiting on the other DataFrame's lifetime. */ - private long otherHandle(String op, DataFrame other) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (other == null) { - throw new IllegalArgumentException(op + " other must be non-null"); + private DataFrame combine(String op, DataFrame other, LongBinaryOperator nativeOp) { + long left = handle.acquire(); + try { + if (other == null) { + throw new IllegalArgumentException(op + " other must be non-null"); + } + long right = other.acquireAsOperand(op + " other DataFrame is closed or already collected"); + try { + return new DataFrame(nativeOp.applyAsLong(left, right)); + } finally { + other.handle.release(); + } + } finally { + handle.release(); } - if (other.nativeHandle == 0) { - throw new IllegalStateException(op + " other DataFrame is closed or already collected"); + } + + /** + * Pin this DataFrame as the non-receiver operand of a binary operation, reporting a closed handle + * with {@code message} so the caller can distinguish which side was already gone. + */ + private long acquireAsOperand(String message) { + try { + return handle.acquire(); + } catch (IllegalStateException e) { + throw new IllegalStateException(message); } - return other.nativeHandle; } /** @@ -495,25 +544,27 @@ private long otherHandle(String op, DataFrame other) { * @throws RuntimeException if a sort column does not exist in this DataFrame's schema. */ public DataFrame sort(SortExpr... exprs) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (exprs == null) { - throw new IllegalArgumentException("sort exprs must be non-null"); - } - String[] columns = new String[exprs.length]; - boolean[] ascending = new boolean[exprs.length]; - boolean[] nullsFirst = new boolean[exprs.length]; - for (int i = 0; i < exprs.length; i++) { - SortExpr e = exprs[i]; - if (e == null) { - throw new IllegalArgumentException("sort exprs[" + i + "] must be non-null"); + long h = handle.acquire(); + try { + if (exprs == null) { + throw new IllegalArgumentException("sort exprs must be non-null"); } - columns[i] = e.column(); - ascending[i] = e.ascending(); - nullsFirst[i] = e.nullsFirst(); + String[] columns = new String[exprs.length]; + boolean[] ascending = new boolean[exprs.length]; + boolean[] nullsFirst = new boolean[exprs.length]; + for (int i = 0; i < exprs.length; i++) { + SortExpr e = exprs[i]; + if (e == null) { + throw new IllegalArgumentException("sort exprs[" + i + "] must be non-null"); + } + columns[i] = e.column(); + ascending[i] = e.ascending(); + nullsFirst[i] = e.nullsFirst(); + } + return new DataFrame(sortRows(h, columns, ascending, nullsFirst)); + } finally { + handle.release(); } - return new DataFrame(sortRows(nativeHandle, columns, ascending, nullsFirst)); } /** @@ -524,13 +575,15 @@ public DataFrame sort(SortExpr... exprs) { * @throws RuntimeException if the underlying repartition plan rejects the request. */ public DataFrame repartitionRoundRobin(int numPartitions) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (numPartitions <= 0) { - throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions); + long h = handle.acquire(); + try { + if (numPartitions <= 0) { + throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions); + } + return new DataFrame(repartitionRoundRobinRows(h, numPartitions)); + } finally { + handle.release(); } - return new DataFrame(repartitionRoundRobinRows(nativeHandle, numPartitions)); } /** @@ -544,24 +597,26 @@ public DataFrame repartitionRoundRobin(int numPartitions) { * @throws RuntimeException if a partition column does not exist in this DataFrame's schema. */ public DataFrame repartitionHash(int numPartitions, String... columns) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (numPartitions <= 0) { - throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions); - } - if (columns == null) { - throw new IllegalArgumentException("repartitionHash columns must be non-null"); - } - if (columns.length == 0) { - throw new IllegalArgumentException("repartitionHash requires at least one column"); - } - for (int i = 0; i < columns.length; i++) { - if (columns[i] == null) { - throw new IllegalArgumentException("repartitionHash columns[" + i + "] must be non-null"); + long h = handle.acquire(); + try { + if (numPartitions <= 0) { + throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions); + } + if (columns == null) { + throw new IllegalArgumentException("repartitionHash columns must be non-null"); } + if (columns.length == 0) { + throw new IllegalArgumentException("repartitionHash requires at least one column"); + } + for (int i = 0; i < columns.length; i++) { + if (columns[i] == null) { + throw new IllegalArgumentException("repartitionHash columns[" + i + "] must be non-null"); + } + } + return new DataFrame(repartitionHashRows(h, numPartitions, columns)); + } finally { + handle.release(); } - return new DataFrame(repartitionHashRows(nativeHandle, numPartitions, columns)); } /** @@ -580,8 +635,17 @@ public DataFrame repartitionHash(int numPartitions, String... columns) { */ public DataFrame join(DataFrame right, JoinType type, String[] leftCols, String[] rightCols) { checkJoinArgs(right, type, leftCols, rightCols); - return new DataFrame( - joinDataFrame(nativeHandle, right.nativeHandle, type.code(), leftCols, rightCols, null)); + long l = handle.acquire(); + try { + long r = right.acquireAsOperand("right DataFrame is closed or already collected"); + try { + return new DataFrame(joinDataFrame(l, r, type.code(), leftCols, rightCols, null)); + } finally { + right.handle.release(); + } + } finally { + handle.release(); + } } /** @@ -601,11 +665,20 @@ public DataFrame join(DataFrame right, JoinType type, String[] leftCols, String[ public DataFrame join( DataFrame right, JoinType type, String[] leftCols, String[] rightCols, String filter) { checkJoinArgs(right, type, leftCols, rightCols); - if (filter == null) { - throw new IllegalArgumentException("join filter must be non-null"); + long l = handle.acquire(); + try { + long r = right.acquireAsOperand("right DataFrame is closed or already collected"); + try { + if (filter == null) { + throw new IllegalArgumentException("join filter must be non-null"); + } + return new DataFrame(joinDataFrame(l, r, type.code(), leftCols, rightCols, filter)); + } finally { + right.handle.release(); + } + } finally { + handle.release(); } - return new DataFrame( - joinDataFrame(nativeHandle, right.nativeHandle, type.code(), leftCols, rightCols, filter)); } /** @@ -640,17 +713,20 @@ public DataFrame joinOn(DataFrame right, JoinType type, String... predicates) { throw new IllegalArgumentException("joinOn predicates must not contain null"); } } - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (right.nativeHandle == 0) { - throw new IllegalStateException("right DataFrame is closed or already collected"); + long l = handle.acquire(); + try { + long r = right.acquireAsOperand("right DataFrame is closed or already collected"); + try { + return new DataFrame(joinOnDataFrame(l, r, type.code(), predicates)); + } finally { + right.handle.release(); + } + } finally { + handle.release(); } - return new DataFrame( - joinOnDataFrame(nativeHandle, right.nativeHandle, type.code(), predicates)); } - private void checkJoinArgs( + private static void checkJoinArgs( DataFrame right, JoinType type, String[] leftCols, String[] rightCols) { if (right == null) { throw new IllegalArgumentException("join right must be non-null"); @@ -671,12 +747,6 @@ private void checkJoinArgs( + " and " + rightCols.length); } - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (right.nativeHandle == 0) { - throw new IllegalStateException("right DataFrame is closed or already collected"); - } } /** @@ -698,15 +768,17 @@ public void writeParquet(String path) { * etc.). */ public void writeParquet(String path, ParquetWriteOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + writeParquetWithOptions( + h, + path, + options.compression(), + options.singleFileOutput() != null, + options.singleFileOutput() != null && options.singleFileOutput()); + } finally { + handle.release(); } - writeParquetWithOptions( - nativeHandle, - path, - options.compression(), - options.singleFileOutput() != null, - options.singleFileOutput() != null && options.singleFileOutput()); } /** @@ -729,16 +801,18 @@ public void writeCsv(String path) { * etc.). */ public void writeCsv(String path, CsvWriteOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (path == null) { - throw new IllegalArgumentException("writeCsv path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("writeCsv options must be non-null"); + long h = handle.acquire(); + try { + if (path == null) { + throw new IllegalArgumentException("writeCsv path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("writeCsv options must be non-null"); + } + writeCsvWithOptions(h, path, options.toBytes()); + } finally { + handle.release(); } - writeCsvWithOptions(nativeHandle, path, options.toBytes()); } /** @@ -761,23 +835,30 @@ public void writeJson(String path) { * etc.). */ public void writeJson(String path, JsonWriteOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (path == null) { - throw new IllegalArgumentException("writeJson path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("writeJson options must be non-null"); + long h = handle.acquire(); + try { + if (path == null) { + throw new IllegalArgumentException("writeJson path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("writeJson options must be non-null"); + } + writeJsonWithOptions(h, path, options.toBytes()); + } finally { + handle.release(); } - writeJsonWithOptions(nativeHandle, path, options.toBytes()); } + /** + * Release the native plan. Blocks until calls already in flight on this DataFrame return, so that + * no native call can be left dereferencing a freed plan. Idempotent, and a no-op if the DataFrame + * was already consumed by {@link #collect} or {@link #executeStream}. + */ @Override public void close() { - if (nativeHandle != 0) { - closeDataFrame(nativeHandle); - nativeHandle = 0; + long h = handle.claimQuietly(); + if (h != 0) { + closeDataFrame(h); } } diff --git a/core/src/main/java/org/apache/datafusion/NativeHandle.java b/core/src/main/java/org/apache/datafusion/NativeHandle.java new file mode 100644 index 0000000..a5f2731 --- /dev/null +++ b/core/src/main/java/org/apache/datafusion/NativeHandle.java @@ -0,0 +1,172 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +package org.apache.datafusion; + +/** + * Owns a pointer to a native allocation and serialises its lifetime against concurrent use. + * + *

A bare {@code long} field guarded by {@code if (handle == 0) throw} is a time-of-check / + * time-of-use bug: a thread can read a live handle, and a {@code close()} on another thread can + * free the allocation before the first thread's JNI call dereferences it. This class closes that + * window by pinning the handle for the duration of each native call and deferring the free until + * every in-flight call has drained. + * + *

Usage is a pin around each native call: + * + * 

{@code
+ * long h = handle.acquire();
+ * try {
+ *   someNativeCall(h, ...);
+ * } finally {
+ *   handle.release();
+ * }
+ * }
+ * + *

and a claim for the operation that hands the pointer back to Rust to be freed or consumed: + * + * 

{@code
+ * long h = handle.claimQuietly();
+ * if (h != 0) {
+ *   closeNative(h);
+ * }
+ * }
+ * + *

Two properties matter for callers: + * + *

+ * + *

The monitor is held only for bookkeeping, never across a native call, so independent + * operations on the same object still run concurrently. + */ +final class NativeHandle { + + /** Message for the {@link IllegalStateException} raised once the handle is gone. */ + private final String closedMessage; + + /** The native pointer, or 0 once claimed. */ + private long handle; + + /** Number of threads currently inside a native call with this handle pinned. */ + private int inFlight; + + /** + * Set once a claim has begun. New pins are refused from this point on, even though {@link + * #handle} stays readable until the drain completes. + */ + private boolean claimed; + + /** + * @param handle a non-zero native pointer + * @param closedMessage the {@link IllegalStateException} message to use once the handle is gone; + * lets each owner keep its own wording + */ + NativeHandle(long handle, String closedMessage) { + if (handle == 0) { + throw new IllegalArgumentException(closedMessage + ": native handle is null"); + } + this.handle = handle; + this.closedMessage = closedMessage; + } + + /** + * Pin the handle for the duration of one native call and return it. Never blocks. + * + * @throws IllegalStateException if the handle has been claimed + */ + synchronized long acquire() { + if (claimed) { + throw new IllegalStateException(closedMessage); + } + inFlight++; + return handle; + } + + /** + * Release a pin taken by {@link #acquire()}. Must be called from a {@code finally} block so that + * a native call which throws still drains. + */ + synchronized void release() { + if (inFlight == 0) { + throw new IllegalStateException("release() without a matching acquire()"); + } + inFlight--; + if (inFlight == 0 && claimed) { + notifyAll(); + } + } + + /** + * Take exclusive ownership of the handle, blocking until in-flight calls drain, and return it for + * freeing or consumption. Subsequent {@link #acquire()} calls throw. + * + * @throws IllegalStateException if the handle has already been claimed + */ + synchronized long claim() { + if (claimed) { + throw new IllegalStateException(closedMessage); + } + return drainAndTake(); + } + + /** + * As {@link #claim()}, but returns 0 rather than throwing when the handle has already been + * claimed. Used by {@code close()}, which is specified to be idempotent. + */ + synchronized long claimQuietly() { + if (claimed) { + return 0; + } + return drainAndTake(); + } + + /** + * Refuse further pins, wait for the outstanding ones to drain, then surrender the pointer. + * + *

Interruption is deferred rather than obeyed: returning early would either leak the native + * allocation or free it under a call that is still using it. The flag is restored so the caller + * can act on it after the handle is safely accounted for. + */ + private long drainAndTake() { + claimed = true; + boolean interrupted = false; + while (inFlight > 0) { + try { + wait(); + } catch (InterruptedException e) { + interrupted = true; + } + } + long claimedHandle = handle; + handle = 0; + if (interrupted) { + Thread.currentThread().interrupt(); + } + return claimedHandle; + } +} diff --git a/core/src/main/java/org/apache/datafusion/SessionContext.java b/core/src/main/java/org/apache/datafusion/SessionContext.java index b68cda5..df5a745 100644 --- a/core/src/main/java/org/apache/datafusion/SessionContext.java +++ b/core/src/main/java/org/apache/datafusion/SessionContext.java @@ -38,30 +38,41 @@ /** * A DataFusion session context. * - *

Instances are not thread-safe. Concurrent calls to any of {@link #sql}, - * {@link #registerParquet}, or {@link #close} from different threads can produce a use-after-free - * on the native side. Callers must externally synchronize, or confine each context to a single - * thread. + *

Instances are safe to share between threads. The underlying native session is itself + * thread-safe, and this class pins its native handle for the duration of every call, so a {@link + * #close} racing with work on another thread cannot free the session out from under it. A call that + * loses such a race throws {@link IllegalStateException}. + * + *

{@link #close} blocks until calls already in flight on this context have returned, then + * releases the native session. Calling it a second time, concurrently or otherwise, is a no-op. + * + *

Thread safety here covers the handle's lifetime, not the semantics of overlapping operations: + * registering a table concurrently with a query that reads it still races in the ordinary way, and + * whether the query observes the registration is undefined. */ public final class SessionContext implements AutoCloseable { static { NativeLibraryLoader.loadLibrary(); } - private long nativeHandle; + private static final String CLOSED = "SessionContext is closed"; + + private final NativeHandle handle; public SessionContext() { - this.nativeHandle = createSessionContext(); - if (this.nativeHandle == 0) { + long nativeHandle = createSessionContext(); + if (nativeHandle == 0) { throw new RuntimeException("Failed to create native SessionContext"); } + this.handle = new NativeHandle(nativeHandle, CLOSED); } SessionContext(byte[] optionsBytes) { - this.nativeHandle = createSessionContextWithOptions(optionsBytes); - if (this.nativeHandle == 0) { + long nativeHandle = createSessionContextWithOptions(optionsBytes); + if (nativeHandle == 0) { throw new RuntimeException("Failed to create native SessionContext"); } + this.handle = new NativeHandle(nativeHandle, CLOSED); } /** Start configuring a {@link SessionContext}. */ @@ -74,11 +85,12 @@ public static SessionContextBuilder builder() { * until {@link DataFrame#collect} is called. */ public DataFrame sql(String query) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + return new DataFrame(createDataFrame(h, query)); + } finally { + handle.release(); } - long dfHandle = createDataFrame(nativeHandle, query); - return new DataFrame(dfHandle); } /** @@ -92,11 +104,12 @@ public DataFrame sql(String query) { * planning fails. */ public DataFrame fromProto(byte[] planBytes) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + return new DataFrame(createDataFrameFromProto(h, planBytes)); + } finally { + handle.release(); } - long dfHandle = createDataFrameFromProto(nativeHandle, planBytes); - return new DataFrame(dfHandle); } /** @@ -126,14 +139,15 @@ public DataFrame fromProto(byte[] planBytes) { * if the native crate was built without the {@code substrait} feature. */ public DataFrame fromSubstrait(byte[] planBytes) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (planBytes == null) { - throw new IllegalArgumentException("fromSubstrait planBytes must be non-null"); + long h = handle.acquire(); + try { + if (planBytes == null) { + throw new IllegalArgumentException("fromSubstrait planBytes must be non-null"); + } + return new DataFrame(createDataFrameFromSubstrait(h, planBytes)); + } finally { + handle.release(); } - long dfHandle = createDataFrameFromSubstrait(nativeHandle, planBytes); - return new DataFrame(dfHandle); } /** @@ -165,11 +179,13 @@ public DataFrame fromSubstrait(byte[] planBytes) { * (should not happen in practice -- tracker registration is done by the constructor). */ public MemoryUsage memoryUsage() { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + long[] values = memoryUsageNative(h); + return new MemoryUsage(values[0], values[1]); + } finally { + handle.release(); } - long[] values = memoryUsageNative(nativeHandle); - return new MemoryUsage(values[0], values[1]); } /** @@ -195,12 +211,14 @@ public MemoryUsage memoryUsage() { * feature. */ public RuntimeStats runtimeStats() { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + long[] s = runtimeStatsNative(h); + return new RuntimeStats( + (int) s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10]); + } finally { + handle.release(); } - long[] s = runtimeStatsNative(nativeHandle); - return new RuntimeStats( - (int) s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10]); } /** @@ -210,10 +228,13 @@ public RuntimeStats runtimeStats() { * @throws RuntimeException if {@code tableName} is not registered in this context. */ public Schema tableSchema(String tableName) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + byte[] ipcBytes; + long h = handle.acquire(); + try { + ipcBytes = tableSchemaIpc(h, tableName); + } finally { + handle.release(); } - byte[] ipcBytes = tableSchemaIpc(nativeHandle, tableName); try { return MessageSerializer.deserializeSchema( new ReadChannel(Channels.newChannel(new ByteArrayInputStream(ipcBytes)))); @@ -241,13 +262,15 @@ public Schema tableSchema(String tableName) { * @throws IllegalStateException if this context is closed. */ public String getOption(String key) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (key == null) { - throw new IllegalArgumentException("getOption key must be non-null"); + long h = handle.acquire(); + try { + if (key == null) { + throw new IllegalArgumentException("getOption key must be non-null"); + } + return getOptionNative(h, key); + } finally { + handle.release(); } - return getOptionNative(nativeHandle, key); } public void registerCsv(String name, String path) { @@ -261,15 +284,17 @@ public void registerCsv(String name, String path) { * @throws RuntimeException if registration fails (path not found, schema inference error, etc.). */ public void registerCsv(String name, String path, CsvReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + registerCsvWithOptions( + h, + name, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null); + } finally { + handle.release(); } - registerCsvWithOptions( - nativeHandle, - name, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); } /** Read a CSV file as a {@link DataFrame} without registering it. */ @@ -283,16 +308,17 @@ public DataFrame readCsv(String path) { * @throws RuntimeException if the read fails. */ public DataFrame readCsv(String path, CsvReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + return new DataFrame( + readCsvWithOptions( + h, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null)); + } finally { + handle.release(); } - long dfHandle = - readCsvWithOptions( - nativeHandle, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); - return new DataFrame(dfHandle); } public void registerJson(String name, String path) { @@ -306,24 +332,26 @@ public void registerJson(String name, String path) { * @throws RuntimeException if registration fails (path not found, schema inference error, etc.). */ public void registerJson(String name, String path, NdJsonReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (name == null) { - throw new IllegalArgumentException("registerJson name must be non-null"); - } - if (path == null) { - throw new IllegalArgumentException("registerJson path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("registerJson options must be non-null"); + long h = handle.acquire(); + try { + if (name == null) { + throw new IllegalArgumentException("registerJson name must be non-null"); + } + if (path == null) { + throw new IllegalArgumentException("registerJson path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("registerJson options must be non-null"); + } + registerJsonWithOptions( + h, + name, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null); + } finally { + handle.release(); } - registerJsonWithOptions( - nativeHandle, - name, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); } /** Read a newline-delimited JSON file as a {@link DataFrame} without registering it. */ @@ -338,22 +366,23 @@ public DataFrame readJson(String path) { * @throws RuntimeException if the read fails. */ public DataFrame readJson(String path, NdJsonReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (path == null) { - throw new IllegalArgumentException("readJson path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("readJson options must be non-null"); + long h = handle.acquire(); + try { + if (path == null) { + throw new IllegalArgumentException("readJson path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("readJson options must be non-null"); + } + return new DataFrame( + readJsonWithOptions( + h, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null)); + } finally { + handle.release(); } - long dfHandle = - readJsonWithOptions( - nativeHandle, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); - return new DataFrame(dfHandle); } public void registerParquet(String name, String path) { @@ -366,15 +395,17 @@ public void registerParquet(String name, String path) { * @throws RuntimeException if registration fails (path not found, schema mismatch, etc.). */ public void registerParquet(String name, String path, ParquetReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + registerParquetWithOptions( + h, + name, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null); + } finally { + handle.release(); } - registerParquetWithOptions( - nativeHandle, - name, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); } /** Read a parquet file as a {@link DataFrame} without registering it. */ @@ -388,16 +419,17 @@ public DataFrame readParquet(String path) { * @throws RuntimeException if the read fails. */ public DataFrame readParquet(String path, ParquetReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + return new DataFrame( + readParquetWithOptions( + h, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null)); + } finally { + handle.release(); } - long dfHandle = - readParquetWithOptions( - nativeHandle, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); - return new DataFrame(dfHandle); } /** Register an Arrow IPC file (or directory of Arrow IPC files) as a table. */ @@ -414,24 +446,26 @@ public void registerArrow(String name, String path) { * @throws RuntimeException if registration fails (path not found, schema mismatch, etc.). */ public void registerArrow(String name, String path, ArrowReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (name == null) { - throw new IllegalArgumentException("registerArrow name must be non-null"); - } - if (path == null) { - throw new IllegalArgumentException("registerArrow path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("registerArrow options must be non-null"); + long h = handle.acquire(); + try { + if (name == null) { + throw new IllegalArgumentException("registerArrow name must be non-null"); + } + if (path == null) { + throw new IllegalArgumentException("registerArrow path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("registerArrow options must be non-null"); + } + registerArrowWithOptions( + h, + name, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null); + } finally { + handle.release(); } - registerArrowWithOptions( - nativeHandle, - name, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); } /** Read an Arrow IPC file as a {@link DataFrame} without registering it. */ @@ -446,22 +480,23 @@ public DataFrame readArrow(String path) { * @throws RuntimeException if the read fails. */ public DataFrame readArrow(String path, ArrowReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (path == null) { - throw new IllegalArgumentException("readArrow path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("readArrow options must be non-null"); + long h = handle.acquire(); + try { + if (path == null) { + throw new IllegalArgumentException("readArrow path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("readArrow options must be non-null"); + } + return new DataFrame( + readArrowWithOptions( + h, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null)); + } finally { + handle.release(); } - long dfHandle = - readArrowWithOptions( - nativeHandle, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); - return new DataFrame(dfHandle); } /** Register an Avro file (or directory of Avro files) as a table. */ @@ -478,24 +513,26 @@ public void registerAvro(String name, String path) { * @throws RuntimeException if registration fails (path not found, schema mismatch, etc.). */ public void registerAvro(String name, String path, AvroReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (name == null) { - throw new IllegalArgumentException("registerAvro name must be non-null"); - } - if (path == null) { - throw new IllegalArgumentException("registerAvro path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("registerAvro options must be non-null"); + long h = handle.acquire(); + try { + if (name == null) { + throw new IllegalArgumentException("registerAvro name must be non-null"); + } + if (path == null) { + throw new IllegalArgumentException("registerAvro path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("registerAvro options must be non-null"); + } + registerAvroWithOptions( + h, + name, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null); + } finally { + handle.release(); } - registerAvroWithOptions( - nativeHandle, - name, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); } /** Read an Avro file as a {@link DataFrame} without registering it. */ @@ -510,22 +547,23 @@ public DataFrame readAvro(String path) { * @throws RuntimeException if the read fails. */ public DataFrame readAvro(String path, AvroReadOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (path == null) { - throw new IllegalArgumentException("readAvro path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("readAvro options must be non-null"); + long h = handle.acquire(); + try { + if (path == null) { + throw new IllegalArgumentException("readAvro path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("readAvro options must be non-null"); + } + return new DataFrame( + readAvroWithOptions( + h, + path, + options.toBytes(), + options.schema() != null ? serializeSchemaIpc(options.schema()) : null)); + } finally { + handle.release(); } - long dfHandle = - readAvroWithOptions( - nativeHandle, - path, - options.toBytes(), - options.schema() != null ? serializeSchemaIpc(options.schema()) : null); - return new DataFrame(dfHandle); } /** @@ -539,19 +577,21 @@ public DataFrame readAvro(String path, AvroReadOptions options) { * incompatible signature, schema serialisation failure). */ public void registerUdf(ScalarUdf udf) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + java.util.Objects.requireNonNull(udf, "udf"); + ScalarFunction impl = udf.impl(); + String name = udf.name(); + Volatility volatility = udf.volatility(); + List fields = new ArrayList<>(udf.argFields().size() + 1); + fields.add(udf.returnField()); + fields.addAll(udf.argFields()); + Schema signatureSchema = new Schema(fields); + byte[] signatureBytes = serializeSchemaIpc(signatureSchema); + registerScalarUdf(h, name, signatureBytes, volatility.code(), impl); + } finally { + handle.release(); } - java.util.Objects.requireNonNull(udf, "udf"); - ScalarFunction impl = udf.impl(); - String name = udf.name(); - Volatility volatility = udf.volatility(); - List fields = new ArrayList<>(udf.argFields().size() + 1); - fields.add(udf.returnField()); - fields.addAll(udf.argFields()); - Schema signatureSchema = new Schema(fields); - byte[] signatureBytes = serializeSchemaIpc(signatureSchema); - registerScalarUdf(nativeHandle, name, signatureBytes, volatility.code(), impl); } /** @@ -572,21 +612,23 @@ public void registerUdf(ScalarUdf udf) { * @throws RuntimeException if native registration fails. */ public void registerTable(String name, TableProvider provider) { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); - } - if (name == null) { - throw new IllegalArgumentException("registerTable name must be non-null"); - } - if (provider == null) { - throw new IllegalArgumentException("registerTable provider must be non-null"); - } - Schema schema = provider.schema(); - if (schema == null) { - throw new IllegalStateException("TableProvider.schema returned null"); + long h = handle.acquire(); + try { + if (name == null) { + throw new IllegalArgumentException("registerTable name must be non-null"); + } + if (provider == null) { + throw new IllegalArgumentException("registerTable provider must be non-null"); + } + Schema schema = provider.schema(); + if (schema == null) { + throw new IllegalStateException("TableProvider.schema returned null"); + } + byte[] schemaIpc = serializeSchemaIpc(schema); + registerTableNative(h, name, schemaIpc, provider); + } finally { + handle.release(); } - byte[] schemaIpc = serializeSchemaIpc(schema); - registerTableNative(nativeHandle, name, schemaIpc, provider); } private static byte[] serializeSchemaIpc(Schema schema) { @@ -610,11 +652,15 @@ private static byte[] serializeSchemaIpc(Schema schema) { * @throws IllegalStateException if this context is closed. */ public boolean tableExists(String name) { - checkOpenSessionContext(); - if (name == null) { - throw new IllegalArgumentException("tableExists name must be non-null"); + long h = handle.acquire(); + try { + if (name == null) { + throw new IllegalArgumentException("tableExists name must be non-null"); + } + return tableExists(h, name); + } finally { + handle.release(); } - return tableExists(nativeHandle, name); } /** @@ -626,24 +672,28 @@ public boolean tableExists(String name) { * @throws IllegalStateException if this context is closed. */ public void deregisterTable(String name) { - checkOpenSessionContext(); - if (name == null) { - throw new IllegalArgumentException("deregisterTable name must be non-null"); - } - deregisterTable(nativeHandle, name); - } - - private void checkOpenSessionContext() { - if (nativeHandle == 0) { - throw new IllegalStateException("SessionContext is closed"); + long h = handle.acquire(); + try { + if (name == null) { + throw new IllegalArgumentException("deregisterTable name must be non-null"); + } + deregisterTable(h, name); + } finally { + handle.release(); } } + /** + * Release the native session. Blocks until calls already in flight on this context return, so + * that no native call can be left dereferencing a freed session. Idempotent, and safe to call + * concurrently with work on other threads: those calls either complete or throw {@link + * IllegalStateException}. + */ @Override public void close() { - if (nativeHandle != 0) { - closeSessionContext(nativeHandle); - nativeHandle = 0; + long h = handle.claimQuietly(); + if (h != 0) { + closeSessionContext(h); } } diff --git a/core/src/test/java/org/apache/datafusion/DataFrameConcurrencyTest.java b/core/src/test/java/org/apache/datafusion/DataFrameConcurrencyTest.java new file mode 100644 index 0000000..d4e404e --- /dev/null +++ b/core/src/test/java/org/apache/datafusion/DataFrameConcurrencyTest.java @@ -0,0 +1,258 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +package org.apache.datafusion; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.channels.Channels; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +/** + * Concurrency contract for {@link DataFrame}: {@link DataFrame#close()} and the consuming + * operations must not release the native plan while another thread is inside a JNI call on it. See + * issue #40. + */ +class DataFrameConcurrencyTest { + + /** Exactly one of several threads racing to consume a DataFrame may win. */ + @Test + @Timeout(120) + void concurrentCollectYieldsExactlyOneWinner() throws Exception { + try (BufferAllocator allocator = new RootAllocator(); + SessionContext ctx = new SessionContext()) { + for (int round = 0; round < 20; round++) { + DataFrame df = ctx.sql("select 1 as a"); + int threads = 6; + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + AtomicInteger winners = new AtomicInteger(); + List unexpected = Collections.synchronizedList(new ArrayList<>()); + List readers = Collections.synchronizedList(new ArrayList<>()); + + for (int i = 0; i < threads; i++) { + new Thread( + () -> { + try { + start.await(); + readers.add(df.collect(allocator)); + winners.incrementAndGet(); + } catch (IllegalStateException lost) { + // Another thread consumed the DataFrame first; that is the contract. + } catch (Throwable t) { + unexpected.add(t); + } finally { + done.countDown(); + } + }, + "collector-" + i) + .start(); + } + + start.countDown(); + assertTrue(done.await(60, TimeUnit.SECONDS)); + assertTrue(unexpected.isEmpty(), "unexpected failures: " + unexpected); + assertEquals(1, winners.get(), "a DataFrame must be consumable exactly once"); + for (ArrowReader reader : readers) { + reader.close(); + } + df.close(); + } + } + } + + /** + * {@code close()} must block until an in-flight execution returns. A {@link TableProvider} whose + * {@code scan} parks on a latch holds the native call open for as long as the test needs. + */ + @Test + @Timeout(120) + void closeWaitsForAnInFlightExecution() throws Exception { + try (SessionContext ctx = new SessionContext()) { + LatchedTableProvider provider = new LatchedTableProvider(); + ctx.registerTable("t", provider); + DataFrame df = ctx.sql("select * from t"); + + Thread counter = new Thread(df::count, "counter"); + counter.start(); + assertTrue(provider.insideScan.await(60, TimeUnit.SECONDS), "scan() was never reached"); + + Thread closer = new Thread(df::close, "closer"); + closer.start(); + closer.join(500); + assertTrue(closer.isAlive(), "close() must not free the plan while a call is in flight"); + + provider.releaseScan.countDown(); + counter.join(60_000); + closer.join(60_000); + assertFalse(counter.isAlive()); + assertFalse(closer.isAlive()); + + assertThrows(IllegalStateException.class, df::count); + } + } + + /** Non-consuming operations are shared, not serialised: concurrent readers all succeed. */ + @Test + @Timeout(120) + void concurrentNonConsumingOperationsAllSucceed() throws Exception { + try (SessionContext ctx = new SessionContext(); + DataFrame df = ctx.sql("select 1 as a")) { + int threads = 6; + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + List unexpected = Collections.synchronizedList(new ArrayList<>()); + + for (int i = 0; i < threads; i++) { + new Thread( + () -> { + try { + start.await(); + for (int j = 0; j < 20; j++) { + assertEquals(1, df.count()); + assertEquals(1, df.schema().getFields().size()); + } + } catch (Throwable t) { + unexpected.add(t); + } finally { + done.countDown(); + } + }, + "reader-" + i) + .start(); + } + + start.countDown(); + assertTrue(done.await(60, TimeUnit.SECONDS)); + assertTrue(unexpected.isEmpty(), "unexpected failures: " + unexpected); + } + } + + /** + * The two-handle set operations pin both DataFrames at once. Pins never block, so opposing pin + * orders across threads cannot deadlock -- this test hangs (and times out) if that ever changes. + */ + @Test + @Timeout(120) + void opposingSetOperationsDoNotDeadlock() throws Exception { + try (SessionContext ctx = new SessionContext(); + DataFrame left = ctx.sql("select 1 as a"); + DataFrame right = ctx.sql("select 2 as a")) { + int threads = 6; + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + List unexpected = Collections.synchronizedList(new ArrayList<>()); + + for (int i = 0; i < threads; i++) { + boolean forward = i % 2 == 0; + new Thread( + () -> { + try { + start.await(); + for (int j = 0; j < 20; j++) { + try (DataFrame u = forward ? left.union(right) : right.union(left)) { + assertEquals(2, u.count()); + } + } + } catch (Throwable t) { + unexpected.add(t); + } finally { + done.countDown(); + } + }, + "unioner-" + i) + .start(); + } + + start.countDown(); + assertTrue(done.await(60, TimeUnit.SECONDS), "set operations deadlocked"); + assertTrue(unexpected.isEmpty(), "unexpected failures: " + unexpected); + } + } + + /** A {@link TableProvider} whose scan parks until the test releases it. */ + private static final class LatchedTableProvider implements TableProvider { + private final Schema schema = + new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); + private final byte[] emptyStream = emptyIpcStream(schema); + private final CountDownLatch insideScan = new CountDownLatch(1); + private final CountDownLatch releaseScan = new CountDownLatch(1); + + @Override + public Schema schema() { + return schema; + } + + @Override + public ArrowReader scan(BufferAllocator allocator) { + insideScan.countDown(); + boolean interrupted = false; + while (true) { + try { + releaseScan.await(); + break; + } catch (InterruptedException e) { + interrupted = true; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + return new ArrowStreamReader(new ByteArrayInputStream(emptyStream), allocator); + } + + /** An Arrow IPC stream carrying {@code schema} and zero batches. */ + private static byte[] emptyIpcStream(Schema schema) { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (BufferAllocator tmp = new RootAllocator(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, tmp); + ArrowStreamWriter writer = new ArrowStreamWriter(root, null, Channels.newChannel(baos))) { + writer.start(); + writer.end(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + return baos.toByteArray(); + } + } +} diff --git a/core/src/test/java/org/apache/datafusion/NativeHandleTest.java b/core/src/test/java/org/apache/datafusion/NativeHandleTest.java new file mode 100644 index 0000000..3b1c6c5 --- /dev/null +++ b/core/src/test/java/org/apache/datafusion/NativeHandleTest.java @@ -0,0 +1,237 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +package org.apache.datafusion; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +/** + * Unit tests for {@link NativeHandle}. These use fabricated handle values and never call into the + * native library, so they exercise the lifetime state machine in isolation. + */ +class NativeHandleTest { + + private static final long HANDLE = 0xDEADBEEFL; + + private static NativeHandle newHandle() { + return new NativeHandle(HANDLE, "test handle is closed"); + } + + @Test + void rejectsZeroHandleAtConstruction() { + assertThrows(IllegalArgumentException.class, () -> new NativeHandle(0, "unused")); + } + + @Test + void acquireReturnsHandleAndIsReentrantAcrossPins() { + NativeHandle h = newHandle(); + assertEquals(HANDLE, h.acquire()); + assertEquals(HANDLE, h.acquire()); + h.release(); + h.release(); + assertEquals(HANDLE, h.claimQuietly()); + } + + @Test + void acquireAfterClaimThrowsWithSuppliedMessage() { + NativeHandle h = newHandle(); + assertEquals(HANDLE, h.claim()); + IllegalStateException e = assertThrows(IllegalStateException.class, h::acquire); + assertEquals("test handle is closed", e.getMessage()); + } + + @Test + void claimYieldsHandleExactlyOnce() { + NativeHandle h = newHandle(); + assertEquals(HANDLE, h.claim()); + assertThrows(IllegalStateException.class, h::claim); + } + + @Test + void claimQuietlyReturnsZeroOnceClaimed() { + NativeHandle h = newHandle(); + assertEquals(HANDLE, h.claimQuietly()); + assertEquals(0L, h.claimQuietly()); + } + + @Test + void releaseWithoutAcquireIsRejected() { + NativeHandle h = newHandle(); + assertThrows(IllegalStateException.class, h::release); + } + + /** + * The core guarantee: a claim (i.e. {@code close()} or a consuming operation) must not hand back + * the raw handle for freeing while another thread is inside a JNI call that holds a pin. + */ + @Test + @Timeout(10) + void claimBlocksUntilInFlightPinsDrain() throws Exception { + NativeHandle h = newHandle(); + assertEquals(HANDLE, h.acquire()); + + CountDownLatch claimStarted = new CountDownLatch(1); + AtomicLong claimed = new AtomicLong(-1); + Thread closer = + new Thread( + () -> { + claimStarted.countDown(); + claimed.set(h.claim()); + }); + closer.start(); + + assertTrue(claimStarted.await(5, TimeUnit.SECONDS)); + // Give the closer a chance to reach the drain wait, then confirm it is still parked. + closer.join(200); + assertTrue(closer.isAlive(), "claim() must not return while a pin is held"); + assertEquals(-1, claimed.get()); + + // A late arrival is rejected immediately rather than queueing behind the claim. + assertThrows(IllegalStateException.class, h::acquire); + + h.release(); + closer.join(5000); + assertFalse(closer.isAlive()); + assertEquals(HANDLE, claimed.get()); + } + + /** + * A claim in progress must survive interruption -- abandoning the drain would leak the native + * allocation or, worse, free it under an in-flight call. The interrupt is deferred to the caller. + */ + @Test + @Timeout(10) + void claimAbsorbsInterruptionAndRestoresTheFlag() throws Exception { + NativeHandle h = newHandle(); + assertEquals(HANDLE, h.acquire()); + + CountDownLatch started = new CountDownLatch(1); + AtomicLong claimed = new AtomicLong(-1); + AtomicReference interrupted = new AtomicReference<>(); + Thread closer = + new Thread( + () -> { + started.countDown(); + claimed.set(h.claim()); + interrupted.set(Thread.currentThread().isInterrupted()); + }); + closer.start(); + + assertTrue(started.await(5, TimeUnit.SECONDS)); + closer.join(200); + closer.interrupt(); + + closer.join(200); + assertTrue(closer.isAlive(), "interrupt must not abandon the drain wait"); + + h.release(); + closer.join(5000); + assertFalse(closer.isAlive()); + assertEquals(HANDLE, claimed.get()); + assertTrue(interrupted.get(), "interrupt status must be restored for the caller"); + } + + /** Only one of many racing claimers may take the handle. */ + @Test + @Timeout(30) + void concurrentClaimsYieldExactlyOneWinner() throws Exception { + for (int round = 0; round < 100; round++) { + NativeHandle h = newHandle(); + int threads = 8; + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + AtomicLong winners = new AtomicLong(); + for (int i = 0; i < threads; i++) { + new Thread( + () -> { + try { + start.await(); + if (h.claimQuietly() != 0) { + winners.incrementAndGet(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + done.countDown(); + } + }) + .start(); + } + start.countDown(); + assertTrue(done.await(10, TimeUnit.SECONDS)); + assertEquals(1, winners.get()); + } + } + + /** + * Pins taken concurrently with a claim either succeed outright or fail; none observe a torn + * state. + */ + @Test + @Timeout(30) + void concurrentPinsNeverOutliveAClaim() throws Exception { + for (int round = 0; round < 100; round++) { + NativeHandle h = newHandle(); + int threads = 8; + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + for (int i = 0; i < threads; i++) { + boolean claimer = i == 0; + new Thread( + () -> { + try { + start.await(); + if (claimer) { + h.claimQuietly(); + } else { + long raw = h.acquire(); + try { + assertEquals(HANDLE, raw); + } finally { + h.release(); + } + } + } catch (IllegalStateException expected) { + // Lost the race with the claim; that is the contract. + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + done.countDown(); + } + }) + .start(); + } + start.countDown(); + assertTrue(done.await(10, TimeUnit.SECONDS)); + // Whatever the interleaving, the handle is gone afterwards. + assertEquals(0L, h.claimQuietly()); + } + } +} diff --git a/core/src/test/java/org/apache/datafusion/SessionContextConcurrencyTest.java b/core/src/test/java/org/apache/datafusion/SessionContextConcurrencyTest.java new file mode 100644 index 0000000..5d12421 --- /dev/null +++ b/core/src/test/java/org/apache/datafusion/SessionContextConcurrencyTest.java @@ -0,0 +1,196 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + * + * http://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. + */ + +package org.apache.datafusion; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +/** + * Concurrency contract for {@link SessionContext}: a {@link SessionContext#close()} racing with + * work on other threads must never free the native session out from under an in-flight JNI call. + * See issue #40. + */ +class SessionContextConcurrencyTest { + + @Test + void useAfterCloseThrowsIllegalState() { + SessionContext ctx = new SessionContext(); + ctx.close(); + assertThrows(IllegalStateException.class, () -> ctx.sql("select 1")); + assertThrows(IllegalStateException.class, () -> ctx.tableExists("t")); + assertThrows(IllegalStateException.class, ctx::memoryUsage); + } + + /** + * The central guarantee. {@code registerTable} calls {@link TableProvider#schema()} on the + * calling thread while the native handle is pinned, which gives us a deterministic way to hold a + * call open across a concurrent {@code close()}. + */ + @Test + @Timeout(60) + void closeWaitsForAnInFlightCall() throws Exception { + SessionContext ctx = new SessionContext(); + CountDownLatch insideSchema = new CountDownLatch(1); + CountDownLatch releaseSchema = new CountDownLatch(1); + + TableProvider blocking = + new TableProvider() { + @Override + public Schema schema() { + insideSchema.countDown(); + awaitUninterruptibly(releaseSchema); + return new Schema( + Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true)))); + } + + @Override + public ArrowReader scan(BufferAllocator allocator) { + throw new UnsupportedOperationException("not scanned by this test"); + } + }; + + Thread registrar = new Thread(() -> ctx.registerTable("t", blocking), "registrar"); + registrar.start(); + assertTrue(insideSchema.await(30, TimeUnit.SECONDS), "registerTable never reached schema()"); + + Thread closer = new Thread(ctx::close, "closer"); + closer.start(); + closer.join(500); + assertTrue(closer.isAlive(), "close() must not free the session while a call is in flight"); + + releaseSchema.countDown(); + registrar.join(30_000); + closer.join(30_000); + assertFalse(registrar.isAlive()); + assertFalse(closer.isAlive()); + + assertThrows(IllegalStateException.class, () -> ctx.sql("select 1")); + } + + /** + * Hammer a context from several threads while closing it. Before the fix this raced on a plain + * {@code long} field and could dereference a freed {@code SessionContext}; the only failure + * permitted now is {@link IllegalStateException} from losing the race. + */ + @Test + @Timeout(300) + void concurrentQueriesRacingCloseOnlyEverFailWithIllegalState() throws Exception { + for (int round = 0; round < 5; round++) { + SessionContext ctx = new SessionContext(); + int threads = 6; + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + CountDownLatch firstSuccess = new CountDownLatch(1); + List unexpected = Collections.synchronizedList(new ArrayList<>()); + + for (int i = 0; i < threads; i++) { + new Thread( + () -> { + try { + start.await(); + for (int j = 0; j < 20; j++) { + try (DataFrame df = ctx.sql("select 1")) { + df.count(); + firstSuccess.countDown(); + } catch (IllegalStateException closedUnderUs) { + return; + } + } + } catch (Throwable t) { + unexpected.add(t); + } finally { + done.countDown(); + } + }, + "querier-" + i) + .start(); + } + + start.countDown(); + // Let real work overlap the close rather than closing an idle context. + firstSuccess.await(30, TimeUnit.SECONDS); + ctx.close(); + + assertTrue(done.await(120, TimeUnit.SECONDS), "workers did not finish"); + assertTrue(unexpected.isEmpty(), "unexpected failures: " + unexpected); + } + } + + /** {@code close()} is idempotent even when several threads call it at once. */ + @Test + @Timeout(60) + void concurrentCloseIsIdempotent() throws Exception { + SessionContext ctx = new SessionContext(); + int threads = 8; + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + List unexpected = Collections.synchronizedList(new ArrayList<>()); + + for (int i = 0; i < threads; i++) { + new Thread( + () -> { + try { + start.await(); + ctx.close(); + } catch (Throwable t) { + unexpected.add(t); + } finally { + done.countDown(); + } + }, + "closer-" + i) + .start(); + } + + start.countDown(); + assertTrue(done.await(30, TimeUnit.SECONDS)); + assertTrue(unexpected.isEmpty(), "close() must be idempotent, saw: " + unexpected); + } + + private static void awaitUninterruptibly(CountDownLatch latch) { + boolean interrupted = false; + while (true) { + try { + latch.await(); + break; + } catch (InterruptedException e) { + interrupted = true; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } +} diff --git a/docs/source/user-guide/quickstart.md b/docs/source/user-guide/quickstart.md index 7d0df38..a569814 100644 --- a/docs/source/user-guide/quickstart.md +++ b/docs/source/user-guide/quickstart.md @@ -55,8 +55,9 @@ per application) and close it in a `try`-with-resources. **Session context.** `SessionContext` is the entry point into DataFusion. It holds the catalog of registered tables and the query planner. It is -`AutoCloseable` and **not thread-safe** — use one per thread, or guard -access externally. +`AutoCloseable` and safe to share across threads; see +[SessionContext](sessioncontext.md) for what that does and does not +guarantee. **Registering data.** `registerParquet(name, path)` reads the file's footer on call and exposes it under the given table name. See diff --git a/docs/source/user-guide/sessioncontext.md b/docs/source/user-guide/sessioncontext.md index f2408c9..6d118e0 100644 --- a/docs/source/user-guide/sessioncontext.md +++ b/docs/source/user-guide/sessioncontext.md @@ -36,9 +36,23 @@ on exception. ## Threading -A `SessionContext` is **not thread-safe**. Do not share one across threads -without external synchronization. The simplest pattern is one context per -thread. +A `SessionContext` is safe to share across threads. Queries and +registrations may run concurrently, and `close()` is safe to call while +other threads are still using the context: it waits for calls already in +flight to return before releasing the native session, and any call that +arrives after the close throws `IllegalStateException`. Closing more than +once, or from several threads at once, is a no-op after the first. + +This covers the *lifetime* of the native session, not the ordering of +overlapping operations. Registering a table concurrently with a query that +reads it still races in the ordinary way — whether the query observes the +registration is undefined. Sequence those calls yourself when the ordering +matters. + +`DataFrame` carries the same guarantees. Because `collect` and +`executeStream` consume the DataFrame, threads racing to consume the same +one resolve cleanly: exactly one succeeds and the rest get +`IllegalStateException`. ## Configuration