[ML] fix race condition with job snapshot upgrade and feature reset (#85121) (#85396)

This addresses a race condition in feature reset && model snapshot upgrade.

It is possible that feature reset is called soon after the task is assigned, meaning the persistent task is cancelled before the model snapshot upgrader is set on the task.

This means that the task could continue executing even after being cancelled as nothing else checks if the feature is being reset or if the task has been cancelled.

closes #84709

Author: benwtrent

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/process/autodetect/AutodetectProcessManager.java

@@ -480,6 +480,17 @@ public void upgradeSnapshot(SnapshotUpgradeTask task, Consumer<Exception> closeH
                     );
                     return;
                 }
+                if (resetInProgress) {
+                    logger.trace(
+                        () -> new ParameterizedMessage(
+                            "Aborted upgrading snapshot [{}] for job [{}] as ML feature is being reset",
+                            snapshotId,
+                            jobId
+                        )
+                    );
+                    closeHandler.accept(null);
+                    return;
+                }
                 // We need to fork, otherwise we restore model state from a network thread (several GET api calls):
                 threadPool.executor(MachineLearning.UTILITY_THREAD_POOL_NAME).execute(new AbstractRunnable() {
                     @Override
@@ -500,6 +511,17 @@ protected void doRun() {
                             closeHandler.accept(null);
                             return;
                         }
+                        if (resetInProgress) {
+                            logger.trace(
+                                () -> new ParameterizedMessage(
+                                    "Aborted upgrading snapshot [{}] for job [{}] as ML feature is being reset",
+                                    snapshotId,
+                                    jobId
+                                )
+                            );
+                            closeHandler.accept(null);
+                            return;
+                        }
                         runSnapshotUpgrade(task, job, params, closeHandler);
                     }
                 });

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/process/autodetect/JobModelSnapshotUpgrader.java

@@ -102,7 +102,10 @@ public final class JobModelSnapshotUpgrader {
     }
 
     synchronized void start() {
-        task.setJobModelSnapshotUpgrader(this);
+        if (task.setJobModelSnapshotUpgrader(this) == false) {
+            this.killProcess(task.getReasonCancelled());
+            return;
+        }
 
         // A TP with no queue, so that we fail immediately if there are no threads available
         ExecutorService autodetectExecutorService = threadPool.executor(MachineLearning.JOB_COMMS_THREAD_POOL_NAME);

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/snapshot/upgrader/SnapshotUpgradeTask.java

@@ -70,7 +70,15 @@ protected synchronized void onCancelled() {
         }
     }
 
-    public synchronized void setJobModelSnapshotUpgrader(JobModelSnapshotUpgrader jobModelSnapshotUpgrader) {
+    /**
+     * @param jobModelSnapshotUpgrader the snapshot upgrader
+     * @return false if the task has been canceled and upgrader not set
+     */
+    public synchronized boolean setJobModelSnapshotUpgrader(JobModelSnapshotUpgrader jobModelSnapshotUpgrader) {
+        if (this.isCancelled()) {
+            return false;
+        }
         this.jobModelSnapshotUpgrader = jobModelSnapshotUpgrader;
+        return true;
     }
 }