Skip to content

Datanode: validate that a query id belongs to the calling session #18643

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ColinLeeo wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Datanode: validate that a query id belongs to the calling session #18643

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@

/** Compile-time i18n constants for DataNode misc subsystems (English). */
public final class DataNodeMiscMessages {
public static final String MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237 =
"The requested query does not belong to the current session.";

public static final String INVALID_PIPE_NAME =
"Invalid pipeName";
Expand Down
2 changes: 2 additions & 0 deletions iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@

/** 编译时国际化常量 - DataNode 杂项子系统（中文）。 */
public final class DataNodeMiscMessages {
public static final String MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237 =
"请求的查询不属于当前会话。";

public static final String INVALID_PIPE_NAME =
"无效的 pipeName";
Expand Down
19 changes: 19 additions & 0 deletions iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,25 @@ public void addQueryId(Long statementId, long queryId) {
queryIds.add(queryId);
}

@Override
public boolean containsQueryId(Long statementId, long queryId) {
return containsQueryId(statementIdToQueryId, statementId, queryId);
}

public static boolean containsQueryId(
Map<Long, Set<Long>> statementIdToQueryId, Long statementId, long queryId) {
if (statementId == null) {
for (Set<Long> queryIds : statementIdToQueryId.values()) {
if (queryIds != null && queryIds.contains(queryId)) {
return true;
}
}
return false;
}
Set<Long> queryIds = statementIdToQueryId.get(statementId);
return queryIds != null && queryIds.contains(queryId);
}

@Override
public void removeQueryId(Long statementId, Long queryId) {
removeQueryId(statementIdToQueryId, statementId, queryId);
Expand Down
3 changes: 3 additions & 0 deletions iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,9 @@ public ConnectionInfo convertToConnectionInfo() {

public abstract void addQueryId(Long statementId, long queryId);

// statementId could be null
public abstract boolean containsQueryId(Long statementId, long queryId);

// statementId could be null
public abstract void removeQueryId(Long statementId, Long queryId);

Expand Down
5 changes: 5 additions & 0 deletions ...re/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,11 @@ public void addQueryId(Long statementId, long queryId) {
queryIds.add(queryId);
}

@Override
public boolean containsQueryId(Long statementId, long queryId) {
return ClientSession.containsQueryId(statementIdToQueryId, statementId, queryId);
}

@Override
public void removeQueryId(Long statementId, Long queryId) {
ClientSession.removeQueryId(statementIdToQueryId, statementId, queryId);
Expand Down
5 changes: 5 additions & 0 deletions ...b-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,11 @@ public void addQueryId(Long statementId, long queryId) {
throw new UnsupportedOperationException();
}

@Override
public boolean containsQueryId(Long statementId, long queryId) {
return false;
}

@Override
public void removeQueryId(Long statementId, Long queryId) {
throw new UnsupportedOperationException();
Expand Down
5 changes: 5 additions & 0 deletions ...b-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,11 @@ public void addQueryId(Long statementId, long queryId) {
throw new UnsupportedOperationException();
}

@Override
public boolean containsQueryId(Long statementId, long queryId) {
return false;
}

@Override
public void removeQueryId(Long statementId, Long queryId) {
throw new UnsupportedOperationException();
Expand Down
92 changes: 66 additions & 26 deletions ...datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1561,6 +1561,7 @@ public TSFetchResultsResp fetchResultsV2(TSFetchResultsReq req) {
String statementType = null;
Throwable t = null;
IQueryExecution queryExecution = null;
boolean queryOwnedBySession = false;
IClientSession clientSession = SESSION_MANAGER.getCurrSessionAndUpdateIdleTime();
Long statementId = req.isSetStatementId() ? req.getStatementId() : null;
try {
Expand All @@ -1570,13 +1571,22 @@ public TSFetchResultsResp fetchResultsV2(TSFetchResultsReq req) {
}

queryExecution = COORDINATOR.getQueryExecution(req.queryId);

if (queryExecution == null) {
TSStatus noQueryExecutionStatus = new TSStatus(QUERY_WAS_KILLED.getStatusCode());
noQueryExecutionStatus.setMessage(NO_QUERY_EXECUTION_ERR_MSG);
return RpcUtils.getTSFetchResultsResp(noQueryExecutionStatus);
}

if (!clientSession.containsQueryId(statementId, req.queryId)) {
// The query is still running, but it was submitted by another session: do not stream its
// result and do not release it, so that the query which owns it is left untouched.
return RpcUtils.getTSFetchResultsResp(
RpcUtils.getStatus(
TSStatusCode.NO_PERMISSION,
DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237));
}
queryOwnedBySession = true;

TSFetchResultsResp resp = RpcUtils.getTSFetchResultsResp(TSStatusCode.SUCCESS_STATUS);

queryExecution.updateCurrentRpcStartTime(startTime);
Expand Down Expand Up @@ -1605,19 +1615,21 @@ public TSFetchResultsResp fetchResultsV2(TSFetchResultsReq req) {
throw error;
} finally {

long currentOperationCost = System.nanoTime() - startTime;
COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);

// record each operation time cost
CommonUtils.addStatementExecutionLatency(
OperationType.FETCH_RESULTS, statementType, currentOperationCost);
if (queryOwnedBySession) {
long currentOperationCost = System.nanoTime() - startTime;
COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);

if (finished) {
// record total time cost for one query
long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
CommonUtils.addQueryLatency(
StatementType.QUERY, executionTime > 0 ? executionTime : currentOperationCost);
clearUp(clientSession, statementId, req.queryId, req, t);
// record each operation time cost
CommonUtils.addStatementExecutionLatency(
OperationType.FETCH_RESULTS, statementType, currentOperationCost);

if (finished) {
// record total time cost for one query
long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
CommonUtils.addQueryLatency(
StatementType.QUERY, executionTime > 0 ? executionTime : currentOperationCost);
clearUp(clientSession, statementId, req.queryId, req, t);
}
}

SESSION_MANAGER.updateIdleTime();
Expand Down Expand Up @@ -1711,8 +1723,22 @@ public TSStatus cancelOperation(TSCancelOperationReq req) {

@Override
public TSStatus closeOperation(TSCloseOperationReq req) {
IClientSession clientSession = SESSION_MANAGER.getCurrSession();
if (req.isSetQueryId()
&& clientSession != null
&& clientSession.isLogin()
&& COORDINATOR.getQueryExecution(req.queryId) != null
&& !clientSession.containsQueryId(
req.isSetStatementId() ? req.getStatementId() : null, req.queryId)) {
// The queryId indexes the process-wide map of running queries, so only the session that
// submitted the query may release it. Queries that are no longer running keep the previous
// behaviour: releasing an unknown queryId stays a no-op.
return RpcUtils.getStatus(
TSStatusCode.NO_PERMISSION,
DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237);
}
return SESSION_MANAGER.closeOperation(
SESSION_MANAGER.getCurrSession(),
clientSession,
req.queryId,
req.statementId,
req.isSetStatementId(),
Expand Down Expand Up @@ -2319,6 +2345,7 @@ public TSFetchResultsResp fetchResults(TSFetchResultsReq req) {
String statementType = null;
Throwable t = null;
IQueryExecution queryExecution = null;
boolean queryOwnedBySession = false;
IClientSession clientSession = SESSION_MANAGER.getCurrSessionAndUpdateIdleTime();
Long statementId = req.isSetStatementId() ? req.getStatementId() : null;
try {
Expand All @@ -2333,6 +2360,17 @@ public TSFetchResultsResp fetchResults(TSFetchResultsReq req) {
noQueryExecutionStatus.setMessage(NO_QUERY_EXECUTION_ERR_MSG);
return RpcUtils.getTSFetchResultsResp(noQueryExecutionStatus);
}

if (!clientSession.containsQueryId(statementId, req.queryId)) {
// The query is still running, but it was submitted by another session: do not stream its
// result and do not release it, so that the query which owns it is left untouched.
return RpcUtils.getTSFetchResultsResp(
RpcUtils.getStatus(
TSStatusCode.NO_PERMISSION,
DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237));
}
queryOwnedBySession = true;

queryExecution.updateCurrentRpcStartTime(startTime);
statementType = queryExecution.getStatementType();

Expand Down Expand Up @@ -2360,19 +2398,21 @@ public TSFetchResultsResp fetchResults(TSFetchResultsReq req) {
throw error;
} finally {

long currentOperationCost = System.nanoTime() - startTime;
COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
if (queryOwnedBySession) {
long currentOperationCost = System.nanoTime() - startTime;
COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);

// record each operation time cost
CommonUtils.addStatementExecutionLatency(
OperationType.FETCH_RESULTS, statementType, currentOperationCost);

if (finished) {
// record total time cost for one query
long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
CommonUtils.addQueryLatency(
StatementType.QUERY, executionTime > 0 ? executionTime : currentOperationCost);
clearUp(clientSession, statementId, req.queryId, req, t);
// record each operation time cost
CommonUtils.addStatementExecutionLatency(
OperationType.FETCH_RESULTS, statementType, currentOperationCost);

if (finished) {
// record total time cost for one query
long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
CommonUtils.addQueryLatency(
StatementType.QUERY, executionTime > 0 ? executionTime : currentOperationCost);
clearUp(clientSession, statementId, req.queryId, req, t);
}
}

SESSION_MANAGER.updateIdleTime();
Expand Down
Loading
Loading