Skip to content

Commit e86ef39

Browse files
JackieTien97JackieTien97
authored
Fix repeated RPC dispatch reusing a released FragmentInstanceContext (NPE) (#17794) (#18235)
1 parent 8b10792 commit e86ef39

6 files changed

Lines changed: 141 additions & 23 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
@@ -126,6 +126,7 @@ public enum TSStatusCode {
126126
QUERY_TIMEOUT(720),
127127
PLAN_FAILED_NETWORK_PARTITION(721),
128128
CANNOT_FETCH_FI_STATE(722),
129+
REPEATED_RPC_CALL(723),
129130

130131
// Arithmetic
131132
NUMERIC_VALUE_OUT_OF_RANGE(750),

‎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;
@@ -124,6 +125,13 @@ public RegionExecutionResult execute(
124125
|| t instanceof InterruptedException) {
125126
resp.setReadNeedRetry(true);
126127
resp.setStatus(new TSStatus(TSStatusCode.RATIS_READ_UNAVAILABLE.getStatusCode()));
128+
} else if (t instanceof IoTDBRuntimeException) {
129+
// Carry the original status code (e.g. REPEATED_RPC_CALL) back to the dispatcher so it is
130+
// not downgraded to EXECUTE_STATEMENT_ERROR; needRetryHelper decides retryability.
131+
TSStatus status = new TSStatus(((IoTDBRuntimeException) t).getErrorCode());
132+
status.setMessage(t.getMessage());
133+
resp.setStatus(status);
134+
resp.setReadNeedRetry(StatusUtils.needRetryHelper(status));
127135
}
128136
return resp;
129137
}
@@ -140,8 +148,19 @@ public RegionExecutionResult execute(FragmentInstance fragmentInstance) {
140148
return RegionExecutionResult.create(!info.getState().isFailed(), info.getMessage(), null);
141149
} catch (Throwable t) {
142150
LOGGER.error("Execute FragmentInstance in QueryExecutor failed.", t);
143-
return RegionExecutionResult.create(
144-
false, String.format(ERROR_MSG_FORMAT, t.getMessage()), null);
151+
RegionExecutionResult resp =
152+
RegionExecutionResult.create(
153+
false, String.format(ERROR_MSG_FORMAT, t.getMessage()), null);
154+
Throwable rootCause = ErrorHandlingCommonUtils.getRootCause(t);
155+
if (rootCause instanceof IoTDBRuntimeException) {
156+
// Carry the original status code (e.g. REPEATED_RPC_CALL) back to the dispatcher so it is
157+
// not downgraded to EXECUTE_STATEMENT_ERROR; needRetryHelper decides retryability.
158+
TSStatus status = new TSStatus(((IoTDBRuntimeException) rootCause).getErrorCode());
159+
status.setMessage(rootCause.getMessage());
160+
resp.setStatus(status);
161+
resp.setReadNeedRetry(StatusUtils.needRetryHelper(status));
162+
}
163+
return resp;
145164
}
146165
}
147166
}

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

Lines changed: 57 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.apache.iotdb.commons.concurrent.threadpool.ScheduledExecutorUtil;
2626
import org.apache.iotdb.commons.conf.CommonDescriptor;
2727
import org.apache.iotdb.commons.exception.IoTDBException;
28+
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
2829
import org.apache.iotdb.db.conf.IoTDBDescriptor;
2930
import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException;
3031
import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
@@ -45,6 +46,7 @@
4546
import org.apache.iotdb.db.storageengine.dataregion.IDataRegionForQuery;
4647
import org.apache.iotdb.db.utils.SetThreadName;
4748
import org.apache.iotdb.mpp.rpc.thrift.TFetchFragmentInstanceStatisticsResp;
49+
import org.apache.iotdb.rpc.TSStatusCode;
4850

4951
import io.airlift.units.Duration;
5052
import org.slf4j.Logger;
@@ -143,22 +145,29 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
143145
FragmentInstanceStateMachine stateMachine =
144146
new FragmentInstanceStateMachine(instanceId, instanceNotificationExecutor);
145147

146-
int dataNodeFINum = instance.getDataNodeFINum();
147-
DataNodeQueryContext dataNodeQueryContext =
148-
getOrCreateDataNodeQueryContext(instanceId.getQueryId(), dataNodeFINum);
149-
148+
boolean[] contextCreated = new boolean[] {false};
149+
DataNodeQueryContext[] dataNodeQueryContexts = new DataNodeQueryContext[1];
150150
FragmentInstanceContext context =
151151
instanceContext.computeIfAbsent(
152152
instanceId,
153-
fragmentInstanceId ->
154-
createFragmentInstanceContext(
155-
fragmentInstanceId,
156-
stateMachine,
157-
instance.getSessionInfo(),
158-
dataRegion,
159-
instance.getGlobalTimePredicate(),
160-
dataNodeQueryContextMap,
161-
instance.isDebug()));
153+
fragmentInstanceId -> {
154+
contextCreated[0] = true;
155+
// Only ensure the DataNodeQueryContext when we actually create the
156+
// FragmentInstanceContext, so the repeated-dispatch path (which rejects
157+
// without creating a context) does not leak a context entry.
158+
dataNodeQueryContexts[0] =
159+
getOrCreateDataNodeQueryContext(
160+
instanceId.getQueryId(), instance.getDataNodeFINum());
161+
return createFragmentInstanceContext(
162+
fragmentInstanceId,
163+
stateMachine,
164+
instance.getSessionInfo(),
165+
dataRegion,
166+
instance.getGlobalTimePredicate(),
167+
dataNodeQueryContextMap,
168+
instance.isDebug());
169+
});
170+
rejectIfRepeatedDispatch(contextCreated[0], instanceId);
162171
context.setHighestPriority(instance.isHighestPriority());
163172

