Fix retry logic in BrokerClient and flakey BrokerClientTest

server/src/main/java/org/apache/druid/discovery/BrokerClient.java

@@ -33,7 +33,6 @@
 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,17 +87,16 @@ public String sendQuery(final Request request) throws Exception
             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();
         },
         (throwable) -> {
           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

server/src/test/java/org/apache/druid/discovery/BrokerClientTest.java

@@ -101,7 +101,7 @@ public void testSimple() throws Exception
   }
 
   @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 void testError() throws Exception
     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 Response redirecting()
         return Response.status(504).build();
       }
     }
+
+    @POST
+    @Path("/error")
+    @Produces(MediaType.APPLICATION_JSON)
+    public Response error()
+    {
+      return Response.status(404).build();
+    }
   }
 }