diff --git a/core/src/main/java/org/apache/iceberg/rest/CatalogHandlers.java b/core/src/main/java/org/apache/iceberg/rest/CatalogHandlers.java index 975cb960096b..e4e4ca324a3c 100644 --- a/core/src/main/java/org/apache/iceberg/rest/CatalogHandlers.java +++ b/core/src/main/java/org/apache/iceberg/rest/CatalogHandlers.java @@ -928,10 +928,19 @@ public static FetchScanTasksResponse fetchScanTasks( Table table = catalog.loadTable(ident); String planTask = request.planTask(); List fileScanTasks = IN_MEMORY_PLANNING_STATE.fileScanTasksForPlanTask(planTask); + List nextPlanTasks = IN_MEMORY_PLANNING_STATE.nextPlanTask(planTask); + + // Release the fetched plan task's state to prevent unbounded memory growth. The FileScanTask + // list is the dominant memory consumer and need not be retained once served (#17427). When + // this was the last plan task in the chain, also release the async planning state. + IN_MEMORY_PLANNING_STATE.releasePlanTask(planTask); + if (nextPlanTasks.isEmpty()) { + IN_MEMORY_PLANNING_STATE.releaseAsyncPlanForTask(planTask); + } return FetchScanTasksResponse.builder() .withFileScanTasks(fileScanTasks) - .withPlanTasks(IN_MEMORY_PLANNING_STATE.nextPlanTask(planTask)) + .withPlanTasks(nextPlanTasks) .withSpecsById(table.specs()) .build(); } diff --git a/core/src/main/java/org/apache/iceberg/rest/InMemoryPlanningState.java b/core/src/main/java/org/apache/iceberg/rest/InMemoryPlanningState.java index b90740c4faba..7af479c06996 100644 --- a/core/src/main/java/org/apache/iceberg/rest/InMemoryPlanningState.java +++ b/core/src/main/java/org/apache/iceberg/rest/InMemoryPlanningState.java @@ -127,6 +127,41 @@ List nextPlanTask(String planTaskKey) { return ImmutableList.of(); } + /** + * Releases the state for a single fetched plan task. Called after a successful fetch to prevent + * unbounded memory growth — the {@code FileScanTask} list is the dominant memory consumer and + * need not be retained once the client has received it (#17427). + * + * @param planTaskKey the plan task key to release + */ + void releasePlanTask(String planTaskKey) { + planTaskToFileScanTasks.remove(planTaskKey); + planTaskToNext.remove(planTaskKey); + } + + /** + * Releases the async planning state for the plan that owns the given plan task key. Called when + * the last plan task in a chain is fetched, so completed plans don't accumulate indefinitely + * (#17427). + * + *

The plan task key format is {@code {planId}-{tableId}-{sequence}}. The planId is extracted by + * stripping the last two hyphen-separated components. + * + * @param planTaskKey a plan task key belonging to the plan whose async state should be released + */ + void releaseAsyncPlanForTask(String planTaskKey) { + int lastHyphen = planTaskKey.lastIndexOf('-'); + if (lastHyphen < 0) { + return; + } + int secondLastHyphen = planTaskKey.lastIndexOf('-', lastHyphen - 1); + if (secondLastHyphen < 0) { + return; + } + String planId = planTaskKey.substring(0, secondLastHyphen); + asyncPlanningStates.remove(planId); + } + /** * Retrieves the initial set of file scan tasks for a plan. PlanIDs are assumed to be separated * with hyphens where the last component indicates the sequencing of plan IDs. diff --git a/core/src/test/java/org/apache/iceberg/rest/TestInMemoryPlanningState.java b/core/src/test/java/org/apache/iceberg/rest/TestInMemoryPlanningState.java new file mode 100644 index 000000000000..8f062d62e34d --- /dev/null +++ b/core/src/test/java/org/apache/iceberg/rest/TestInMemoryPlanningState.java @@ -0,0 +1,94 @@ +/* + * 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.iceberg.rest; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.util.List; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.exceptions.NoSuchPlanIdException; +import org.apache.iceberg.exceptions.NoSuchPlanTaskException; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +public class TestInMemoryPlanningState { + + private final InMemoryPlanningState state = InMemoryPlanningState.getInstance(); + + @AfterEach + public void cleanup() { + state.clear(); + } + + @Test + public void releasePlanTaskRemovesFetchedState() { + FileScanTask task = Mockito.mock(FileScanTask.class); + String planTaskKey = "plan-1-table-1-0"; + state.addPlanTask(planTaskKey, ImmutableList.of(task)); + state.addNextPlanTask(planTaskKey, "plan-1-table-1-1"); + + // State is readable before release + assertThat(state.fileScanTasksForPlanTask(planTaskKey)).containsExactly(task); + assertThat(state.nextPlanTask(planTaskKey)).containsExactly("plan-1-table-1-1"); + + state.releasePlanTask(planTaskKey); + + // Both the file scan tasks and the next-task link are gone after release + assertThatThrownBy(() -> state.fileScanTasksForPlanTask(planTaskKey)) + .isInstanceOf(NoSuchPlanTaskException.class); + assertThat(state.nextPlanTask(planTaskKey)).isEmpty(); + } + + @Test + public void releaseAsyncPlanForTaskRemovesAsyncState() { + // planId is "async-plan-1"; the key format is {planId}-{tableId}-{sequence} + String planId = "async-plan-1"; + String planTaskKey = planId + "-table-1-0"; + state.addAsyncPlan(planId); + assertThat(state.asyncPlanStatus(planId)).isEqualTo(PlanStatus.SUBMITTED); + + state.releaseAsyncPlanForTask(planTaskKey); + + assertThatThrownBy(() -> state.asyncPlanStatus(planId)) + .isInstanceOf(NoSuchPlanIdException.class); + } + + @Test + public void releaseAsyncPlanForTaskIgnoresMalformedKeys() { + // Keys without two hyphens must not throw and must not remove unrelated state + String planId = "some-plan"; + state.addAsyncPlan(planId); + + state.releaseAsyncPlanForTask("no-hyphens"); + state.releaseAsyncPlanForTask("only-one-hyphen"); + + // Unrelated plan state is untouched + assertThat(state.asyncPlanStatus(planId)).isEqualTo(PlanStatus.SUBMITTED); + } + + @Test + public void releasePlanTaskIsIdempotentForUnknownKeys() { + // Releasing a key that was never added must be a no-op, not an error + state.releasePlanTask("plan-x-table-y-0"); + assertThat(state.nextPlanTask("plan-x-table-y-0")).isEmpty(); + } +}