[ML] Mitigate IOSession timeouts (#115414) (#115526)

prwhelan · web-flow · commit d2eebff6e912 · 2024-10-25T00:30:48.000+11:00
We are seeing exceptions ~0.03% of the time in our integration tests: ``` org.apache.http.ConnectionClosedException: Connection closed unexpectedly ``` The `contentDecoder` does not always fully consume the body within `SimpleInputBuffer.consumeContent`. When we return back to Apache, the rest of the body is never delivered, and the IOSession eventually times out and gets cleaned up. During that cleanup process, Apache calls our Consumer with the above exception. If we read 0 bytes and return back immediately, Apache has a better chance to load the rest of the body/footer, and it will call `consumeContent` again. This reduces the exception rate down to ~0.001%. Fix #114105 Fix #114232 Fix #114327 Fix #114385
diff --git a/docs/changelog/115414.yaml b/docs/changelog/115414.yaml
@@ -0,0 +1,9 @@
+pr: 115414
+summary: Mitigate IOSession timeouts
+area: Machine Learning
+type: bug
+issues:
+ - 114385
+ - 114327
+ - 114105
+ - 114232
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/StreamingHttpResultPublisher.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/StreamingHttpResultPublisher.java
@@ -96,11 +96,10 @@ public void consumeContent(ContentDecoder contentDecoder, IOControl ioControl) t
 
         try {
             var consumed = inputBuffer.consumeContent(contentDecoder);
-            var allBytes = new byte[consumed];
-            inputBuffer.read(allBytes);
-
-            // we can have empty bytes, don't bother sending them
-            if (allBytes.length > 0) {
+            // we could have read 0 bytes if the body was delayed getting in, we need to return out so apache can load the body/footer
+            if (consumed > 0) {
+                var allBytes = new byte[consumed];
+                inputBuffer.read(allBytes);
                 queue.offer(() -> {
                     subscriber.onNext(new HttpResult(response, allBytes));
                     var currentBytesInQueue = bytesInQueue.updateAndGet(current -> Long.max(0, current - allBytes.length));
@@ -111,18 +110,17 @@ public void consumeContent(ContentDecoder contentDecoder, IOControl ioControl) t
                         }
                     }
                 });
-            }
 
-            // always check if totalByteSize > the configured setting in case the settings change
-            if (bytesInQueue.accumulateAndGet(allBytes.length, Long::sum) >= settings.getMaxResponseSize().getBytes()) {
-                pauseProducer(ioControl);
-            }
+                // always check if totalByteSize > the configured setting in case the settings change
+                if (bytesInQueue.accumulateAndGet(allBytes.length, Long::sum) >= settings.getMaxResponseSize().getBytes()) {
+                    pauseProducer(ioControl);
+                }
 
-            // always run in case we're waking up from a pause and need to start a new thread
-            taskRunner.requestNextRun();
+                taskRunner.requestNextRun();
 
-            if (listenerCalled.compareAndSet(false, true)) {
-                listener.onResponse(this);
+                if (listenerCalled.compareAndSet(false, true)) {
+                    listener.onResponse(this);
+                }
             }
         } finally {
             inputBuffer.reset();
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/InferenceEventsAssertion.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/InferenceEventsAssertion.java
@@ -12,7 +12,6 @@
 import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.inference.InferenceServiceResults;
-import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xcontent.XContentFactory;
 import org.hamcrest.MatcherAssert;
 import org.hamcrest.Matchers;
@@ -26,6 +25,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 
+import static org.elasticsearch.test.ESTestCase.fail;
 import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS;
 import static org.hamcrest.CoreMatchers.is;
 
@@ -47,7 +47,9 @@ public InferenceEventsAssertion hasFinishedStream() {
     }
 
     public InferenceEventsAssertion hasNoErrors() {
-        MatcherAssert.assertThat("Expected no errors from stream.", error, Matchers.nullValue());
+        if (error != null) {
+            fail(error, "Expected no errors from stream.");
+        }
         return this;
     }
 
@@ -66,7 +68,7 @@ public InferenceEventsAssertion hasErrorWithStatusCode(int statusCode) {
             }
             t = t.getCause();
         }
-        ESTestCase.fail(error, "Expected an underlying ElasticsearchStatusException.");
+        fail(error, "Expected an underlying ElasticsearchStatusException.");
         return this;
     }
 
@@ -79,7 +81,7 @@ public InferenceEventsAssertion hasErrorContaining(String message) {
             }
             t = t.getCause();
         }
