Retry token removal only on 5xx response codes

wendigo · kokosing · commit 932f153eacd5 · 2022-02-03T21:06:57.000+01:00
diff --git a/client/trino-client/src/main/java/io/trino/client/auth/external/HttpTokenPoller.java b/client/trino-client/src/main/java/io/trino/client/auth/external/HttpTokenPoller.java
@@ -37,6 +37,7 @@
 import static io.trino.client.JsonCodec.jsonCodec;
 import static io.trino.client.JsonResponse.execute;
 import static java.lang.String.format;
+import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
 import static java.net.HttpURLConnection.HTTP_OK;
 import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
 import static java.time.temporal.ChronoUnit.MILLIS;
@@ -96,7 +97,7 @@ public void tokenReceived(URI tokenUri)
                     .withMaxAttempts(-1)
                     .withMaxDuration(Duration.ofSeconds(4))
                     .withBackoff(100, 500, MILLIS)
-                    .handleResultIf(code -> code != HTTP_OK))
+                    .handleResultIf(code -> code >= HTTP_INTERNAL_ERROR))
                     .get(() -> {
                         Request request = prepareRequestBuilder(tokenUri)
                                 .delete()
diff --git a/client/trino-client/src/test/java/io/trino/client/auth/external/TestHttpTokenPoller.java b/client/trino-client/src/test/java/io/trino/client/auth/external/TestHttpTokenPoller.java
@@ -176,20 +176,16 @@ public void testTokenReceived()
     }
 
     @Test
-    public void testTokenReceivedRetriesUntilHTTP_OK()
+    public void testTokenReceivedRetriesUntilNotErrorReturned()
     {
         server.enqueue(status(HTTP_UNAVAILABLE));
         server.enqueue(status(HTTP_UNAVAILABLE));
         server.enqueue(status(HTTP_UNAVAILABLE));
         server.enqueue(status(202));
-        server.enqueue(status(303));
-        server.enqueue(status(HTTP_OK));
-        server.enqueue(status(HTTP_OK));
-        server.enqueue(status(HTTP_OK));
 
         tokenPoller.tokenReceived(tokenUri());
 
-        assertThat(server.getRequestCount()).isEqualTo(6);
+        assertThat(server.getRequestCount()).isEqualTo(4);
     }
 
     @Test

Original file line number	Diff line number	Diff line change
`@@ -176,20 +176,16 @@ public void testTokenReceived()`
`176`	`176`	`}`
`177`	`177`
`178`	`178`	`@Test`
`179`		`- public void testTokenReceivedRetriesUntilHTTP_OK()`
	`179`	`+ public void testTokenReceivedRetriesUntilNotErrorReturned()`
`180`	`180`	`{`
`181`	`181`	`server.enqueue(status(HTTP_UNAVAILABLE));`
`182`	`182`	`server.enqueue(status(HTTP_UNAVAILABLE));`
`183`	`183`	`server.enqueue(status(HTTP_UNAVAILABLE));`
`184`	`184`	`server.enqueue(status(202));`
`185`		`- server.enqueue(status(303));`
`186`		`- server.enqueue(status(HTTP_OK));`
`187`		`- server.enqueue(status(HTTP_OK));`
`188`		`- server.enqueue(status(HTTP_OK));`
`189`	`185`
`190`	`186`	`tokenPoller.tokenReceived(tokenUri());`
`191`	`187`
`192`		`- assertThat(server.getRequestCount()).isEqualTo(6);`
	`188`	`+ assertThat(server.getRequestCount()).isEqualTo(4);`
`193`	`189`	`}`
`194`	`190`
`195`	`191`	`@Test`