Fix PIT (Point in Time) resource leaks in v2 query engine

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

core/src/main/java/org/opensearch/sql/planner/physical/CursorCloseOperator.java

@@ -9,6 +9,7 @@
 import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.executor.ExecutionEngine;
+import org.opensearch.sql.storage.TableScanOperator;
 
 /**
  * A plan node which blocks issuing a request in {@link #open} and getting results in {@link
@@ -51,4 +52,22 @@ public ExecutionEngine.Schema schema() {
   public void open() {
     // no-op, no search should be invoked.
   }
+
+  /**
+   * Force cleanup of server-side resources. When a cursor is explicitly closed, any underlying
+   * table scan must release its resources (e.g. PIT) unconditionally, even if pagination is not
+   * complete.
+   */
+  @Override
+  public void close() {
+    forceCloseChildren(input);
+  }
+
+  private void forceCloseChildren(PhysicalPlan node) {
+    if (node instanceof TableScanOperator scan) {
+      scan.forceClose();
+    } else {
+      node.getChild().forEach(this::forceCloseChildren);
+    }
+  }
 }

core/src/main/java/org/opensearch/sql/storage/TableScanOperator.java

@@ -26,6 +26,15 @@ public List<PhysicalPlan> getChild() {
     return Collections.emptyList();
   }
 
+  /**
+   * Force cleanup of server-side resources (e.g. PIT, scroll) regardless of pagination state. Used
+   * when the client explicitly closes a cursor mid-pagination. Default implementation delegates to
+   * {@link #close()}.
+   */
+  public void forceClose() {
+    close();
+  }
+
   /**
    * Explain the execution plan.
    *

core/src/test/java/org/opensearch/sql/planner/physical/CursorCloseOperatorTest.java

@@ -15,6 +15,7 @@
 import org.junit.jupiter.api.DisplayNameGeneration;
 import org.junit.jupiter.api.DisplayNameGenerator;
 import org.junit.jupiter.api.Test;
+import org.opensearch.sql.storage.TableScanOperator;
 
 @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class)
 public class CursorCloseOperatorTest {
@@ -36,11 +37,23 @@ public void open_is_not_propagated() {
   }
 
   @Test
-  public void close_is_propagated() {
-    var child = mock(PhysicalPlan.class);
+  public void close_calls_forceClose_on_table_scan() {
+    var child = mock(TableScanOperator.class);
     var plan = new CursorCloseOperator(child);
     plan.close();
-    verify(child).close();
+    verify(child).forceClose();
+    verify(child, never()).close();
+  }
+
+  @Test
+  public void close_traverses_tree_to_find_table_scan() {
+    var scan = mock(TableScanOperator.class);
+    // Wrap the scan in a regular PhysicalPlan node
+    var middle = mock(PhysicalPlan.class);
+    org.mockito.Mockito.when(middle.getChild()).thenReturn(java.util.List.of(scan));
+    var plan = new CursorCloseOperator(middle);
+    plan.close();
+    verify(scan).forceClose();
   }
 
   @Test

opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java

@@ -145,6 +145,8 @@ public ExplainResponseNode visitTableScan(
             listener.onResponse(openSearchExplain.apply(plan));
           } catch (Exception e) {
             listener.onFailure(e);
+          } finally {
+            plan.close();
           }
         });
   }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java

@@ -108,7 +108,26 @@ private void fetchNextBatch() {
   public void close() {
     super.close();
 
-    client.cleanup(request);
+    if (request.hasAnotherBatch()) {
+      // PIT is still needed for the next page — only clean up the in-memory state
+      // without deleting the PIT from OpenSearch.
+      client.cleanup(request);
+    } else {
+      // No more pages expected (last page consumed, non-paginated query, or error path).
+      // Force delete the PIT to prevent leaking.
+      client.forceCleanup(request);
+    }
+  }
+
+  /**
+   * Force cleanup of server-side resources (PIT) regardless of pagination state. Used by {@link
+   * org.opensearch.sql.planner.physical.CursorCloseOperator} when the client explicitly closes a
+   * cursor mid-pagination.
+   */
+  @Override
+  public void forceClose() {
+    super.close();
+    client.forceCleanup(request);
   }
 
   @Override

opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java

@@ -76,7 +76,7 @@ void query_empty_result() {
       indexScan.open();
       assertFalse(indexScan.hasNext());
     }
-    verify(client).cleanup(any());
+    verify(client).forceCleanup(any());
   }
 
   @Test

opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java

@@ -174,7 +174,7 @@ void query_empty_result() {
       indexScan.open();
       assertFalse(indexScan.hasNext());
     }
-    verify(client).cleanup(any());
+    verify(client).forceCleanup(any());
   }
 
   @Test
@@ -200,7 +200,7 @@ void query_all_results_with_query() {
           () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()),
           () -> assertFalse(indexScan.hasNext()));
     }
-    verify(client).cleanup(any());
+    verify(client).forceCleanup(any());
   }
 
   static final OpenSearchRequest.IndexName EMPLOYEES_INDEX =
@@ -228,7 +228,7 @@ void query_all_results_with_scroll() {
           () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()),
           () -> assertFalse(indexScan.hasNext()));
     }
-    verify(client).cleanup(any());
+    verify(client).forceCleanup(any());
   }
 
   @Test
@@ -255,7 +255,7 @@ void query_some_results_with_query() {
           () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()),
           () -> assertFalse(indexScan.hasNext()));
     }
-    verify(client).cleanup(any());
+    verify(client).forceCleanup(any());
   }
 
   @Test
@@ -278,7 +278,7 @@ void query_some_results_with_scroll() {
           () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()),
           () -> assertFalse(indexScan.hasNext()));
     }
-    verify(client).cleanup(any());
+    verify(client).forceCleanup(any());
   }
 
   static void mockTwoPageResponse(OpenSearchClient client) {
@@ -311,7 +311,7 @@ void query_results_limited_by_query_size() {
           () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()),
           () -> assertFalse(indexScan.hasNext()));
     }
-    verify(client).cleanup(any());
+    verify(client).forceCleanup(any());
   }
 
   @Test