Code review comment: The most risky bug in this code is: Potential loss of TaskRuns when `replayCreateTaskRun` rejects a task due to `arrangeTaskRun` returning false.  You can modify the code like this: ```java if (!taskRunManager.arrangeTaskRun(taskRun)) {     LOG.warn("Submit task run to pending queue failed, reject the submit:{}", taskRun);     // Handle failure case properly, e.g., requeue, notify, or log as needed. } ``` Explanation: When `taskRunManager.arrangeTaskRun(taskRun)` returns false and the task submission is rejected, it's simply logged but there's no action taken to either re-attempt the arrangement or handle the failure case adequately (e.g., by adding it to a retry queue, notifying an administrator, or otherwise managing the failed submission). This could lead to lost task runs where tasks that fail to be arranged are neither processed nor properly accounted for. Such losses in a task scheduling system could lead to data inconsistencies, delays in processing, or other operational issues depending on the context and purpose of the system. Properly handling these failures would mitigate the risk of losing track of tasks that cannot be immediately arranged.

Code before review:
public void replayUpdateTaskRun(TaskRunStatusChange statusChange) {
Constants.TaskRunState fromStatus = statusChange.getFromStatus();
Constants.TaskRunState toStatus = statusChange.getToStatus();
Long taskId = statusChange.getTaskId();
LOG.info("replayUpdateTaskRun:" + statusChange);
if (fromStatus == Constants.TaskRunState.PENDING) {
Queue<TaskRun> taskRunQueue = taskRunManager.getPendingTaskRunMap().get(taskId);
if (taskRunQueue == null) {
return;
}
if (taskRunQueue.size() == 0) {
taskRunManager.getPendingTaskRunMap().remove(taskId);
return;
}
TaskRun pendingTaskRun = null;
List<TaskRun> tempQueue = Lists.newArrayList();
while (!taskRunQueue.isEmpty()) {
TaskRun taskRun = taskRunQueue.poll();
if (taskRun.getStatus().getQueryId().equals(statusChange.getQueryId())) {
pendingTaskRun = taskRun;
break;
} else {
tempQueue.add(taskRun);
}
}
taskRunQueue.addAll(tempQueue);
if (pendingTaskRun == null) {
LOG.warn("could not find query_id:{}, taskId:{}, when replay update pendingTaskRun",
statusChange.getQueryId(), taskId);
return;
}
TaskRunStatus status = pendingTaskRun.getStatus();
if (toStatus == Constants.TaskRunState.RUNNING) {
if (status.getQueryId().equals(statusChange.getQueryId())) {
status.setState(Constants.TaskRunState.RUNNING);
taskRunManager.getRunningTaskRunMap().put(taskId, pendingTaskRun);
}
} else if (toStatus == Constants.TaskRunState.FAILED) {
status.setErrorMessage(statusChange.getErrorMessage());
status.setErrorCode(statusChange.getErrorCode());
status.setState(Constants.TaskRunState.FAILED);
taskRunManager.getTaskRunHistory().addHistory(status);
} else if (toStatus == Constants.TaskRunState.SUCCESS) {
LOG.info("Replay update pendingTaskRun which is merged by others, query_id:{}, taskId:{}",
statusChange.getQueryId(), taskId);
status.setErrorMessage(statusChange.getErrorMessage());
status.setErrorCode(statusChange.getErrorCode());
status.setState(Constants.TaskRunState.SUCCESS);
status.setProgress(100);
status.setFinishTime(statusChange.getFinishTime());
taskRunManager.getTaskRunHistory().addHistory(status);
}
if (taskRunQueue.size() == 0) {
taskRunManager.getPendingTaskRunMap().remove(taskId);
}
} else if (fromStatus == Constants.TaskRunState.RUNNING &&
(toStatus == Constants.TaskRunState.SUCCESS || toStatus == Constants.TaskRunState.FAILED)) {
TaskRun runningTaskRun = taskRunManager.getRunningTaskRunMap().remove(taskId);
if (runningTaskRun != null) {
TaskRunStatus status = runningTaskRun.getStatus();
if (status.getQueryId().equals(statusChange.getQueryId())) {
if (toStatus == Constants.TaskRunState.FAILED) {
status.setErrorMessage(statusChange.getErrorMessage());
status.setErrorCode(statusChange.getErrorCode());
}
status.setState(toStatus);
status.setProgress(100);
status.setFinishTime(statusChange.getFinishTime());
status.setExtraMessage(statusChange.getExtraMessage());
taskRunManager.getTaskRunHistory().addHistory(status);
}
} else {
String queryId = statusChange.getQueryId();
TaskRunStatus status = taskRunManager.getTaskRunHistory().getTask(queryId);
if (status == null) {
return;
}
status.setExtraMessage(statusChange.getExtraMessage());
}
} else {
LOG.warn("Illegal TaskRun queryId:{} status transform from {} to {}",
statusChange.getQueryId(), fromStatus, toStatus);
}
}

