From 51b2f6cb4524a79fd9ea4abbe4e35416f30b1cdb Mon Sep 17 00:00:00 2001
From: Abhishek Radhakrishnan <abhishek.rb19@gmail.com>
Date: Mon, 17 Jun 2024 12:48:15 -0700
Subject: [PATCH] Fix retry logic in `BrokerClient` and flakey
 `BrokerClientTest` (#16618)

* Fix flakey BrokerClientTest.

The testError() method reliably fails in the IDE. This is because the
the test runner has

<surefire.rerunFailingTestsCount>3</surefire.rerunFailingTestsCount> is set to 3, so maven
retries this "flaky test" multiple times and the test code returns a successful response
in the third attempt.

The exception handling in BrokerClientTest was broken:
- All non-2xx errors were being turned as 5xx errors. Remove that block of
code. If we need to handle retries of more specific 5xx error codes, that should be
hanlded explicitly. Or if there's a source of 4xx class error that needs to be 5xx,
fix that in the source of error.

* Fix CodeQL warning for unused parameter.
---
 .../apache/druid/discovery/BrokerClient.java  |  8 ++---
 .../druid/discovery/BrokerClientTest.java     | 30 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/server/src/main/java/org/apache/druid/discovery/BrokerClient.java b/server/src/main/java/org/apache/druid/discovery/BrokerClient.java
index a64c8d670e4..bdee9b8dfe4 100644
--- a/server/src/main/java/org/apache/druid/discovery/BrokerClient.java
+++ b/server/src/main/java/org/apache/druid/discovery/BrokerClient.java
@@ -33,7 +33,6 @@ import org.jboss.netty.channel.ChannelException;
 import org.jboss.netty.handler.codec.http.HttpMethod;
 import org.jboss.netty.handler.codec.http.HttpResponseStatus;
 
-import javax.servlet.http.HttpServletResponse;
 import java.io.IOException;
 import java.net.MalformedURLException;
 import java.net.URL;
@@ -88,10 +87,6 @@ public class BrokerClient
             throw DruidException.forPersona(DruidException.Persona.OPERATOR)
                                 .ofCategory(DruidException.Category.RUNTIME_FAILURE)
                                 .build("Request to broker failed due to failed response status: [%s]", responseStatus);
-          } else if (responseStatus.getCode() != HttpServletResponse.SC_OK) {
-            throw DruidException.forPersona(DruidException.Persona.OPERATOR)
-                                .ofCategory(DruidException.Category.RUNTIME_FAILURE)
-                                .build("Request to broker failed due to failed response code: [%s]", responseStatus.getCode());
           }
           return fullResponseHolder.getContent();
         },
@@ -99,6 +94,9 @@ public class BrokerClient
           if (throwable instanceof ExecutionException) {
             return throwable.getCause() instanceof IOException || throwable.getCause() instanceof ChannelException;
           }
+          if (throwable instanceof DruidException) {
+            return ((DruidException) throwable).getCategory() == DruidException.Category.RUNTIME_FAILURE;
+          }
           return throwable instanceof IOE;
         },
         MAX_RETRIES
diff --git a/server/src/test/java/org/apache/druid/discovery/BrokerClientTest.java b/server/src/test/java/org/apache/druid/discovery/BrokerClientTest.java
index 333882a4305..de03877a9b0 100644
--- a/server/src/test/java/org/apache/druid/discovery/BrokerClientTest.java
+++ b/server/src/test/java/org/apache/druid/discovery/BrokerClientTest.java
@@ -101,7 +101,7 @@ public class BrokerClientTest extends BaseJettyTest
   }
 
   @Test
-  public void testError() throws Exception
+  public void testRetryableError() throws Exception
   {
     DruidNodeDiscovery druidNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
     EasyMock.expect(druidNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(discoveryDruidNode)).anyTimes();
@@ -121,6 +121,26 @@ public class BrokerClientTest extends BaseJettyTest
     Assert.assertEquals("hello", brokerClient.sendQuery(request));
   }
 
+  @Test
+  public void testNonRetryableError() throws Exception
+  {
+    DruidNodeDiscovery druidNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
+    EasyMock.expect(druidNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(discoveryDruidNode)).anyTimes();
+
+    DruidNodeDiscoveryProvider druidNodeDiscoveryProvider = EasyMock.createMock(DruidNodeDiscoveryProvider.class);
+    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.BROKER)).andReturn(druidNodeDiscovery);
+
+    EasyMock.replay(druidNodeDiscovery, druidNodeDiscoveryProvider);
+
+    BrokerClient brokerClient = new BrokerClient(
+        httpClient,
+        druidNodeDiscoveryProvider
+    );
+
+    Request request = brokerClient.makeRequest(HttpMethod.POST, "/simple/error");
+    Assert.assertEquals("", brokerClient.sendQuery(request));
+  }
+
   @Path("/simple")
   public static class SimpleResource
   {
@@ -150,5 +170,13 @@ public class BrokerClientTest extends BaseJettyTest
         return Response.status(504).build();
       }
     }
+
+    @POST
+    @Path("/error")
+    @Produces(MediaType.APPLICATION_JSON)
+    public Response error()
+    {
+      return Response.status(404).build();
+    }
   }
 }