From 83bfb4a8f21a0f5cc654f72d5cda4b0f1175107b Mon Sep 17 00:00:00 2001 From: bharath-techie Date: Wed, 6 May 2026 23:37:09 +0530 Subject: [PATCH 1/2] changes to handle session context handle close Signed-off-by: bharath-techie --- .../backend/jni/ConsumableNativeHandle.java | 86 +++++++++++++++++++ .../spi/BackendExecutionContext.java | 26 ++++-- .../rust/src/ffm.rs | 14 ++- .../rust/src/query_executor.rs | 12 +-- .../be/datafusion/DataFusionSessionState.java | 17 +++- .../be/datafusion/DatafusionContext.java | 14 ++- .../be/datafusion/DatafusionSearcher.java | 8 +- .../be/datafusion/nativelib/NativeBridge.java | 38 ++++++-- .../nativelib/SessionContextHandle.java | 36 ++++---- .../DataFusionNativeBridgeTests.java | 14 ++- .../exec/AnalyticsSearchService.java | 12 ++- .../exec/stage/LocalStageScheduler.java | 37 +++++++- 12 files changed, 268 insertions(+), 46 deletions(-) create mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandle.java diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandle.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandle.java new file mode 100644 index 0000000000000..2618ac6735d18 --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandle.java @@ -0,0 +1,86 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.backend.jni; + +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Specialisation of {@link NativeHandle} for pointers whose ownership is transferred to the + * native side by a specific FFM call (for example, Rust's {@code Box::from_raw} inside a + * consuming function). After the consuming call the native resource is freed internally; + * calling the matching {@code close_X} entry a second time would be a double-free, while + * not calling it on the error path would leak. + * + *

The bridge method that performs the consuming FFM call must invoke + * {@link #markConsumed()} after the downcall returns (typically in a {@code finally} block). + * This: + *

+ * + *

On paths where the consuming call never happened (pre-dispatch Java error, aborted flow, + * Cleaner-at-GC fallback), {@link #doClose()} delegates to {@link #doCloseNative()} which + * subclasses implement to free the native resource via the appropriate {@code close_X} FFM entry. + * + *

{@link #markConsumed()} is idempotent and safe to call after {@link #close()}. + */ +public abstract class ConsumableNativeHandle extends NativeHandle { + + /** + * Set once the native side has taken ownership of {@link #ptr} via the consuming FFM call. + * When {@code true}, {@link #doClose()} skips the call to {@link #doCloseNative()} to avoid + * a double-free. + */ + private final AtomicBoolean consumed = new AtomicBoolean(false); + + protected ConsumableNativeHandle(long ptr) { + super(ptr); + } + + /** + * Marks this handle as having had its native pointer consumed by the bridge's + * ownership-transferring FFM call, then closes the Java wrapper. See the class javadoc + * for the full contract and typical call pattern. + */ + public final void markConsumed() { + consumed.set(true); + close(); + } + + /** + * @return {@code true} if {@link #markConsumed()} has been called. + */ + protected final boolean isConsumed() { + return consumed.get(); + } + + /** + * Template method: short-circuits to a no-op when {@link #isConsumed()} is {@code true} + * (the native side already freed the resource), otherwise delegates to + * {@link #doCloseNative()}. Marked {@code final} so subclasses cannot bypass the guard. + */ + @Override + protected final void doClose() { + if (isConsumed()) { + return; + } + doCloseNative(); + } + + /** + * Releases the native resource via the appropriate {@code close_X} FFM entry. + * Called by {@link #doClose()} only when the handle has not been marked consumed, + * i.e. on the error / never-executed path. Must be safe to call at most once per pointer. + */ + protected abstract void doCloseNative(); +} diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/BackendExecutionContext.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/BackendExecutionContext.java index ac3ca2508a2c7..cffa1e972ef9a 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/BackendExecutionContext.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/BackendExecutionContext.java @@ -8,15 +8,31 @@ package org.opensearch.analytics.spi; +import java.io.IOException; + /** - * Marker interface for backend-specific execution context that flows between - * successive instruction handler calls. The first handler in the chain receives - * {@code null} and bootstraps the context; subsequent handlers receive and build - * upon the previous handler's output. + * Backend-specific execution context that flows between successive instruction handler + * calls. The first handler in the chain receives {@code null} and bootstraps the context; + * subsequent handlers receive and build upon the previous handler's output. * *

Each backend defines its own concrete implementation (e.g., * {@code DataFusionSessionState} holding a native SessionContext handle). * + *

Lifecycle

+ *

Extends {@link AutoCloseable} with a narrowed {@code throws IOException} signature so + * backends can attach native / resource-holding handles to the context and rely on the + * orchestrator (e.g. {@code AnalyticsSearchService} or {@code LocalStageScheduler}) to + * close it if the fragment aborts before ownership is transferred to the + * {@code SearchExecEngine}. Implementations that hold no resources should leave the default + * no-op {@link #close()}. {@code close()} must be idempotent; in particular it must + * tolerate being called after the resources have already been handed off to a + * successfully-constructed engine. + * * @opensearch.internal */ -public interface BackendExecutionContext {} +public interface BackendExecutionContext extends AutoCloseable { + @Override + default void close() throws IOException { + // Default: no resources to release. + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 85404ff2f49ce..c208006eb871b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -598,12 +598,24 @@ pub unsafe extern "C" fn df_execute_with_context( plan_ptr: *const u8, plan_len: i64, ) -> i64 { + // Consume the session context handle on entry. Ownership transfers here + // regardless of whether the remainder of this function succeeds, returns an + // error via `?`, or panics — RAII (or `catch_unwind` drop-during-unwind) + // drops `session_handle` and frees the underlying SessionContext resources. + // + // This matches the Java-side contract: SessionContextHandle.markConsumed() is + // invoked in a `finally` after the FFM downcall, so every observable path from + // Java's perspective ("call.invoke ran") maps to "Rust consumed the handle". + // If we were to run fallible or panic-prone code (e.g. `get_rt_manager()?`) + // before Box::from_raw, the handle would leak on those paths. + let session_handle = *Box::from_raw(session_ctx_ptr as *mut crate::session_context::SessionContextHandle); + let mgr = get_rt_manager()?; let plan_bytes = slice::from_raw_parts(plan_ptr, plan_len as usize); let cpu_executor = mgr.cpu_executor(); mgr.io_runtime .block_on(crate::query_executor::execute_with_context( - session_ctx_ptr, + session_handle, plan_bytes, cpu_executor, )) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs index 14c8d172add9a..33d47b9bb556a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs @@ -156,14 +156,16 @@ pub async fn execute_query( } /// Executes a Substrait plan against a pre-configured SessionContext. -/// Consumes the handle — SessionContext lifetime is tied to the returned stream. -pub async unsafe fn execute_with_context( - session_ctx_ptr: i64, +/// +/// Takes ownership of the handle by value. The ownership transfer (consuming the +/// raw Java pointer) happens at the FFM entry in `df_execute_with_context`, so +/// by the time this function is reached the pointer is already invalidated from +/// Java's perspective and cleanup is pure RAII. +pub async fn execute_with_context( + handle: SessionContextHandle, plan_bytes: &[u8], cpu_executor: DedicatedExecutor, ) -> Result { - let handle = *Box::from_raw(session_ctx_ptr as *mut SessionContextHandle); - let substrait_plan = Plan::decode(plan_bytes).map_err(|e| { DataFusionError::Execution(format!("Failed to decode Substrait: {}", e)) })?; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionSessionState.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionSessionState.java index edd48ddb11f22..c807dcf3978a5 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionSessionState.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionSessionState.java @@ -12,8 +12,21 @@ import org.opensearch.be.datafusion.nativelib.SessionContextHandle; /** - * Backend-specific execution context produced by ShardScanInstructionHandler, - * consumed by DatafusionSearcher at execute time. + * Backend-specific execution context produced by {@link ShardScanInstructionHandler}, + * consumed by {@link DatafusionSearcher} at execute time. + * + *

{@link #close()} closes the underlying {@link SessionContextHandle} as the + * fragment-orchestrator's safety net for error paths that never reach the execute step. + * The handle's close is idempotent and cooperates with {@link DatafusionContext#close()} + * (which also closes it once the handle is handed off to an engine), so it is safe to call + * from both places — whichever runs first wins. */ public record DataFusionSessionState(SessionContextHandle sessionContextHandle) implements BackendExecutionContext { + + @Override + public void close() { + if (sessionContextHandle != null) { + sessionContextHandle.close(); + } + } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java index cd30c64e0758e..1d7a17352f4ff 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java @@ -54,7 +54,19 @@ public void close() throws IOException { streamHandle = null; } } finally { - engineSearcher.close(); + try { + // Safety net for aborted-search paths: if the SessionContext was created but + // executeWithContextAsync never ran (or ran and the context is being closed + // without handing off the handle), doClose() calls df_close_session_context. + // On the happy path the handle is already marked consumed and this close() + // is a no-op. + if (sessionContextHandle != null) { + sessionContextHandle.close(); + sessionContextHandle = null; + } + } finally { + engineSearcher.close(); + } } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearcher.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearcher.java index c8bb98991f10e..b6f8abc339101 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearcher.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearcher.java @@ -58,7 +58,7 @@ private void searchWithSessionContext(DatafusionContext context, SessionContextH DatafusionQuery query = context.getDatafusionQuery(); NativeRuntimeHandle runtimeHandle = context.getNativeRuntime(); CompletableFuture future = new CompletableFuture<>(); - NativeBridge.executeWithContextAsync(sessionCtx.getPointer(), query.getSubstraitBytes(), new ActionListener<>() { + NativeBridge.executeWithContextAsync(sessionCtx, query.getSubstraitBytes(), new ActionListener<>() { @Override public void onResponse(Long streamPtr) { future.complete(streamPtr); @@ -75,8 +75,10 @@ public void onFailure(Exception exception) { } catch (Exception exception) { throw new IOException("Query execution with session context failed", exception); } - // Rust consumed the session context — unregister from live handle set - sessionCtx.close(); + // NativeBridge#executeWithContextAsync has already marked the handle consumed (which + // closes the Java wrapper) on both success and native-error paths; no explicit close + // is needed here. The owning DatafusionContext#close() closes it as a safety net for + // paths that never reach this method (e.g. aborted search). context.setStreamHandle(new StreamHandle(streamPtr, runtimeHandle)); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index 558f1e76bdb64..7bd33a0b22d2d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -734,18 +734,44 @@ public static SessionContextHandle createSessionContext(long readerPtr, long run } /** - * Executes a Substrait plan against the configured SessionContext. - * Consumes the session context handle (freed internally when stream closes). + * Frees a native {@code SessionContext} handle. Invoked from + * {@link SessionContextHandle#doCloseNative()} ()} on error / never-executed paths; not called on the + * happy path where Rust's {@code execute_with_context} consumes the handle itself. + * Safe to call at most once per pointer. */ - /** Frees a native SessionContext handle. Safe to call once. */ public static void closeSessionContext(long ptr) { NativeCall.invokeVoid(CLOSE_SESSION_CONTEXT, ptr); } - public static void executeWithContextAsync(long sessionCtxPtr, byte[] substraitPlan, ActionListener listener) { - NativeHandle.validatePointer(sessionCtxPtr, "sessionContext"); + /** + * Executes a Substrait plan against the configured SessionContext. + * + *

Rust's {@code execute_with_context} takes ownership of the {@code SessionContext} via + * {@code Box::from_raw} on entry, regardless of whether the rest of the call then succeeds or + * returns an error. The handle is therefore marked consumed in a {@code finally} block so + * that both success and native-error paths skip {@code df_close_session_context} (which + * would otherwise double-free). Only a Java-side failure before the downcall dispatches + * (argument marshalling) leaves the handle unconsumed, in which case its + * {@link SessionContextHandle#doCloseNative()} ()} will free it. + */ + public static void executeWithContextAsync(SessionContextHandle sessionContext, byte[] substraitPlan, ActionListener listener) { + final long sessionCtxPtr; + try { + sessionCtxPtr = sessionContext.getPointer(); + } catch (Exception e) { + listener.onFailure(e); + return; + } try (var call = new NativeCall()) { - long result = call.invoke(EXECUTE_WITH_CONTEXT, sessionCtxPtr, call.bytes(substraitPlan), (long) substraitPlan.length); + var plan = call.bytes(substraitPlan); + long planLen = (long) substraitPlan.length; + long result; + try { + result = call.invoke(EXECUTE_WITH_CONTEXT, sessionCtxPtr, plan, planLen); + } finally { + // Rust took ownership via Box::from_raw; do not let doClose() double-free. + sessionContext.markConsumed(); + } listener.onResponse(result); } catch (Throwable throwable) { listener.onFailure(throwable instanceof Exception ? (Exception) throwable : new RuntimeException(throwable)); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/SessionContextHandle.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/SessionContextHandle.java index c26d1799611eb..08d8ae515e45a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/SessionContextHandle.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/SessionContextHandle.java @@ -8,31 +8,35 @@ package org.opensearch.be.datafusion.nativelib; -import org.opensearch.analytics.backend.jni.NativeHandle; +import org.opensearch.analytics.backend.jni.ConsumableNativeHandle; /** - * Type-safe wrapper for a native SessionContext pointer returned by - * {@link NativeBridge#createSessionContext}. The Rust side consumes this - * handle when {@link NativeBridge#executeWithContextAsync} is called, - * so {@link #doClose()} is a no-op — the pointer is freed by Rust internally. + * Type-safe wrapper for a native {@code SessionContext} pointer returned by + * {@link NativeBridge#createSessionContext}. * - *

This handle exists to participate in the {@link NativeHandle} live-pointer - * registry so that {@link NativeHandle#validatePointer} passes for FFM calls. + *

Ownership

+ *

On the happy path, {@link NativeBridge#executeWithContextAsync} transfers ownership of the + * pointer to Rust, which takes it via {@code Box::from_raw} on the first line of + * {@code df_execute_with_context} and drops it when the stream finishes. The bridge method + * calls {@link ConsumableNativeHandle#markConsumed()} after the FFM downcall so that the + * inherited {@link #doClose()} short-circuits without calling + * {@code df_close_session_context} — doing so would be a double-free. + * + *

On any path where execute is never reached (Java-side error before the downcall, aborted + * search, context closed before execution), {@link #doCloseNative()} calls + * {@link NativeBridge#closeSessionContext(long)} which invokes the Rust + * {@code df_close_session_context} entry to free the handle. Both the explicit + * {@link #close()} call from {@link org.opensearch.be.datafusion.DatafusionContext#close()} and + * the {@link java.lang.ref.Cleaner} GC-time fallback route through this path. */ -public class SessionContextHandle extends NativeHandle { +public class SessionContextHandle extends ConsumableNativeHandle { public SessionContextHandle(long ptr) { super(ptr); } @Override - protected void doClose() { - // TODO: Handle error-path cleanup. Currently Rust consumes the handle in - // execute_with_context (moves QueryTrackingContext into the stream). If execute - // fails or is never called, this handle leaks on the Rust side. - // Options: (a) AtomicBool 'consumed' flag on Rust handle — close_session_context - // checks flag before freeing, (b) don't consume in Rust and use no-op tracking - // on the stream, (c) markConsumed() on NativeHandle to skip doClose on happy path. - // See df_close_session_context FFM entry which exists but is not yet wired here. + protected void doCloseNative() { + NativeBridge.closeSessionContext(ptr); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java index 0e5c78087986f..23fa446681911 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java @@ -8,6 +8,7 @@ package org.opensearch.be.datafusion; +import org.opensearch.analytics.backend.jni.NativeHandle; import org.opensearch.be.datafusion.nativelib.NativeBridge; import org.opensearch.be.datafusion.nativelib.ReaderHandle; import org.opensearch.be.datafusion.nativelib.SessionContextHandle; @@ -95,8 +96,13 @@ public void testSessionContextCreationAndTableRegistration() throws Exception { "SELECT message FROM test_table", runtimeHandle.get() ); + // Capture the pointer value BEFORE execute — after execute the handle is marked consumed + // (which closes the Java wrapper), so getPointer() would throw IllegalStateException. + long sessionCtxPtrBefore = sessionCtx.getPointer(); + assertTrue("SessionContext pointer should be live before execute", NativeHandle.isLivePointer(sessionCtxPtrBefore)); + CompletableFuture future = new CompletableFuture<>(); - NativeBridge.executeWithContextAsync(sessionCtx.getPointer(), substrait, new ActionListener<>() { + NativeBridge.executeWithContextAsync(sessionCtx, substrait, new ActionListener<>() { @Override public void onResponse(Long streamPtr) { future.complete(streamPtr); @@ -110,8 +116,10 @@ public void onFailure(Exception exception) { long streamPtr = future.join(); assertTrue("Stream pointer should be non-zero", streamPtr != 0); - // Session context is consumed by execute — close the Java handle - sessionCtx.close(); + // executeWithContextAsync marks the handle consumed (which closes the Java wrapper). + // Verify the pointer is no longer in the live registry and the wrapper rejects getPointer(). + assertFalse("SessionContextHandle pointer must no longer be live after execute", NativeHandle.isLivePointer(sessionCtxPtrBefore)); + expectThrows(IllegalStateException.class, sessionCtx::getPointer); NativeBridge.streamClose(streamPtr); readerHandle.close(); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java index 804635187cf4d..0b2e0e28dbb6a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java @@ -113,12 +113,12 @@ private FragmentResources startFragment(FragmentExecutionRequest request, Resolv GatedCloseable gatedReader = resolved.readerProvider.acquireReader(); SearchExecEngine engine = null; EngineResultStream stream = null; + BackendExecutionContext backendContext = null; try { ShardScanExecutionContext ctx = buildContext(request, gatedReader.get(), resolved.plan, task); AnalyticsSearchBackendPlugin backend = backends.get(resolved.plan.getBackendId()); // Apply instruction handlers in order — each builds upon the previous handler's backend context - BackendExecutionContext backendContext = null; List instructions = resolved.plan.getInstructions(); if (!instructions.isEmpty()) { FragmentInstructionHandlerFactory factory = backend.getInstructionHandlerFactory(); @@ -137,6 +137,16 @@ private FragmentResources startFragment(FragmentExecutionRequest request, Resolv } catch (Exception suppressed) { e.addSuppressed(suppressed); } + // Close the backend execution context as a safety net for failure paths that + // never reached / never finished the engine construction — if the handle was + // already transferred, close() is a no-op (implementations must be idempotent). + if (backendContext != null) { + try { + backendContext.close(); + } catch (Exception suppressed) { + e.addSuppressed(suppressed); + } + } throw e; } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java index d9205260000ef..c934bc1c6b76b 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java @@ -61,9 +61,40 @@ public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryConte FragmentInstructionHandlerFactory factory = stage.getInstructionHandlerFactory(); if (factory != null) { BackendExecutionContext backendContext = null; - for (InstructionNode node : stage.getPlanAlternatives().getFirst().instructions()) { - FragmentInstructionHandler handler = factory.createHandler(node); - backendContext = handler.apply(node, context, backendContext); + Throwable primaryFailure = null; + try { + for (InstructionNode node : stage.getPlanAlternatives().getFirst().instructions()) { + FragmentInstructionHandler handler = factory.createHandler(node); + BackendExecutionContext previous = backendContext; + backendContext = handler.apply(node, context, backendContext); + // A handler that returns a new reference implicitly abandons the previous + // context — close it now so its resources aren't orphaned. + if (previous != null && previous != backendContext) { + previous.close(); + } + } + } catch (Throwable t) { + primaryFailure = t; + } finally { + // The reduce path does not currently hand backendContext off to the sink + // provider — any resources attached by instruction handlers must be released + // here. Close is idempotent so a future handoff can coexist with this call. + if (backendContext != null) { + try { + backendContext.close(); + } catch (Exception closeFailure) { + if (primaryFailure != null) { + primaryFailure.addSuppressed(closeFailure); + } else { + primaryFailure = closeFailure; + } + } + } + } + if (primaryFailure != null) { + if (primaryFailure instanceof RuntimeException re) throw re; + if (primaryFailure instanceof Error err) throw err; + throw new RuntimeException("Instruction handler failed for stageId=" + stage.getStageId(), primaryFailure); } } From e011d3c3c9b5d67d5a81738e3d643a99547529e3 Mon Sep 17 00:00:00 2001 From: bharath-techie Date: Thu, 7 May 2026 11:16:50 +0530 Subject: [PATCH 2/2] fixing sandbox check Signed-off-by: bharath-techie --- .../backend/jni/ConsumableNativeHandle.java | 4 +- .../jni/ConsumableNativeHandleTests.java | 142 ++++++++++++++++++ 2 files changed, 144 insertions(+), 2 deletions(-) create mode 100644 sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandleTests.java diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandle.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandle.java index 2618ac6735d18..033c5487b85a8 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandle.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandle.java @@ -22,8 +22,8 @@ * This: *

diff --git a/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandleTests.java b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandleTests.java new file mode 100644 index 0000000000000..a2fd03d7901bc --- /dev/null +++ b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/backend/jni/ConsumableNativeHandleTests.java @@ -0,0 +1,142 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.backend.jni; + +import org.opensearch.test.OpenSearchTestCase; + +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Tests for {@link ConsumableNativeHandle}'s ownership-transfer contract. + * + *

The class guards against two specific failure modes: + *

+ * + *

Both paths rely on the {@code doCloseNative()} callback being invoked + * exactly zero or one times, never twice. These tests nail that contract + * down with a counting subclass so a future change to + * {@link ConsumableNativeHandle} that accidentally re-introduces a + * double-close will fail loudly. + * + *

Reference: the real subclass + * {@code org.opensearch.be.datafusion.nativelib.SessionContextHandle} is used + * from {@code DatafusionContext#close()} and + * {@code DataFusionSessionState#close()} — both paths can reach + * {@code close()} on the same instance, so idempotency is load-bearing. + */ +public class ConsumableNativeHandleTests extends OpenSearchTestCase { + + /** + * Counts calls to {@link #doCloseNative()} so tests can assert exact + * invocation counts. + */ + private static final class CountingHandle extends ConsumableNativeHandle { + final AtomicInteger nativeCloses = new AtomicInteger(0); + + CountingHandle(long ptr) { + super(ptr); + } + + @Override + protected void doCloseNative() { + nativeCloses.incrementAndGet(); + } + } + + // ---- close() without consumption ------------------------------------ + + public void testCloseWithoutConsumeCallsNativeOnce() { + CountingHandle handle = new CountingHandle(100L); + handle.close(); + assertEquals("doCloseNative should run once on the never-consumed path", 1, handle.nativeCloses.get()); + } + + public void testDoubleCloseWithoutConsumeStillCallsNativeOnce() { + CountingHandle handle = new CountingHandle(101L); + handle.close(); + handle.close(); + assertEquals("close() must be idempotent — second call is a no-op", 1, handle.nativeCloses.get()); + } + + // ---- markConsumed() ownership-transferred path ---------------------- + + public void testMarkConsumedSkipsNativeClose() { + CountingHandle handle = new CountingHandle(200L); + handle.markConsumed(); + assertEquals( + "markConsumed() must not call doCloseNative — the native side already freed the pointer", + 0, + handle.nativeCloses.get() + ); + } + + public void testCloseAfterMarkConsumedIsNoOp() { + CountingHandle handle = new CountingHandle(201L); + handle.markConsumed(); + handle.close(); + assertEquals( + "An explicit close() after markConsumed() must remain a no-op — otherwise Rust's Box::from_raw would be followed by a second free", + 0, + handle.nativeCloses.get() + ); + } + + public void testMarkConsumedAfterCloseDoesNotRunNativeTwice() { + // Order reversed from the normal happy path. The bridge always calls + // markConsumed() after the FFM downcall returns, but the test ensures + // that even if some future caller inverted the sequence, the native + // close is never invoked twice. + CountingHandle handle = new CountingHandle(202L); + handle.close(); + assertEquals(1, handle.nativeCloses.get()); + handle.markConsumed(); + assertEquals("markConsumed() after close() must not trigger another native close", 1, handle.nativeCloses.get()); + } + + public void testMarkConsumedIsIdempotent() { + CountingHandle handle = new CountingHandle(203L); + handle.markConsumed(); + handle.markConsumed(); + handle.close(); + assertEquals(0, handle.nativeCloses.get()); + } + + // ---- State observation --------------------------------------------- + + public void testGetPointerAfterMarkConsumedThrows() { + CountingHandle handle = new CountingHandle(300L); + handle.markConsumed(); + // markConsumed() closes the Java wrapper eagerly; subsequent getPointer + // should refuse to hand out the now-dangling value. + expectThrows(IllegalStateException.class, handle::getPointer); + } + + public void testIsLivePointerFalseAfterMarkConsumed() { + CountingHandle handle = new CountingHandle(301L); + assertTrue(NativeHandle.isLivePointer(301L)); + handle.markConsumed(); + assertFalse( + "markConsumed() must remove the pointer from the live registry so validatePointer rejects it on a stale re-use", + NativeHandle.isLivePointer(301L) + ); + } + + public void testValidatePointerAfterMarkConsumedThrows() { + CountingHandle handle = new CountingHandle(302L); + handle.markConsumed(); + expectThrows(IllegalStateException.class, () -> NativeHandle.validatePointer(302L, "consumed")); + } +}