Differentiate different types of Drift failures

drift-client/src/main/java/com/facebook/drift/client/DriftMethodInvocation.java

@@ -19,6 +19,11 @@
 import com.facebook.airlift.log.Logger;
 import com.facebook.drift.TException;
 import com.facebook.drift.client.address.AddressSelector;
+import com.facebook.drift.client.exceptions.DriftNonRetryableException;
+import com.facebook.drift.client.exceptions.DriftNoHostsAvailableException;
+import com.facebook.drift.client.exceptions.DriftRetriesFailedException;
+import com.facebook.drift.client.exceptions.DriftRetryTimeExceededException;
+import com.facebook.drift.client.exceptions.DriftMaxRetryAttemptsExceededException;
 import com.facebook.drift.client.stats.MethodInvocationStat;
 import com.facebook.drift.protocol.TTransportException;
 import com.facebook.drift.transport.MethodMetadata;
@@ -158,7 +163,12 @@ private synchronized void nextAttempt(boolean noConnectDelay)
 
             Optional<A> address = addressSelector.selectAddress(addressSelectionContext, attemptedAddresses);
             if (!address.isPresent()) {
-                fail("No hosts available");
+                fail(new DriftNoHostsAvailableException(
+                        invocationAttempts,
+                        succinctNanos(ticker.read() - startTime),
+                        failedConnections,
+                        overloadedRejects,
+                        attemptedAddresses));
                 return;
             }
 
@@ -252,15 +262,32 @@ else if (exceptionClassification.getHostStatus() == DOWN || exceptionClassificat
             if (!exceptionClassification.isRetry().orElse(FALSE)) {
                 // always store exception if non-retryable, so it is added to the exception chain
                 lastException = throwable;
-                fail("Non-retryable exception");
+                fail(new DriftNonRetryableException(
+                        invocationAttempts,
+                        succinctNanos(ticker.read() - startTime),
+                        failedConnections,
+                        overloadedRejects,
+                        attemptedAddresses));
                 return;
             }
             if (invocationAttempts > retryPolicy.getMaxRetries()) {
-                fail(format("Max retry attempts (%s) exceeded", retryPolicy.getMaxRetries()));
+                fail(new DriftMaxRetryAttemptsExceededException(
+                        retryPolicy.getMaxRetries(),
+                        invocationAttempts,
+                        succinctNanos(ticker.read() - startTime),
+                        failedConnections,
+                        overloadedRejects,
+                        attemptedAddresses));
                 return;
             }
             if (duration.compareTo(retryPolicy.getMaxRetryTime()) >= 0) {
-                fail(format("Max retry time (%s) exceeded", retryPolicy.getMaxRetryTime()));
+                fail(new DriftRetryTimeExceededException(
+                        retryPolicy.getMaxRetryTime(),
+                        invocationAttempts,
+                        succinctNanos(ticker.read() - startTime),
+                        failedConnections,
+                        overloadedRejects,
+                        attemptedAddresses));
                 return;
             }
 
@@ -323,28 +350,20 @@ private synchronized void onCancel(boolean wasInterrupted)
         }
     }
 