-        ESTestCase.fail(error, "Expected exception to contain string: " + message);
+        fail(error, "Expected exception to contain string: " + message);
         return this;
     }
 
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/anthropic/AnthropicServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/anthropic/AnthropicServiceTests.java
@@ -532,7 +532,6 @@ public void testInfer_SendsCompletionRequest() throws IOException {
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest() throws Exception {
         String responseJson = """
             data: {"type": "message_start", "message": {"model": "claude, probably"}}
@@ -578,7 +577,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException {
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest_ErrorResponse() throws Exception {
         String responseJson = """
             data: {"type": "error", "error": {"type": "request_too_large", "message": "blah"}}
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java
@@ -1197,7 +1197,6 @@ public void testInfer_UnauthorisedResponse() throws IOException {
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest() throws Exception {
         String responseJson = """
             data: {\
@@ -1253,7 +1252,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException, URISy
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest_ErrorResponse() throws Exception {
         String responseJson = """
             {
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java
@@ -1322,7 +1322,6 @@ private void testChunkedInfer(AzureOpenAiEmbeddingsModel model) throws IOExcepti
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest() throws Exception {
         String responseJson = """
             data: {\
@@ -1381,7 +1380,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException, URISy
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest_ErrorResponse() throws Exception {
         String responseJson = """
             {
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java
@@ -1532,7 +1532,6 @@ public void testDefaultSimilarity() {
         assertEquals(SimilarityMeasure.DOT_PRODUCT, CohereService.defaultSimilarity());
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest() throws Exception {
         String responseJson = """
             {"event_type":"text-generation", "text":"hello"}
@@ -1566,7 +1565,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException {
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest_ErrorResponse() throws Exception {
         String responseJson = """
             { "event_type":"stream-end", "finish_reason":"ERROR", "response":{ "text": "how dare you" } }
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java
@@ -914,7 +914,6 @@ public void testInfer_SendsRequest() throws IOException {
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest() throws Exception {
         String responseJson = """
             data: {\
@@ -964,7 +963,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException {
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")
     public void testInfer_StreamRequest_ErrorResponse() throws Exception {
         String responseJson = """
             {

Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,6 @@`
`12`	`12`	`import org.elasticsearch.common.xcontent.ChunkedToXContent;`
`13`	`13`	`import org.elasticsearch.common.xcontent.XContentHelper;`
`14`	`14`	`import org.elasticsearch.inference.InferenceServiceResults;`
`15`		`-import org.elasticsearch.test.ESTestCase;`
`16`	`15`	`import org.elasticsearch.xcontent.XContentFactory;`
`17`	`16`	`import org.hamcrest.MatcherAssert;`
`18`	`17`	`import org.hamcrest.Matchers;`
`@@ -26,6 +25,7 @@`
`26`	`25`	`import java.util.concurrent.atomic.AtomicInteger;`
`27`	`26`	`import java.util.stream.Stream;`
`28`	`27`
	`28`	`+import static org.elasticsearch.test.ESTestCase.fail;`
`29`	`29`	`import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS;`
`30`	`30`	`import static org.hamcrest.CoreMatchers.is;`
`31`	`31`
`@@ -47,7 +47,9 @@ public InferenceEventsAssertion hasFinishedStream() {`
`47`	`47`	`}`
`48`	`48`
`49`	`49`	`public InferenceEventsAssertion hasNoErrors() {`
`50`		`- MatcherAssert.assertThat("Expected no errors from stream.", error, Matchers.nullValue());`
	`50`	`+ if (error != null) {`
	`51`	`+ fail(error, "Expected no errors from stream.");`
	`52`	`+ }`
`51`	`53`	`return this;`
`52`	`54`	`}`
`53`	`55`
`@@ -66,7 +68,7 @@ public InferenceEventsAssertion hasErrorWithStatusCode(int statusCode) {`
`66`	`68`	`}`
`67`	`69`	`t = t.getCause();`
`68`	`70`	`}`
`69`		`- ESTestCase.fail(error, "Expected an underlying ElasticsearchStatusException.");`
	`71`	`+ fail(error, "Expected an underlying ElasticsearchStatusException.");`
`70`	`72`	`return this;`
`71`	`73`	`}`
`72`	`74`
`@@ -79,7 +81,7 @@ public InferenceEventsAssertion hasErrorContaining(String message) {`
`79`	`81`	`}`
`80`	`82`	`t = t.getCause();`
`81`	`83`	`}`
`82`		`- ESTestCase.fail(error, "Expected exception to contain string: " + message);`
	`84`	`+ fail(error, "Expected exception to contain string: " + message);`
`83`	`85`	`return this;`
`84`	`86`	`}`
`85`	`87`
Original file line number	Diff line number	Diff line change
`@@ -532,7 +532,6 @@ public void testInfer_SendsCompletionRequest() throws IOException {`
`532`	`532`	`}`
`533`	`533`	`}`
`534`	`534`
`535`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`536`	`535`	`public void testInfer_StreamRequest() throws Exception {`
`537`	`536`	`String responseJson = """`
`538`	`537`	`data: {"type": "message_start", "message": {"model": "claude, probably"}}`
`@@ -578,7 +577,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException {`
`578`	`577`	`}`
`579`	`578`	`}`
`580`	`579`
`581`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`582`	`580`	`public void testInfer_StreamRequest_ErrorResponse() throws Exception {`
`583`	`581`	`String responseJson = """`
`584`	`582`	`data: {"type": "error", "error": {"type": "request_too_large", "message": "blah"}}`
Original file line number	Diff line number	Diff line change
`@@ -1197,7 +1197,6 @@ public void testInfer_UnauthorisedResponse() throws IOException {`
`1197`	`1197`	`}`
`1198`	`1198`	`}`
`1199`	`1199`
`1200`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`1201`	`1200`	`public void testInfer_StreamRequest() throws Exception {`
`1202`	`1201`	`String responseJson = """`
`1203`	`1202`	`data: {\`
`@@ -1253,7 +1252,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException, URISy`
`1253`	`1252`	`}`
`1254`	`1253`	`}`
`1255`	`1254`
`1256`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`1257`	`1255`	`public void testInfer_StreamRequest_ErrorResponse() throws Exception {`
`1258`	`1256`	`String responseJson = """`
`1259`	`1257`	`{`
Original file line number	Diff line number	Diff line change
`@@ -1322,7 +1322,6 @@ private void testChunkedInfer(AzureOpenAiEmbeddingsModel model) throws IOExcepti`
`1322`	`1322`	`}`
`1323`	`1323`	`}`
`1324`	`1324`
`1325`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`1326`	`1325`	`public void testInfer_StreamRequest() throws Exception {`
`1327`	`1326`	`String responseJson = """`
`1328`	`1327`	`data: {\`
`@@ -1381,7 +1380,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException, URISy`
`1381`	`1380`	`}`
`1382`	`1381`	`}`
`1383`	`1382`
`1384`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`1385`	`1383`	`public void testInfer_StreamRequest_ErrorResponse() throws Exception {`
`1386`	`1384`	`String responseJson = """`
`1387`	`1385`	`{`
Original file line number	Diff line number	Diff line change
`@@ -1532,7 +1532,6 @@ public void testDefaultSimilarity() {`
`1532`	`1532`	`assertEquals(SimilarityMeasure.DOT_PRODUCT, CohereService.defaultSimilarity());`
`1533`	`1533`	`}`
`1534`	`1534`
`1535`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`1536`	`1535`	`public void testInfer_StreamRequest() throws Exception {`
`1537`	`1536`	`String responseJson = """`
`1538`	`1537`	`{"event_type":"text-generation", "text":"hello"}`
`@@ -1566,7 +1565,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException {`
`1566`	`1565`	`}`
`1567`	`1566`	`}`
`1568`	`1567`
`1569`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`1570`	`1568`	`public void testInfer_StreamRequest_ErrorResponse() throws Exception {`
`1571`	`1569`	`String responseJson = """`
`1572`	`1570`	`{ "event_type":"stream-end", "finish_reason":"ERROR", "response":{ "text": "how dare you" } }`
Original file line number	Diff line number	Diff line change
`@@ -914,7 +914,6 @@ public void testInfer_SendsRequest() throws IOException {`
`914`	`914`	`}`
`915`	`915`	`}`
`916`	`916`
`917`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`918`	`917`	`public void testInfer_StreamRequest() throws Exception {`
`919`	`918`	`String responseJson = """`
`920`	`919`	`data: {\`
`@@ -964,7 +963,6 @@ private InferenceServiceResults streamChatCompletion() throws IOException {`
`964`	`963`	`}`
`965`	`964`	`}`
`966`	`965`
`967`		`- @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114385")`
`968`	`966`	`public void testInfer_StreamRequest_ErrorResponse() throws Exception {`
`969`	`967`	`String responseJson = """`
`970`	`968`	`{`