From 7a577946d217a1bcbf891ed558b6521889be5600 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:44:55 +0800 Subject: [PATCH] Fix retry for schema query internal writes --- .../schema/TableSchemaQueryWriteVisitor.java | 31 +++++-- .../TableSchemaQueryWriteVisitorTest.java | 89 +++++++++++++++++++ 2 files changed, 115 insertions(+), 5 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/schema/TableSchemaQueryWriteVisitorTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/schema/TableSchemaQueryWriteVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/schema/TableSchemaQueryWriteVisitor.java index 29e9b8d46139b..b7e96eeb1a65e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/schema/TableSchemaQueryWriteVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/schema/TableSchemaQueryWriteVisitor.java @@ -21,25 +21,46 @@ import org.apache.iotdb.commons.consensus.ConsensusGroupId; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.utils.RetryUtils; +import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult; import org.apache.iotdb.db.queryengine.execution.executor.RegionWriteExecutor; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.TableDeviceSourceNode; +import org.apache.iotdb.rpc.TSStatusCode; import java.util.Objects; public class TableSchemaQueryWriteVisitor extends AbstractTableSchemaQueryAttributeSecurityVisitor { + private final RegionWriteExecutor regionWriteExecutor; + + public TableSchemaQueryWriteVisitor() { + this(new RegionWriteExecutor()); + } + + @TestOnly + public TableSchemaQueryWriteVisitor(final RegionWriteExecutor regionWriteExecutor) { + this.regionWriteExecutor = regionWriteExecutor; + } + @Override protected RegionExecutionResult visitTableDeviceSourceNode( final TableDeviceSourceNode node, final ConsensusGroupId context) { if (Objects.nonNull(node.getSenderLocation())) { final RegionExecutionResult result = - new RegionWriteExecutor() - .execute( - context, - new TableNodeLocationAddNode(new PlanNodeId(""), node.getSenderLocation())); - return !result.isAccepted() ? result : null; + regionWriteExecutor.execute( + context, new TableNodeLocationAddNode(new PlanNodeId(""), node.getSenderLocation())); + if (!result.isAccepted() + && Objects.nonNull(result.getStatus()) + && RetryUtils.needRetryForWrite(result.getStatus().getCode())) { + // This write is an internal step of a schema query. Convert a retryable write failure to a + // dispatch failure so that the query scheduler retries the query instead of reporting the + // transient consensus error to the client. + result.setReadNeedRetry(true); + result.getStatus().setCode(TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } + return result.isAccepted() ? null : result; } return null; } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/schema/TableSchemaQueryWriteVisitorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/schema/TableSchemaQueryWriteVisitorTest.java new file mode 100644 index 0000000000000..384afcfbdb73c --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/schema/TableSchemaQueryWriteVisitorTest.java @@ -0,0 +1,89 @@ +/* + * 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.iotdb.db.queryengine.plan.relational.planner.node.schema; + +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.consensus.SchemaRegionId; +import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult; +import org.apache.iotdb.db.queryengine.execution.executor.RegionWriteExecutor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.TableDeviceSourceNode; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class TableSchemaQueryWriteVisitorTest { + + private static final SchemaRegionId SCHEMA_REGION_ID = new SchemaRegionId(1); + + @Test + public void testRetryableWriteFailureTriggersQueryRetry() { + final RegionWriteExecutor regionWriteExecutor = mock(RegionWriteExecutor.class); + final TableDeviceSourceNode node = mock(TableDeviceSourceNode.class); + when(node.getSenderLocation()).thenReturn(new TDataNodeLocation()); + + final TSStatus status = + new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode()) + .setMessage("LeaderSteppingDownException"); + final RegionExecutionResult writeResult = + RegionExecutionResult.create(false, status.getMessage(), status); + when(regionWriteExecutor.execute(eq(SCHEMA_REGION_ID), any(TableNodeLocationAddNode.class))) + .thenReturn(writeResult); + + final RegionExecutionResult result = + new TableSchemaQueryWriteVisitor(regionWriteExecutor) + .visitTableDeviceSourceNode(node, SCHEMA_REGION_ID); + + assertSame(writeResult, result); + assertTrue(result.isReadNeedRetry()); + assertEquals(TSStatusCode.DISPATCH_ERROR.getStatusCode(), result.getStatus().getCode()); + assertEquals("LeaderSteppingDownException", result.getStatus().getMessage()); + } + + @Test + public void testNonRetryableWriteFailureIsPreserved() { + final RegionWriteExecutor regionWriteExecutor = mock(RegionWriteExecutor.class); + final TableDeviceSourceNode node = mock(TableDeviceSourceNode.class); + when(node.getSenderLocation()).thenReturn(new TDataNodeLocation()); + + final TSStatus status = new TSStatus(TSStatusCode.SEMANTIC_ERROR.getStatusCode()); + final RegionExecutionResult writeResult = + RegionExecutionResult.create(false, "semantic error", status); + when(regionWriteExecutor.execute(eq(SCHEMA_REGION_ID), any(TableNodeLocationAddNode.class))) + .thenReturn(writeResult); + + final RegionExecutionResult result = + new TableSchemaQueryWriteVisitor(regionWriteExecutor) + .visitTableDeviceSourceNode(node, SCHEMA_REGION_ID); + + assertSame(writeResult, result); + assertFalse(result.isReadNeedRetry()); + assertEquals(TSStatusCode.SEMANTIC_ERROR.getStatusCode(), result.getStatus().getCode()); + } +}