-    private synchronized void fail(String reason)
+    private synchronized void fail(DriftRetriesFailedException e)
     {
         Throwable cause = lastException;
         if (cause == null) {
             // There are no hosts or all hosts are marked down
-            cause = new TTransportException(reason);
+            cause = new TTransportException(e.getMessage());
         }
 
-        RetriesFailedException retriesFailedException = new RetriesFailedException(
-                reason,
-                invocationAttempts,
-                succinctNanos(ticker.read() - startTime),
-                failedConnections,
-                overloadedRejects,
-                attemptedAddresses);
-
         // attach message exception to the exception thrown to caller
         if (cause instanceof DriftApplicationException) {
-            cause.getCause().addSuppressed(retriesFailedException);
+            cause.getCause().addSuppressed(e);
         }
         else {
-            cause.addSuppressed(retriesFailedException);
+            cause.addSuppressed(e);
         }
 
         setException(cause);

drift-client/src/main/java/com/facebook/drift/client/exceptions/DriftMaxRetryAttemptsExceededException.java

@@ -0,0 +1,23 @@
+package com.facebook.drift.client.exceptions;
+
+import com.facebook.drift.transport.client.Address;
+import io.airlift.units.Duration;
+
+import java.util.Set;
+
+import static java.lang.String.format;
+
+public class DriftMaxRetryAttemptsExceededException
+    extends DriftRetriesFailedException
+{
+    public DriftMaxRetryAttemptsExceededException(
+            int maxRetries,
+            int invocationAttempts,
+            Duration retryTime,
+            int failedConnections,
+            int overloadedRejects,
+            Set<? extends Address> attemptedAddresses)
+    {
+        super(format("Max retry attempts (%s) exceeded", maxRetries), invocationAttempts, retryTime, failedConnections, overloadedRejects, attemptedAddresses);
+    }
+}

drift-client/src/main/java/com/facebook/drift/client/exceptions/DriftNoHostsAvailableException.java

@@ -0,0 +1,15 @@
+package com.facebook.drift.client.exceptions;
+
+import com.facebook.drift.transport.client.Address;
+import io.airlift.units.Duration;
+
+import java.util.Set;
+
+public class DriftNoHostsAvailableException
+    extends DriftRetriesFailedException
+{
+    public DriftNoHostsAvailableException(int invocationAttempts, Duration retryTime, int failedConnections, int overloadedRejects, Set<? extends Address> attemptedAddresses)
+    {
+        super("No hosts available", invocationAttempts, retryTime, failedConnections, overloadedRejects, attemptedAddresses);
+    }
+}

drift-client/src/main/java/com/facebook/drift/client/exceptions/DriftNonRetryableException.java

@@ -0,0 +1,15 @@
+package com.facebook.drift.client.exceptions;
+
+import com.facebook.drift.transport.client.Address;
+import io.airlift.units.Duration;
+
+import java.util.Set;
+
+public class DriftNonRetryableException
+    extends DriftRetriesFailedException
+{
+    public DriftNonRetryableException(int invocationAttempts, Duration retryTime, int failedConnections, int overloadedRejects, Set<? extends Address> attemptedAddresses)
+    {
+        super("Non-retryable exception", invocationAttempts, retryTime, failedConnections, overloadedRejects, attemptedAddresses);
+    }
+}

drift-client/src/main/java/com/facebook/drift/client/RetriesFailedException.java → drift-client/src/main/java/com/facebook/drift/client/exceptions/DriftRetriesFailedException.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.drift.client;
+package com.facebook.drift.client.exceptions;
 
 import com.facebook.drift.transport.client.Address;
 import com.google.common.collect.ImmutableSet;
@@ -24,7 +24,7 @@
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
-public class RetriesFailedException
+public class DriftRetriesFailedException
         extends Exception
 {
     private final int invocationAttempts;
@@ -33,7 +33,7 @@ public class RetriesFailedException
     private final int overloadedRejects;
     private final Set<? extends Address> attemptedAddresses;
 
-    public RetriesFailedException(
+    DriftRetriesFailedException(
             String reason,
             int invocationAttempts,
             Duration retryTime,

drift-client/src/main/java/com/facebook/drift/client/exceptions/DriftRetryTimeExceededException.java

@@ -0,0 +1,23 @@
+package com.facebook.drift.client.exceptions;
+
+import com.facebook.drift.transport.client.Address;
+import io.airlift.units.Duration;
+
+import java.util.Set;
+
+import static java.lang.String.format;
+
+public class DriftRetryTimeExceededException
+    extends DriftRetriesFailedException
+{
+    public DriftRetryTimeExceededException(
+            Duration maxRetryTime,
+            int invocationAttempts,
+            Duration retryTime,
+            int failedConnections,
+            int overloadedRejects,
+            Set<? extends Address> attemptedAddresses)
+    {
+        super(format("Max retry time (%s) exceeded", maxRetryTime), invocationAttempts, retryTime, failedConnections, overloadedRejects, attemptedAddresses);
+    }
+}

drift-client/src/test/java/com/facebook/drift/client/TestDriftMethodInvocation.java

@@ -21,6 +21,11 @@
 import com.facebook.drift.client.ExceptionClassification.HostStatus;
 import com.facebook.drift.client.address.AddressSelector;
 import com.facebook.drift.client.address.SimpleAddressSelector.SimpleAddress;
+import com.facebook.drift.client.exceptions.DriftMaxRetryAttemptsExceededException;
+import com.facebook.drift.client.exceptions.DriftNoHostsAvailableException;
+import com.facebook.drift.client.exceptions.DriftNonRetryableException;
+import com.facebook.drift.client.exceptions.DriftRetryTimeExceededException;
+import com.facebook.drift.client.exceptions.DriftRetriesFailedException;
 import com.facebook.drift.codec.ThriftCodec;
 import com.facebook.drift.codec.internal.builtin.ShortThriftCodec;
 import com.facebook.drift.protocol.TTransportException;
@@ -210,7 +215,7 @@ private static void testBasicRetriesToFailure(int expectedRetries, boolean wrapW
         }
         catch (ExecutionException e) {
             assertEquals(attempts.get(), expectedRetries + 1);
-            assertClassifiedException(e.getCause(), new ExceptionClassification(Optional.of(false), NORMAL), expectedRetries);
+            assertClassifiedException(e.getCause(), DriftNonRetryableException.class, new ExceptionClassification(Optional.of(false), NORMAL), expectedRetries);
         }
         stat.assertFailure(expectedRetries);
         assertDelays(invoker, retryPolicy, expectedRetries);
@@ -241,7 +246,7 @@ public void testBasicRetriesToNoHosts()
         }
         catch (ExecutionException e) {
             assertEquals(attempts.get(), expectedRetries + 1);
-            assertClassifiedException(e.getCause(), new ExceptionClassification(Optional.of(true), NORMAL), expectedRetries);
+            assertClassifiedException(e.getCause(), DriftNoHostsAvailableException.class, new ExceptionClassification(Optional.of(true), NORMAL), expectedRetries);
         }
         stat.assertFailure(expectedRetries);
     }
@@ -266,7 +271,7 @@ public void testMaxRetries()
         }
         catch (ExecutionException e) {
             assertEquals(attempts.get(), maxRetries + 1);
-            assertClassifiedException(e.getCause(), new ExceptionClassification(Optional.of(true), NORMAL), maxRetries);
+            assertClassifiedException(e.getCause(), DriftMaxRetryAttemptsExceededException.class, new ExceptionClassification(Optional.of(true), NORMAL), maxRetries);
         }
         stat.assertFailure(maxRetries);
     }
