Changes to handle session context handle close

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.
+ *
+ * <p>The bridge method that performs the consuming FFM call must invoke
+ * {@link #markConsumed()} after the downcall returns (typically in a {@code finally} block).
+ * This:
+ * <ul>
+ *   <li>flips an internal flag so the inherited {@link #doClose()} short-circuits;</li>
+ *   <li>eagerly closes the Java wrapper — the pointer is removed from LIVE_HANDLES in
+ *       {@link NativeHandle}, subsequent {@link #getPointer()} calls
+ *       throw, and {@link NativeHandle#validatePointer(long, String) validatePointer} rejects
+ *       the now-dangling pointer value.</li>
+ * </ul>
+ *
+ * <p>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.
+ *
+ * <p>{@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 <b>not</b> 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();
+}

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.
  *
  * <p>Each backend defines its own concrete implementation (e.g.,
  * {@code DataFusionSessionState} holding a native SessionContext handle).
  *
+ * <h2>Lifecycle</h2>
+ * <p>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.
+    }
+}

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.
+ *
+ * <p>The class guards against two specific failure modes:
+ * <ul>
+ *   <li><b>Double-free</b>: the Rust side consumed the pointer via
+ *       {@code Box::from_raw}, then the Java-side {@code close()} calls
+ *       {@code df_close_X} which tries to free the same memory again.</li>
+ *   <li><b>Leak</b>: the consuming FFM call never dispatched (pre-invoke
+ *       Java failure, aborted flow), so the Java wrapper is responsible for
+ *       calling {@code df_close_X} exactly once.</li>
+ * </ul>
+ *
+ * <p>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.
+ *
+ * <p>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"));
+    }
+}

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,
         ))

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<i64, DataFusionError> {
-    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))
     })?;

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.
+ *
+ * <p>{@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();
+        }
+    }
 }

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();
+            }
         }
     }
 

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<Long> 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));
     }