[ML] Close results stream before data frame analytics job stops

x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/ClassificationIT.java

@@ -481,7 +481,6 @@ public void testWithOnlyTrainingRowsAndTrainingPercentIsFifty_DependentVariableI
             "classification_training_percent_is_50_boolean", BOOLEAN_FIELD, BOOLEAN_FIELD_VALUES, "boolean");
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/67581")
     public void testStopAndRestart() throws Exception {
         initialize("classification_stop_and_restart");
         String predictedClassField = KEYWORD_FIELD + "_prediction";

x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/RegressionIT.java

@@ -33,9 +33,7 @@
 import org.elasticsearch.xpack.core.ml.inference.TrainedModelConfig;
 import org.elasticsearch.xpack.core.ml.inference.preprocessing.OneHotEncoding;
 import org.elasticsearch.xpack.core.ml.inference.preprocessing.PreProcessor;
-import org.elasticsearch.xpack.core.ml.utils.PhaseProgress;
 import org.junit.After;
-import org.junit.Before;
 
 import java.io.IOException;
 import java.time.Instant;
@@ -69,23 +67,9 @@ public class RegressionIT extends MlNativeDataFrameAnalyticsIntegTestCase {
     private String sourceIndex;
     private String destIndex;
 
-    @Before
-    public void setupLogging() {
-        client().admin().cluster()
-            .prepareUpdateSettings()
-            .setTransientSettings(Settings.builder()
-                .put("logger.org.elasticsearch.xpack.ml.dataframe", "DEBUG"))
-            .get();
-    }
-
     @After
     public void cleanup() {
         cleanUp();
-        client().admin().cluster()
-            .prepareUpdateSettings()
-            .setTransientSettings(Settings.builder()
-                .putNull("logger.org.elasticsearch.xpack.ml.dataframe"))
-            .get();
     }
 
     @Override

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/process/AnalyticsProcess.java

@@ -25,14 +25,6 @@ public interface AnalyticsProcess<ProcessResult> extends NativeProcess {
      */
     Iterator<ProcessResult> readAnalyticsResults();
 
-    /**
-     * Read anything left in the stream before
-     * closing the stream otherwise if the process
-     * tries to write more after the close it gets
-     * a SIGPIPE
-     */
-    void consumeAndCloseOutputStream();
-
     /**
      *
      * @return the process config

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/process/AnalyticsResultProcessor.java

@@ -23,8 +23,8 @@
 import org.elasticsearch.xpack.ml.dataframe.process.results.TrainedModelDefinitionChunk;
 import org.elasticsearch.xpack.ml.dataframe.stats.StatsHolder;
 import org.elasticsearch.xpack.ml.dataframe.stats.StatsPersister;
-import org.elasticsearch.xpack.ml.inference.modelsize.ModelSizeInfo;
 import org.elasticsearch.xpack.ml.extractor.ExtractedFields;
+import org.elasticsearch.xpack.ml.inference.modelsize.ModelSizeInfo;
 import org.elasticsearch.xpack.ml.inference.persistence.TrainedModelProvider;
 import org.elasticsearch.xpack.ml.notifications.DataFrameAnalyticsAuditor;
 
@@ -117,7 +117,6 @@ public void process(AnalyticsProcess<AnalyticsResult> process) {
                 completeResultsProgress();
             }
             completionLatch.countDown();
-            process.consumeAndCloseOutputStream();
         }
     }
 

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/process/MemoryUsageEstimationProcessManager.java

@@ -98,7 +98,6 @@ private MemoryUsageEstimationResult runJob(String jobId,
                     jobId, e.getMessage(), process.readError()).getFormattedMessage();
             throw ExceptionsHelper.serverError(errorMsg, e);
         } finally {
-            process.consumeAndCloseOutputStream();
             try {
                 LOGGER.debug("[{}] Closing process", jobId);
                 process.close();

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

@@ -110,12 +110,4 @@ public interface AutodetectProcess extends NativeProcess {
      * @return stream of autodetect results.
      */
     Iterator<AutodetectResult> readAutodetectResults();
-
-    /**
-     * Read anything left in the stream before
-     * closing the stream otherwise if the process
-     * tries to write more after the close it gets
-     * a SIGPIPE
-     */
-    void consumeAndCloseOutputStream();
 }

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

@@ -173,10 +173,6 @@ public AutodetectResult next() {
         };
     }
 
-    @Override
-    public void consumeAndCloseOutputStream() {
-    }
-
     @Override
     public ZonedDateTime getProcessStartTime() {
         return startTime;

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

@@ -187,24 +187,20 @@ public void process() {
 
     private void readResults() {
         currentRunBucketCount = 0;
-        try {
-            Iterator<AutodetectResult> iterator = process.readAutodetectResults();
-            while (iterator.hasNext()) {
-                try {
-                    AutodetectResult result = iterator.next();
-                    processResult(result);
-                    if (result.getBucket() != null) {
-                        LOGGER.trace("[{}] Bucket number {} parsed from output", jobId, currentRunBucketCount);
-                    }
-                } catch (Exception e) {
-                    if (isAlive() == false) {
-                        throw e;
-                    }
-                    LOGGER.warn(new ParameterizedMessage("[{}] Error processing autodetect result", jobId), e);
+        Iterator<AutodetectResult> iterator = process.readAutodetectResults();
+        while (iterator.hasNext()) {
+            try {
+                AutodetectResult result = iterator.next();
+                processResult(result);
+                if (result.getBucket() != null) {
+                    LOGGER.trace("[{}] Bucket number {} parsed from output", jobId, currentRunBucketCount);
                 }
+            } catch (Exception e) {
+                if (isAlive() == false) {
+                    throw e;
+                }
+                LOGGER.warn(new ParameterizedMessage("[{}] Error processing autodetect result", jobId), e);
             }
-        } finally {
-            process.consumeAndCloseOutputStream();
         }
     }
 

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

@@ -115,23 +115,19 @@ public void process() {
     }
 
     private void readResults() {
-        try {
-            Iterator<AutodetectResult> iterator = process.readAutodetectResults();
-            while (iterator.hasNext()) {
-                try {
-                    AutodetectResult result = iterator.next();
-                    processResult(result);
-                } catch (Exception e) {
-                    if (isAlive() == false) {
-                        throw e;
-                    }
-                    LOGGER.warn(
-                        new ParameterizedMessage("[{}] [{}] Error processing model snapshot upgrade result", jobId, snapshotId),
-                        e);
+        Iterator<AutodetectResult> iterator = process.readAutodetectResults();
+        while (iterator.hasNext()) {
+            try {
+                AutodetectResult result = iterator.next();
+                processResult(result);
+            } catch (Exception e) {
+                if (isAlive() == false) {
+                    throw e;
                 }
+                LOGGER.warn(
+                    new ParameterizedMessage("[{}] [{}] Error processing model snapshot upgrade result", jobId, snapshotId),
+                    e);
             }
-        } finally {
-            process.consumeAndCloseOutputStream();
         }
     }
 

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

@@ -176,12 +176,16 @@ public void close() throws IOException {
             if (processInStream() != null) {
                 processInStream().close();
             }
+
             // wait for the process to exit by waiting for end-of-file on the named pipe connected
             // to the state processor - it may take a long time for all the model state to be
             // indexed
             if (stateProcessorFuture != null) {
                 stateProcessorFuture.get(MachineLearningField.STATE_PERSIST_RESTORE_TIMEOUT.getMinutes(), TimeUnit.MINUTES);
             }
+
+            consumeAndCloseOutputStream();
+
             // the log processor should have stopped by now too - assume processing the logs will
             // take no more than 5 seconds longer than processing the state (usually it should
             // finish first)
@@ -312,13 +316,12 @@ protected boolean isProcessKilled() {
         return processKilled;
     }
 
-    public void consumeAndCloseOutputStream() {
-        try {
+    void consumeAndCloseOutputStream() {
+        try (InputStream outStream = processOutStream()) {
             byte[] buff = new byte[512];
-            while (processOutStream().read(buff) >= 0) {
+            while (outStream.read(buff) >= 0) {
                 // Do nothing
             }
-            processOutStream().close();
         } catch (IOException e) {
             // Given we are closing down the process there is no point propagating IO exceptions here
         }

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/dataframe/process/MemoryUsageEstimationProcessManagerTests.java

@@ -108,7 +108,6 @@ public void testRunJob_NoResults() throws Exception {
         InOrder inOrder = inOrder(process);
         inOrder.verify(process).readAnalyticsResults();
         inOrder.verify(process).readError();
-        inOrder.verify(process).consumeAndCloseOutputStream();
         inOrder.verify(process).close();
         verifyNoMoreInteractions(process, listener);
     }
@@ -127,7 +126,6 @@ public void testRunJob_MultipleResults() throws Exception {
         InOrder inOrder = inOrder(process);
         inOrder.verify(process).readAnalyticsResults();
         inOrder.verify(process).readError();
-        inOrder.verify(process).consumeAndCloseOutputStream();
         inOrder.verify(process).close();
         verifyNoMoreInteractions(process, listener);
     }
@@ -146,7 +144,6 @@ public void testRunJob_OneResult_ParseException() throws Exception {
         InOrder inOrder = inOrder(process);
         inOrder.verify(process).readAnalyticsResults();
         inOrder.verify(process).readError();
-        inOrder.verify(process).consumeAndCloseOutputStream();
         inOrder.verify(process).close();
         verifyNoMoreInteractions(process, listener);
     }
@@ -164,7 +161,6 @@ public void testRunJob_FailsOnClose() throws Exception {
 
         InOrder inOrder = inOrder(process);
         inOrder.verify(process).readAnalyticsResults();
-        inOrder.verify(process).consumeAndCloseOutputStream();
         inOrder.verify(process).close();
         inOrder.verify(process).readError();
         verifyNoMoreInteractions(process, listener);
@@ -186,7 +182,6 @@ public void testRunJob_FailsOnClose_ProcessReportsError() throws Exception {
 
         InOrder inOrder = inOrder(process);
         inOrder.verify(process).readAnalyticsResults();
-        inOrder.verify(process).consumeAndCloseOutputStream();
         inOrder.verify(process).close();
         inOrder.verify(process).readError();
         verifyNoMoreInteractions(process, listener);
@@ -201,7 +196,6 @@ public void testRunJob_Ok() throws Exception {
 
         InOrder inOrder = inOrder(process);
         inOrder.verify(process).readAnalyticsResults();
-        inOrder.verify(process).consumeAndCloseOutputStream();
         inOrder.verify(process).close();
         verifyNoMoreInteractions(process, listener);
     }

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/job/process/autodetect/NativeAutodetectProcessTests.java

@@ -8,15 +8,15 @@
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.core.ml.job.config.ModelPlotConfig;
-import org.elasticsearch.xpack.ml.process.IndexingStateProcessor;
 import org.elasticsearch.xpack.ml.job.process.autodetect.params.DataLoadParams;
 import org.elasticsearch.xpack.ml.job.process.autodetect.params.FlushJobParams;
 import org.elasticsearch.xpack.ml.job.process.autodetect.params.TimeRange;
 import org.elasticsearch.xpack.ml.job.process.autodetect.writer.AutodetectControlMsgWriter;
 import org.elasticsearch.xpack.ml.job.results.AutodetectResult;
+import org.elasticsearch.xpack.ml.process.IndexingStateProcessor;
+import org.elasticsearch.xpack.ml.process.NativeController;
 import org.elasticsearch.xpack.ml.process.ProcessPipes;
 import org.elasticsearch.xpack.ml.process.ProcessResultsParser;
-import org.elasticsearch.xpack.ml.process.NativeController;
 import org.elasticsearch.xpack.ml.process.logging.CppLogMessageHandler;
 import org.junit.Assert;
 import org.junit.Before;
@@ -37,7 +37,6 @@
 import java.util.concurrent.Future;
 import java.util.function.Consumer;
 
-import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.Matchers.any;
@@ -152,19 +151,6 @@ public void testPersistJob() throws IOException {
         testWriteMessage(NativeAutodetectProcess::persistState, AutodetectControlMsgWriter.BACKGROUND_PERSIST_MESSAGE_CODE);
     }
 
-    @SuppressWarnings("unchecked")
-    public void testConsumeAndCloseOutputStream() throws IOException {
-
-        try (NativeAutodetectProcess process = new NativeAutodetectProcess("foo", mock(NativeController.class),
-            processPipes, NUMBER_FIELDS, Collections.emptyList(),
-            new ProcessResultsParser<>(AutodetectResult.PARSER, NamedXContentRegistry.EMPTY), mock(Consumer.class))) {
-
-            process.start(executorService);
-            process.consumeAndCloseOutputStream();
-            assertThat(outputStream.available(), equalTo(0));
-        }
-    }
-
     @SuppressWarnings("unchecked")
     private void testWriteMessage(CheckedConsumer<NativeAutodetectProcess> writeFunction, String expectedMessageCode) throws IOException {
         try (NativeAutodetectProcess process = new NativeAutodetectProcess("foo", mock(NativeController.class),

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/process/AbstractNativeProcessTests.java

@@ -23,6 +23,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
@@ -78,8 +79,9 @@ public void initialize() throws IOException {
     }
 
     @After
-    public void terminateExecutorService() {
+    public void terminateExecutorService() throws IOException {
         ThreadPool.terminate(executorService, 10, TimeUnit.SECONDS);
+        assertThat(processPipes.getProcessOutStream().get().available(), equalTo(0));
         verifyNoMoreInteractions(onProcessCrash);
     }