From 0ee193adf8df1cde292285af5e1322540a53cbaf Mon Sep 17 00:00:00 2001 From: Solaris-star <820622658@qq.com> Date: Thu, 30 Jul 2026 16:51:34 +0800 Subject: [PATCH] fix(rest): release fetched FileScanTask state to prevent unbounded memory growth The reference REST server-side scan planning implementation retains all planned FileScanTask objects in the singleton InMemoryPlanningState even after clients have successfully fetched every plan task. Neither fileScanTasksForPlanTask nor nextPlanTask removes the fetched entries, and the state is only released via cancelPlan. Per the REST Catalog OpenAPI spec, cancellation is not required after all plan tasks have been fetched, so a successful fetch lifecycle must release state without an explicit cancel request (#17427). Add releasePlanTask() to remove a single fetched plan task's FileScanTask list and next-task link, and releaseAsyncPlanForTask() to remove the async planning state when the last plan task in a chain is fetched. Call both from CatalogHandlers.fetchScanTasks after building the response. Adds TestInMemoryPlanningState with unit tests covering: - releasePlanTask removes both file scan tasks and next-task link - releaseAsyncPlanForTask removes async planning state - malformed keys are handled gracefully - releasing unknown keys is a no-op Note: could not run the full Gradle test suite locally (partial clone conflicts with the multi-module build), but the changes are minimal Map.remove() calls following existing patterns in the codebase. --- .../apache/iceberg/rest/CatalogHandlers.java | 11 ++- .../iceberg/rest/InMemoryPlanningState.java | 35 +++++++ .../rest/TestInMemoryPlanningState.java | 94 +++++++++++++++++++ 3 files changed, 139 insertions(+), 1 deletion(-) create mode 100644 core/src/test/java/org/apache/iceberg/rest/TestInMemoryPlanningState.java 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(); + } +}