Move clearing stats queue to reset action

Author: jan-elastic

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearning.java

@@ -806,7 +806,6 @@ public void loadExtensions(ExtensionLoader loader) {
     private final SetOnce<MlAutoscalingDeciderService> mlAutoscalingDeciderService = new SetOnce<>();
     private final SetOnce<DeploymentManager> deploymentManager = new SetOnce<>();
     private final SetOnce<TrainedModelAssignmentClusterService> trainedModelAllocationClusterService = new SetOnce<>();
-    private final SetOnce<TrainedModelStatsService> trainedModelStatsService = new SetOnce<>();
 
     private final SetOnce<MachineLearningExtension> machineLearningExtension = new SetOnce<>();
 
@@ -1165,14 +1164,12 @@ public Collection<?> createComponents(PluginServices services) {
         this.datafeedRunner.set(datafeedRunner);
 
         // Inference components
-        trainedModelStatsService.set(
-            new TrainedModelStatsService(
-                resultsPersisterService,
-                originSettingClient,
-                indexNameExpressionResolver,
-                clusterService,
-                threadPool
-            )
+        final TrainedModelStatsService trainedModelStatsService = new TrainedModelStatsService(
+            resultsPersisterService,
+            originSettingClient,
+            indexNameExpressionResolver,
+            clusterService,
+            threadPool
         );
         final TrainedModelCacheMetadataService trainedModelCacheMetadataService = new TrainedModelCacheMetadataService(
             clusterService,
@@ -1188,7 +1185,7 @@ public Collection<?> createComponents(PluginServices services) {
             inferenceAuditor,
             threadPool,
             clusterService,
-            trainedModelStatsService.get(),
+            trainedModelStatsService,
             settings,
             clusterService.getNodeName(),
             inferenceModelBreaker.get(),
@@ -1395,6 +1392,7 @@ public Collection<?> createComponents(PluginServices services) {
             trainedModelProvider,
             trainedModelAssignmentService,
             trainedModelAllocationClusterService.get(),
+            trainedModelStatsService,
             deploymentManager.get(),
             nodeAvailabilityZoneMapper,
             new MachineLearningExtensionHolder(machineLearningExtension.get()),
@@ -2156,7 +2154,6 @@ public void cleanUpFeature(
         ActionListener<ResetFeatureStateResponse.ResetFeatureStateStatus> unsetResetModeListener = ActionListener.wrap(success -> {
 
             client.execute(SetResetModeAction.INSTANCE, SetResetModeActionRequest.disabled(true), ActionListener.wrap(resetSuccess -> {
-                trainedModelStatsService.get().clearQueue();
                 finalListener.onResponse(success);
                 logger.info("Finished machine learning feature reset");
             }, resetFailure -> {

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportResetMlComponentsAction.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.xpack.ml.inference.TrainedModelStatsService;
 import org.elasticsearch.xpack.ml.notifications.AnomalyDetectionAuditor;
 import org.elasticsearch.xpack.ml.notifications.DataFrameAnalyticsAuditor;
+import org.elasticsearch.xpack.ml.notifications.InferenceAuditor;
 
 import java.io.IOException;
 import java.util.List;
@@ -34,6 +35,7 @@ public class TransportResetMlComponentsAction extends TransportNodesAction<
 
     private final AnomalyDetectionAuditor anomalyDetectionAuditor;
     private final DataFrameAnalyticsAuditor dfaAuditor;
+    private final InferenceAuditor inferenceAuditor;
     private final TrainedModelStatsService trainedModelStatsService;
 
     @Inject
@@ -44,6 +46,7 @@ public TransportResetMlComponentsAction(
         ActionFilters actionFilters,
         AnomalyDetectionAuditor anomalyDetectionAuditor,
         DataFrameAnalyticsAuditor dfaAuditor,
+        InferenceAuditor inferenceAuditor,
         TrainedModelStatsService trainedModelStatsService
     ) {
         super(
@@ -56,6 +59,7 @@ public TransportResetMlComponentsAction(
         );
         this.anomalyDetectionAuditor = anomalyDetectionAuditor;
         this.dfaAuditor = dfaAuditor;
+        this.inferenceAuditor = inferenceAuditor;
         this.trainedModelStatsService = trainedModelStatsService;
     }
 
@@ -82,6 +86,7 @@ protected ResetMlComponentsAction.Response.ResetResponse newNodeResponse(StreamI
     protected ResetMlComponentsAction.Response.ResetResponse nodeOperation(ResetMlComponentsAction.NodeRequest request, Task task) {
         anomalyDetectionAuditor.reset();
         dfaAuditor.reset();
+        inferenceAuditor.reset();
         trainedModelStatsService.clearQueue();
         return new ResetMlComponentsAction.Response.ResetResponse(clusterService.localNode(), true);
     }