Skip to content

Commit 958ce45

Browse files
JackieTien97JackieTien97
authored
Fix repeated RPC dispatch reusing a released FragmentInstanceContext (NPE) (#17794)
1 parent 0848379 commit 958ce45

6 files changed

Lines changed: 141 additions & 25 deletions

File tree

‎iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ public enum TSStatusCode {
146146
QUERY_TIMEOUT(720),
147147
PLAN_FAILED_NETWORK_PARTITION(721),
148148
CANNOT_FETCH_FI_STATE(722),
149+
REPEATED_RPC_CALL(723),
149150

150151
// OBJECT
151152
OBJECT_NOT_EXISTS(740),

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/executor/RegionReadExecutor.java‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import org.apache.iotdb.common.rpc.thrift.TSStatus;
2323
import org.apache.iotdb.commons.consensus.ConsensusGroupId;
2424
import org.apache.iotdb.commons.consensus.DataRegionId;
25+
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
2526
import org.apache.iotdb.commons.utils.ErrorHandlingCommonUtils;
2627
import org.apache.iotdb.commons.utils.StatusUtils;
2728
import org.apache.iotdb.commons.utils.TestOnly;
@@ -127,6 +128,13 @@ public RegionExecutionResult execute(
127128
|| t instanceof InterruptedException) {
128129
resp.setReadNeedRetry(true);
129130
resp.setStatus(new TSStatus(TSStatusCode.RATIS_READ_UNAVAILABLE.getStatusCode()));
131+
} else if (t instanceof IoTDBRuntimeException) {
132+
// Carry the original status code (e.g. REPEATED_RPC_CALL) back to the dispatcher so it is
133+
// not downgraded to EXECUTE_STATEMENT_ERROR; needRetryHelper decides retryability.
134+
TSStatus status = new TSStatus(((IoTDBRuntimeException) t).getErrorCode());
135+
status.setMessage(t.getMessage());
136+
resp.setStatus(status);
137+
resp.setReadNeedRetry(StatusUtils.needRetryHelper(status));
130138
}
131139
return resp;
132140
}
@@ -156,8 +164,19 @@ public RegionExecutionResult execute(FragmentInstance fragmentInstance) {
156164
}
157165
} catch (Throwable t) {
158166
LOGGER.warn(DataNodeQueryMessages.EXECUTE_FRAGMENTINSTANCE_IN_QUERYEXECUTOR_FAILED, t);
159-
return RegionExecutionResult.create(
160-
false, String.format(ERROR_MSG_FORMAT, t.getMessage()), null);
167+
RegionExecutionResult resp =
168+
RegionExecutionResult.create(
169+
false, String.format(ERROR_MSG_FORMAT, t.getMessage()), null);
170+
Throwable rootCause = ErrorHandlingCommonUtils.getRootCause(t);
171+
if (rootCause instanceof IoTDBRuntimeException) {
172+
// Carry the original status code (e.g. REPEATED_RPC_CALL) back to the dispatcher so it is
173+
// not downgraded to EXECUTE_STATEMENT_ERROR; needRetryHelper decides retryability.
174+
TSStatus status = new TSStatus(((IoTDBRuntimeException) rootCause).getErrorCode());
175+
status.setMessage(rootCause.getMessage());
176+
resp.setStatus(status);
177+
resp.setReadNeedRetry(StatusUtils.needRetryHelper(status));
178+
}
179+
return resp;
161180
}
162181
}
163182
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java‎

Lines changed: 57 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -147,23 +147,30 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
147147
FragmentInstanceStateMachine stateMachine =
148148
new FragmentInstanceStateMachine(instanceId, instanceNotificationExecutor);
149149

150-
int dataNodeFINum = instance.getDataNodeFINum();
151-
DataNodeQueryContext dataNodeQueryContext =
152-
getOrCreateDataNodeQueryContext(instanceId.getQueryId(), dataNodeFINum);
153-
150+
boolean[] contextCreated = new boolean[] {false};
151+
DataNodeQueryContext[] dataNodeQueryContexts = new DataNodeQueryContext[1];
154152
FragmentInstanceContext context =
155153
instanceContext.computeIfAbsent(
156154
instanceId,
157-
fragmentInstanceId ->
158-
createFragmentInstanceContext(
159-
fragmentInstanceId,
160-
stateMachine,
161-
instance.getSessionInfo(),
162-
dataRegion,
163-
instance.getGlobalTimePredicate(),
164-
dataNodeQueryContextMap,
165-
instance.isDebug(),
166-
instance.isVerbose()));
155+
fragmentInstanceId -> {
156+
contextCreated[0] = true;
157+
// Only ensure the DataNodeQueryContext when we actually create the
158+
// FragmentInstanceContext, so the repeated-dispatch path (which rejects
159+
// without creating a context) does not leak a context entry.
160+
dataNodeQueryContexts[0] =
161+
getOrCreateDataNodeQueryContext(
162+
instanceId.getQueryId(), instance.getDataNodeFINum());
163+
return createFragmentInstanceContext(
164+
fragmentInstanceId,
165+
stateMachine,
166+
instance.getSessionInfo(),
167+
dataRegion,
168+
instance.getGlobalTimePredicate(),
169+
dataNodeQueryContextMap,
170+
instance.isDebug(),
171+
instance.isVerbose());
172+
});
173+
rejectIfRepeatedDispatch(contextCreated[0], instanceId);
167174
context.setHighestPriority(instance.isHighestPriority());
168175