@@ -302,7 +307,7 @@ public void testMaxRetryTime()
         }
         catch (ExecutionException e) {
             assertEquals(attempts.get(), maxRetries + 1);
-            assertClassifiedException(e.getCause(), new ExceptionClassification(Optional.of(true), NORMAL), maxRetries);
+            assertClassifiedException(e.getCause(), DriftRetryTimeExceededException.class, new ExceptionClassification(Optional.of(true), NORMAL), maxRetries);
         }
         stat.assertFailure(maxRetries);
         assertDelays(invoker, retryPolicy, 7);
@@ -350,7 +355,7 @@ private static void testExhaustHosts(int expectedRetries, boolean overloaded)
             assertTrue(e.getCause() instanceof TTransportException);
             TTransportException transportException = (TTransportException) e.getCause();
             assertTrue(transportException.getMessage().startsWith("No hosts available"));
-            assertRetriesFailedInformation(transportException, 0, 0, overloaded ? expectedRetries : 0);
+            assertRetriesFailedInformation(transportException, DriftNoHostsAvailableException.class, 0, 0, overloaded ? expectedRetries : 0);
         }
         stat.assertNoHostsAvailable(expectedRetries);
         addressSelector.assertAllDown();
@@ -398,7 +403,7 @@ private static void testConnectionFailed(int expectedInvocationAttempts, int fai
             DriftApplicationException applicationException = (DriftApplicationException) e.getCause();
             assertTrue(applicationException.getCause() instanceof ClassifiedException);
             ClassifiedException classifiedException = (ClassifiedException) applicationException.getCause();
-            assertRetriesFailedInformation(classifiedException, failedConnections, expectedInvocationAttempts, 0);
+            assertRetriesFailedInformation(classifiedException, DriftNonRetryableException.class, failedConnections, expectedInvocationAttempts, 0);
         }
     }
 
@@ -750,32 +755,34 @@ private static DriftMethodInvocation<?> createDriftMethodInvocation(
                 retryService);
     }
 
-    private static void assertClassifiedException(Throwable cause, ExceptionClassification exceptionClassification, int expectedRetries)
+    private static <E extends DriftRetriesFailedException> void assertClassifiedException(Throwable cause, Class<E> rootCauseType, ExceptionClassification exceptionClassification, int expectedRetries)
     {
         if (cause instanceof DriftApplicationException) {
             cause = cause.getCause();
         }
         assertTrue(cause instanceof ClassifiedException);
         ClassifiedException classifiedException = (ClassifiedException) cause;
         assertEquals(classifiedException.getClassification(), exceptionClassification);
-        assertRetriesFailedInformation(classifiedException, 0, expectedRetries + 1, 0);
+        assertRetriesFailedInformation(classifiedException, rootCauseType, 0, expectedRetries + 1, 0);
     }
 