Line under review: taskRunManager.getPendingTaskRunMap().remove(taskId);

Code after addressing the review:
public void replayUpdateTaskRun(TaskRunStatusChange statusChange) {
Constants.TaskRunState fromStatus = statusChange.getFromStatus();
Constants.TaskRunState toStatus = statusChange.getToStatus();
Long taskId = statusChange.getTaskId();
LOG.info("replayUpdateTaskRun:" + statusChange);
if (fromStatus == Constants.TaskRunState.PENDING) {
Queue<TaskRun> taskRunQueue = taskRunManager.getPendingTaskRunMap().get(taskId);
if (taskRunQueue == null) {
return;
}
if (taskRunQueue.size() == 0) {
taskRunManager.getPendingTaskRunMap().remove(taskId);
return;
}
TaskRun pendingTaskRun = null;
List<TaskRun> tempQueue = Lists.newArrayList();
while (!taskRunQueue.isEmpty()) {
TaskRun taskRun = taskRunQueue.poll();
if (taskRun.getStatus().getQueryId().equals(statusChange.getQueryId())) {
pendingTaskRun = taskRun;
break;
} else {
tempQueue.add(taskRun);
}
}
taskRunQueue.addAll(tempQueue);
if (pendingTaskRun == null) {
LOG.warn("could not find query_id:{}, taskId:{}, when replay update pendingTaskRun",
statusChange.getQueryId(), taskId);
return;
}
TaskRunStatus status = pendingTaskRun.getStatus();
if (toStatus == Constants.TaskRunState.RUNNING) {
if (status.getQueryId().equals(statusChange.getQueryId())) {
status.setState(Constants.TaskRunState.RUNNING);
taskRunManager.getRunningTaskRunMap().put(taskId, pendingTaskRun);
}
} else if (toStatus == Constants.TaskRunState.FAILED) {
status.setErrorMessage(statusChange.getErrorMessage());
status.setErrorCode(statusChange.getErrorCode());
status.setState(Constants.TaskRunState.FAILED);
taskRunManager.getTaskRunHistory().addHistory(status);
} else if (toStatus == Constants.TaskRunState.SUCCESS) {
LOG.info("Replay update pendingTaskRun which is merged by others, query_id:{}, taskId:{}",
statusChange.getQueryId(), taskId);
status.setErrorMessage(statusChange.getErrorMessage());
status.setErrorCode(statusChange.getErrorCode());
status.setState(Constants.TaskRunState.SUCCESS);
status.setProgress(100);
status.setFinishTime(statusChange.getFinishTime());
taskRunManager.getTaskRunHistory().addHistory(status);
}
if (taskRunQueue.size() == 0) {
taskRunManager.getPendingTaskRunMap().remove(taskId);
}
} else if (fromStatus == Constants.TaskRunState.RUNNING &&
(toStatus == Constants.TaskRunState.SUCCESS || toStatus == Constants.TaskRunState.FAILED)) {
TaskRun runningTaskRun = taskRunManager.getRunningTaskRunMap().remove(taskId);
if (runningTaskRun != null) {
TaskRunStatus status = runningTaskRun.getStatus();
if (status.getQueryId().equals(statusChange.getQueryId())) {
if (toStatus == Constants.TaskRunState.FAILED) {
status.setErrorMessage(statusChange.getErrorMessage());
status.setErrorCode(statusChange.getErrorCode());
}
status.setState(toStatus);
status.setProgress(100);
status.setFinishTime(statusChange.getFinishTime());
status.setExtraMessage(statusChange.getExtraMessage());
taskRunManager.getTaskRunHistory().addHistory(status);
}
} else {
String queryId = statusChange.getQueryId();
TaskRunStatus status = taskRunManager.getTaskRunHistory().getTask(queryId);
if (status == null) {
return;
}
status.setExtraMessage(statusChange.getExtraMessage());
}
} else {
LOG.warn("Illegal TaskRun queryId:{} status transform from {} to {}",
statusChange.getQueryId(), fromStatus, toStatus);
}
}