169176
try {
@@ -172,7 +179,7 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
172179
instance.getFragment().getPlanNodeTree(),
173180
instance.getFragment().getTypeProvider(),
174181
context,
175-
dataNodeQueryContext);
182+
dataNodeQueryContexts[0]);
176183

177184
List<IDriver> drivers = new ArrayList<>();
178185
driverFactories.forEach(factory -> drivers.add(factory.createDriver()));
@@ -246,6 +253,30 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
246253
}
247254
}
248255

256+
/**
257+
* If {@code instanceContext.computeIfAbsent} returned an existing {@link FragmentInstanceContext}
258+
* for this {@code instanceId} (i.e. {@code contextCreated} is false), the same FragmentInstance
259+
* has been dispatched before (e.g. an RPC retry in {@code
260+
* FragmentInstanceDispatcherImpl#dispatchRemote}). The previous execution may have already
261+
* released its resources (dataRegion == null), so reusing this cached context would run a fresh
262+
* driver against a released context and trigger an NPE. Reject the duplicated dispatch with
263+
* REPEATED_RPC_CALL instead of reusing it.
264+
*
265+
* <p>This must be called before the planning try block on purpose, so it propagates up
266+
* (RegionReadExecutor carries the status code) without touching the first execution's cached
267+
* resources.
268+
*/
269+
private static void rejectIfRepeatedDispatch(
270+
boolean contextCreated, FragmentInstanceId instanceId) {
271+
if (!contextCreated) {
272+
throw new IoTDBRuntimeException(
273+
String.format(
274+
"Repeated RPC call detected for FragmentInstance %s, reject the duplicated dispatch.",
275+
instanceId.getFullId()),
276+
TSStatusCode.REPEATED_RPC_CALL.getStatusCode());
277+
}
278+
}
279+
249280
private void clearFIRelatedResources(FragmentInstanceId instanceId) {
250281
// close and remove all the handles of the fragment instance
251282
exchangeManager.forceDeregisterFragmentInstance(instanceId.toThrift());
@@ -270,16 +301,20 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance(
270301
FragmentInstanceStateMachine stateMachine =
271302
new FragmentInstanceStateMachine(instanceId, instanceNotificationExecutor);
272303

304+
boolean[] contextCreated = new boolean[] {false};
273305
FragmentInstanceContext context =
274306
instanceContext.computeIfAbsent(
275307
instanceId,
276-
fragmentInstanceId ->
277-
createFragmentInstanceContext(
278-
fragmentInstanceId,
279-
stateMachine,
280-
instance.getSessionInfo(),
281-
instance.isDebug(),
282-
instance.isVerbose()));
308+
fragmentInstanceId -> {
309+
contextCreated[0] = true;
310+
return createFragmentInstanceContext(
311+
fragmentInstanceId,
312+
stateMachine,
313+
instance.getSessionInfo(),
314+
instance.isDebug(),
315+
instance.isVerbose());
316+
});
317+
rejectIfRepeatedDispatch(contextCreated[0], instanceId);
283318
context.setHighestPriority(instance.isHighestPriority());
284319

285320
try {

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import org.apache.iotdb.consensus.exception.RatisReadUnavailableException;
3838
import org.apache.iotdb.db.conf.IoTDBDescriptor;
3939
import org.apache.iotdb.db.exception.mpp.FragmentInstanceDispatchException;
40+
import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException;
4041
import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
4142
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
4243
import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult;
@@ -589,6 +590,20 @@ private void dispatchRemote(FragmentInstance instance, TEndPoint endPoint)
589590
"can't execute request on node {}, error msg is {}, and we try to reconnect this node.",
590591
endPoint,
591592
ExceptionUtils.getRootCause(e).toString());
593+
// If the query has already timed out, do not retry. Re-dispatching the same FragmentInstance
594+
// may cause it to be executed twice on the remote node (see the REPEATED_RPC_CALL handling in
595+
// FragmentInstanceManager), so we fail fast with a timeout status instead.
596+
long currentTime = System.currentTimeMillis();
597+
if (currentTime - queryContext.getStartTime() >= queryContext.getTimeOut()) {
598+
throw new FragmentInstanceDispatchException(
599+
RpcUtils.getStatus(
600+
TSStatusCode.QUERY_TIMEOUT,
601+
String.format(
602+
QueryTimeoutRuntimeException.QUERY_TIMEOUT_EXCEPTION_MESSAGE,
603+
queryContext.getStartTime(),
604+
queryContext.getStartTime() + queryContext.getTimeOut(),
605+
currentTime)));
606+
}
592607
// we just retry once to clear stale connection for a restart node.
593608
try {
594609
dispatchRemoteHelper(instance, endPoint);

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/ErrorHandlingUtils.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ private static TSStatus tryCatchQueryException(Exception e) {
160160

161161
Throwable t = e instanceof ExecutionException ? e.getCause() : e;
162162
if (t instanceof QueryTimeoutRuntimeException) {
163-
return RpcUtils.getStatus(TSStatusCode.INTERNAL_REQUEST_TIME_OUT, rootCause.getMessage());
163+
return RpcUtils.getStatus(TSStatusCode.QUERY_TIMEOUT, rootCause.getMessage());
164164
} else if (t instanceof ParseCancellationException) {
165165
return RpcUtils.getStatus(
166166
TSStatusCode.SQL_PARSE_ERROR, INFO_PARSING_SQL_ERROR + rootCause.getMessage());

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/executor/RegionReadExecutorTest.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import org.apache.iotdb.commons.consensus.ConsensusGroupId;
2323
import org.apache.iotdb.commons.consensus.DataRegionId;
2424
import org.apache.iotdb.commons.consensus.SchemaRegionId;
25+
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
2526
import org.apache.iotdb.consensus.IConsensus;
2627
import org.apache.iotdb.consensus.exception.ConsensusException;
2728
import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
@@ -30,6 +31,7 @@
3031
import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceManager;
3132
import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance;
3233
import org.apache.iotdb.db.storageengine.dataregion.VirtualDataRegion;
34+
import org.apache.iotdb.rpc.TSStatusCode;
3335

3436
import org.junit.Test;
3537
import org.mockito.Mockito;
@@ -161,6 +163,50 @@ public void testFailedExecute() throws ConsensusException {
161163
assertEquals(String.format(ERROR_MSG_FORMAT, "schema-exception"), res.getMessage());
162164
}
163165

166+
@Test
167+
public void testRepeatedRpcCall() throws ConsensusException {
168+
// A repeated RPC dispatch of the same FragmentInstance makes FragmentInstanceManager throw
169+
// IoTDBRuntimeException(REPEATED_RPC_CALL). RegionReadExecutor must carry that status code back
170+
// (instead of dropping it, which would downgrade it to EXECUTE_STATEMENT_ERROR) and mark the
171+
// result as non-retryable.
172+
ConsensusGroupId dataRegionGroupId = new DataRegionId(1);
173+
FragmentInstanceId fragmentInstanceId =
174+
new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 0), "0");
175+
FragmentInstance fragmentInstance = Mockito.mock(FragmentInstance.class);
176+
Mockito.when(fragmentInstance.getId()).thenReturn(fragmentInstanceId);
177+
178+
IConsensus dataRegionConsensus = Mockito.mock(IConsensus.class);
179+
IConsensus schemaRegionConsensus = Mockito.mock(IConsensus.class);
180+
FragmentInstanceManager fragmentInstanceManager = Mockito.mock(FragmentInstanceManager.class);
181+
182+
RegionReadExecutor executor =
183+
new RegionReadExecutor(dataRegionConsensus, schemaRegionConsensus, fragmentInstanceManager);
184+
185+
// consensus read path (covers both data and schema region queries)
186+
Mockito.when(dataRegionConsensus.read(dataRegionGroupId, fragmentInstance))
187+
.thenThrow(
188+
new IoTDBRuntimeException("repeated", TSStatusCode.REPEATED_RPC_CALL.getStatusCode()));
189+
190+
RegionExecutionResult res = executor.execute(dataRegionGroupId, fragmentInstance);
191+
192+
assertFalse(res.isAccepted());
193+
assertEquals(TSStatusCode.REPEATED_RPC_CALL.getStatusCode(), res.getStatus().getCode());
194+
assertFalse(res.isReadNeedRetry());
195+
196+
// VirtualDataRegion path (FI executed directly through FragmentInstanceManager)
197+
Mockito.when(
198+
fragmentInstanceManager.execDataQueryFragmentInstance(
199+
fragmentInstance, VirtualDataRegion.getInstance()))
200+
.thenThrow(
201+
new IoTDBRuntimeException("repeated", TSStatusCode.REPEATED_RPC_CALL.getStatusCode()));
202+
203+
res = executor.execute(fragmentInstance);
204+
205+
assertFalse(res.isAccepted());
206+
assertEquals(TSStatusCode.REPEATED_RPC_CALL.getStatusCode(), res.getStatus().getCode());
207+
assertFalse(res.isReadNeedRetry());
208+
}
209+
164210
@Test
165211
public void testExceptionHappened() throws ConsensusException {
166212

0 commit comments

Comments
 (0)