164173
try {
@@ -167,7 +176,7 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
167176
instance.getFragment().getPlanNodeTree(),
168177
instance.getFragment().getTypeProvider(),
169178
context,
170-
dataNodeQueryContext);
179+
dataNodeQueryContexts[0]);
171180

172181
List<IDriver> drivers = new ArrayList<>();
173182
driverFactories.forEach(factory -> drivers.add(factory.createDriver()));
@@ -231,6 +240,30 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
231240
}
232241
}
233242

243+
/**
244+
* If {@code instanceContext.computeIfAbsent} returned an existing {@link FragmentInstanceContext}
245+
* for this {@code instanceId} (i.e. {@code contextCreated} is false), the same FragmentInstance
246+
* has been dispatched before (e.g. an RPC retry in {@code
247+
* FragmentInstanceDispatcherImpl#dispatchRemote}). The previous execution may have already
248+
* released its resources (dataRegion == null), so reusing this cached context would run a fresh
249+
* driver against a released context and trigger an NPE. Reject the duplicated dispatch with
250+
* REPEATED_RPC_CALL instead of reusing it.
251+
*
252+
* <p>This must be called before the planning try block on purpose, so it propagates up
253+
* (RegionReadExecutor carries the status code) without touching the first execution's cached
254+
* resources.
255+
*/
256+
private static void rejectIfRepeatedDispatch(
257+
boolean contextCreated, FragmentInstanceId instanceId) {
258+
if (!contextCreated) {
259+
throw new IoTDBRuntimeException(
260+
String.format(
261+
"Repeated RPC call detected for FragmentInstance %s, reject the duplicated dispatch.",
262+
instanceId.getFullId()),
263+
TSStatusCode.REPEATED_RPC_CALL.getStatusCode());
264+
}
265+
}
266+
234267
private void clearFIRelatedResources(FragmentInstanceId instanceId) {
235268
// close and remove all the handles of the fragment instance
236269
exchangeManager.forceDeregisterFragmentInstance(instanceId.toThrift());
@@ -255,15 +288,19 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance(
255288
FragmentInstanceStateMachine stateMachine =
256289
new FragmentInstanceStateMachine(instanceId, instanceNotificationExecutor);
257290

291+
boolean[] contextCreated = new boolean[] {false};
258292
FragmentInstanceContext context =
259293
instanceContext.computeIfAbsent(
260294
instanceId,
261-
fragmentInstanceId ->
262-
createFragmentInstanceContext(
263-
fragmentInstanceId,
264-
stateMachine,
265-
instance.getSessionInfo(),
266-
instance.isDebug()));
295+
fragmentInstanceId -> {
296+
contextCreated[0] = true;
297+
return createFragmentInstanceContext(
298+
fragmentInstanceId,
299+
stateMachine,
300+
instance.getSessionInfo(),
301+
instance.isDebug());
302+
});
303+
rejectIfRepeatedDispatch(contextCreated[0], instanceId);
267304
context.setHighestPriority(instance.isHighestPriority());
268305

269306
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
@@ -34,6 +34,7 @@
3434
import org.apache.iotdb.consensus.exception.RatisReadUnavailableException;
3535
import org.apache.iotdb.db.conf.IoTDBDescriptor;
3636
import org.apache.iotdb.db.exception.mpp.FragmentInstanceDispatchException;
37+
import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException;
3738
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
3839
import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult;
3940
import org.apache.iotdb.db.queryengine.execution.executor.RegionReadExecutor;
@@ -550,6 +551,20 @@ private void dispatchRemote(FragmentInstance instance, TEndPoint endPoint)
550551
"can't execute request on node {}, error msg is {}, and we try to reconnect this node.",
551552
endPoint,
552553
ExceptionUtils.getRootCause(e).toString());
554+
// If the query has already timed out, do not retry. Re-dispatching the same FragmentInstance
555+
// may cause it to be executed twice on the remote node (see the REPEATED_RPC_CALL handling in
556+
// FragmentInstanceManager), so we fail fast with a timeout status instead.
557+
long currentTime = System.currentTimeMillis();
558+
if (currentTime - queryContext.getStartTime() >= queryContext.getTimeOut()) {
559+
throw new FragmentInstanceDispatchException(
560+
RpcUtils.getStatus(
561+
TSStatusCode.QUERY_TIMEOUT,
562+
String.format(
563+
QueryTimeoutRuntimeException.QUERY_TIMEOUT_EXCEPTION_MESSAGE,
564+
queryContext.getStartTime(),
565+
queryContext.getStartTime() + queryContext.getTimeOut(),
566+
currentTime)));
567+
}
553568
// we just retry once to clear stale connection for a restart node.
554569
try {
555570
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
@@ -159,7 +159,7 @@ private static TSStatus tryCatchQueryException(Exception e) {
159159

160160
Throwable t = e instanceof ExecutionException ? e.getCause() : e;
161161
if (t instanceof QueryTimeoutRuntimeException) {
162-
return RpcUtils.getStatus(TSStatusCode.INTERNAL_REQUEST_TIME_OUT, rootCause.getMessage());
162+
return RpcUtils.getStatus(TSStatusCode.QUERY_TIMEOUT, rootCause.getMessage());
163163
} else if (t instanceof ParseCancellationException) {
164164
return RpcUtils.getStatus(
165165
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)