-    private static void assertRetriesFailedInformation(Throwable exception, int expectedFailedConnections, int expectedInvocationAttempts, int expectedOverloaded)
+    private static <E extends DriftRetriesFailedException> void assertRetriesFailedInformation(Throwable exception, Class<E> rootCauseType, int expectedFailedConnections, int expectedInvocationAttempts, int expectedOverloaded)
     {
-        RetriesFailedException retriesFailedException = getRetriesFailedException(exception);
+        DriftRetriesFailedException retriesFailedException = getRetriesFailedException(exception, rootCauseType);
         assertEquals(retriesFailedException.getFailedConnections(), expectedFailedConnections);
         assertEquals(retriesFailedException.getInvocationAttempts(), expectedInvocationAttempts);
         assertEquals(retriesFailedException.getOverloadedRejects(), expectedOverloaded);
     }
 
-    private static RetriesFailedException getRetriesFailedException(Throwable exception)
+    private static <E extends DriftRetriesFailedException> DriftRetriesFailedException getRetriesFailedException(Throwable exception, Class<E> rootCauseType)
     {
         // method invocation attaches retry information using a suppressed exception
         Throwable[] suppressed = exception.getSuppressed();
         assertEquals(suppressed.length, 1);
-        assertTrue(suppressed[0] instanceof RetriesFailedException);
-        return (RetriesFailedException) suppressed[0];
+        assertTrue(suppressed[0] instanceof DriftRetriesFailedException);
+        assertEquals(suppressed[0].getClass(), rootCauseType);
+        assertTrue(rootCauseType.isAssignableFrom(suppressed[0].getClass()));
+        return (DriftRetriesFailedException) suppressed[0];
     }
 
     private static void assertUnexpectedException(Throwable cause)

drift-integration-tests/src/test/java/com/facebook/drift/integration/guice/TestGuiceIntegration.java

@@ -20,7 +20,7 @@
 import com.facebook.drift.TApplicationException;
 import com.facebook.drift.TException;
 import com.facebook.drift.client.ExceptionClassification;
-import com.facebook.drift.client.RetriesFailedException;
+import com.facebook.drift.client.exceptions.DriftRetriesFailedException;
 import com.facebook.drift.integration.guice.EchoService.EmptyOptionalException;
 import com.facebook.drift.integration.guice.EchoService.NullValueException;
 import com.facebook.drift.integration.scribe.drift.DriftLogEntry;
@@ -250,7 +250,7 @@ private static void assertThrowingService(ThrowingService service, boolean singl
             assertFalse(e.isRetryable());
             assertEquals(e.getSuppressed().length, 1);
             Throwable t = e.getSuppressed()[0];
-            assertThat(t).isInstanceOf(RetriesFailedException.class)
+            assertThat(t).isInstanceOf(DriftRetriesFailedException.class)
                     .hasMessageContaining("Non-retryable exception")
                     .hasMessageContaining("invocationAttempts: 1,");
         }
@@ -264,7 +264,7 @@ private static void assertThrowingService(ThrowingService service, boolean singl
             assertTrue(e.isRetryable());
             assertEquals(e.getSuppressed().length, 1);
             Throwable t = e.getSuppressed()[0];
-            assertThat(t).isInstanceOf(RetriesFailedException.class)
+            assertThat(t).isInstanceOf(DriftRetriesFailedException.class)
                     .hasMessageContaining("Max retry attempts (5) exceeded")
                     .hasMessageContaining("invocationAttempts: 6,");
         }
@@ -281,7 +281,7 @@ private static void receiveTooLargeMessage(ThrowingService service)
                     .hasMessageMatching("Frame size .+ exceeded max size .+");
             assertEquals(e.getSuppressed().length, 1);
             Throwable t = e.getSuppressed()[0];
-            assertThat(t).isInstanceOf(RetriesFailedException.class)
+            assertThat(t).isInstanceOf(DriftRetriesFailedException.class)
                     .hasMessageContaining("Non-retryable exception")
                     .hasMessageContaining("invocationAttempts: 1,");
         }