Check response status of the abort task results request

trinodb · Nov 4, 2021 · 92387da · 92387da
1 parent c3f73b8
commit 92387da
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/core/trino-main/src/main/java/io/trino/operator/HttpPageBufferClient.java b/core/trino-main/src/main/java/io/trino/operator/HttpPageBufferClient.java
@@ -60,6 +60,7 @@
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
+import static io.airlift.http.client.HttpStatus.NO_CONTENT;
 import static io.airlift.http.client.HttpStatus.familyForStatusCode;
 import static io.airlift.http.client.Request.Builder.prepareDelete;
 import static io.airlift.http.client.Request.Builder.prepareGet;
@@ -469,6 +470,15 @@ private synchronized void sendDelete()
             public void onSuccess(@Nullable StatusResponse result)
             {
                 assertNotHoldsLock(this);
+
+                if (result.getStatusCode() != NO_CONTENT.code()) {
+                    onFailure(new TrinoTransportException(
+                            REMOTE_BUFFER_CLOSE_FAILED,
+                            fromUri(location),
+                            format("Error closing remote buffer, expected %s got %s", NO_CONTENT.code(), result.getStatusCode())));
+                    return;
+                }
+
                 backoff.success();
                 synchronized (HttpPageBufferClient.this) {
                     closed = true;

diff --git a/core/trino-main/src/test/java/io/trino/operator/TestHttpPageBufferClient.java b/core/trino-main/src/test/java/io/trino/operator/TestHttpPageBufferClient.java
@@ -293,6 +293,7 @@ public void testInvalidResponses()
         assertStatus(client, location, "queued", 0, 3, 3, 3, "not scheduled");
 
         // close client and verify
+        processor.setResponse(new TestingResponse(HttpStatus.NO_CONTENT, ImmutableListMultimap.of(), new byte[0]));
         client.close();
         requestComplete.await(10, TimeUnit.SECONDS);
         assertStatus(client, location, "closed", 0, 3, 4, 3, "not scheduled");

diff --git a/core/trino-main/src/test/java/io/trino/operator/TestingExchangeHttpClientHandler.java b/core/trino-main/src/test/java/io/trino/operator/TestingExchangeHttpClientHandler.java
@@ -58,7 +58,7 @@ public Response handle(Request request)
         ImmutableList<String> parts = ImmutableList.copyOf(Splitter.on("/").omitEmptyStrings().split(request.getUri().getPath()));
         if (request.getMethod().equals("DELETE")) {
             assertEquals(parts.size(), 1);
-            return new TestingResponse(HttpStatus.OK, ImmutableListMultimap.of(), new byte[0]);
+            return new TestingResponse(HttpStatus.NO_CONTENT, ImmutableListMultimap.of(), new byte[0]);
         }
 
         assertEquals(parts.size(), 2);