[FLINK-21735][coordination] Harden JobMaster#updateTaskExecutionState()

This closes apache#15196.

flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java

@@ -428,17 +428,26 @@ public CompletableFuture<Acknowledge> cancel(Time timeout) {
  @Override
  public CompletableFuture<Acknowledge> updateTaskExecutionState(
  final TaskExecutionState taskExecutionState) {
- checkNotNull(taskExecutionState, "taskExecutionState");
+ FlinkException taskExecutionException;
+ try {
+ checkNotNull(taskExecutionState, "taskExecutionState");
 
- if (schedulerNG.updateTaskExecutionState(taskExecutionState)) {
- return CompletableFuture.completedFuture(Acknowledge.get());
- } else {
- return FutureUtils.completedExceptionally(
- new ExecutionGraphException(
- "The execution attempt "
- + taskExecutionState.getID()
- + " was not found."));
+ if (schedulerNG.updateTaskExecutionState(taskExecutionState)) {
+ return CompletableFuture.completedFuture(Acknowledge.get());
+ } else {
+ taskExecutionException =
+ new ExecutionGraphException(
+ "The execution attempt "
+ + taskExecutionState.getID()
+ + " was not found.");
+ }
+ } catch (Exception e) {
+ taskExecutionException =
+ new JobMasterException(
+ "Could not update the state of task execution for JobMaster.", e);
+ handleJobMasterError(taskExecutionException);
  }
+ return FutureUtils.completedExceptionally(taskExecutionException);
  }
 
  @Override