Allow configuring retry on a per method basis

The classifier is used in preference to the per method retry value,
effectively making it a default, as the classifier can take advantage
of information in the actual exception instance to make a decision.

drift-api/src/main/java/io/airlift/drift/annotations/ThriftException.java

@@ -30,4 +30,11 @@
     Class<? extends Throwable> type();
 
     short id();
+
+    Retryable retryable() default Retryable.UNKNOWN;
+
+    enum Retryable
+    {
+        UNKNOWN, FALSE, TRUE
+    }
 }

drift-api/src/main/java/io/airlift/drift/annotations/ThriftRetryable.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2020 The Drift Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.airlift.drift.annotations;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.TYPE_USE;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+/**
+ * Indicates if an exception thrown by a method is retryable.
+ */
+@Documented
+@Retention(RUNTIME)
+@Target(TYPE_USE)
+public @interface ThriftRetryable
+{
+    boolean value();
+}

drift-client/src/main/java/io/airlift/drift/client/RetryPolicy.java

@@ -17,6 +17,7 @@
 
 import io.airlift.drift.protocol.TTransportException;
 import io.airlift.drift.transport.client.ConnectionFailedException;
+import io.airlift.drift.transport.client.DriftApplicationException;
 import io.airlift.drift.transport.client.DriftClientConfig;
 import io.airlift.drift.transport.client.MessageTooLargeException;
 import io.airlift.drift.transport.client.RequestTimeoutException;
@@ -125,6 +126,14 @@ public ExceptionClassification classifyException(Throwable throwable, boolean id
             return result;
         }
 
+        // use provided result if available
+        if (throwable instanceof DriftApplicationException) {
+            Optional<Boolean> retryable = ((DriftApplicationException) throwable).isRetryable();
+            if (retryable.isPresent()) {
+                return new ExceptionClassification(retryable, NORMAL);
+            }
+        }
+
         if (idempotent && throwable instanceof TTransportException) {
             // We don't know if there is a problem with this server or if this
             // is a general network error, so just mark the server as normal.

drift-client/src/test/java/io/airlift/drift/client/TestDriftMethodInvocation.java

@@ -75,6 +75,7 @@ public class TestDriftMethodInvocation
             ImmutableList.of(),
             (ThriftCodec<Object>) (Object) new ShortThriftCodec(),
             ImmutableMap.of(),
+            ImmutableMap.of(),
             false,
             true);
     private static final Error UNEXPECTED_EXCEPTION = new Error("unexpected exception");
@@ -795,7 +796,7 @@ public static Exception createClassifiedException(boolean retry, HostStatus host
         {
             Exception exception = new ClassifiedException(new ExceptionClassification(Optional.of(retry), hostStatus));
             if (wrapWithApplicationException) {
-                exception = new DriftApplicationException(exception);
+                exception = new DriftApplicationException(exception, Optional.empty());
             }
             return exception;
         }

drift-client/src/test/java/io/airlift/drift/client/TestRetryPolicy.java

@@ -21,8 +21,10 @@
 
 import java.util.Optional;
 
+import static io.airlift.drift.client.ExceptionClassification.HostStatus.NORMAL;
 import static io.airlift.drift.client.ExceptionClassification.HostStatus.OVERLOADED;
 import static io.airlift.drift.client.ExceptionClassification.NORMAL_EXCEPTION;
+import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertSame;
 
 public class TestRetryPolicy
@@ -37,10 +39,37 @@ public void testRetryUserException()
             }
             return NORMAL_EXCEPTION;
         });
-        assertSame(policy.classifyException(new DriftApplicationException(new TestingUserException()), true), overloaded);
+        assertSame(policy.classifyException(new DriftApplicationException(new TestingUserException(), Optional.empty()), true), overloaded);
         assertSame(policy.classifyException(new TestingUserException(), true), overloaded);
     }
 
+    @Test
+    public void testProvidedResult()
+    {
+        // classifier has precedence over provided result
+        assertEquals(classify(Optional.of(true), Optional.of(false)), Optional.of(true));
+        assertEquals(classify(Optional.of(false), Optional.of(true)), Optional.of(false));
+
+        // both set to same
+        assertEquals(classify(Optional.of(true), Optional.of(true)), Optional.of(true));
+        assertEquals(classify(Optional.of(false), Optional.of(false)), Optional.of(false));
+
+        // only one set
+        assertEquals(classify(Optional.empty(), Optional.of(true)), Optional.of(true));
+        assertEquals(classify(Optional.empty(), Optional.of(false)), Optional.of(false));
+        assertEquals(classify(Optional.of(true), Optional.empty()), Optional.of(true));
+        assertEquals(classify(Optional.of(false), Optional.empty()), Optional.of(false));
+
+        // neither set
+        assertEquals(classify(Optional.empty(), Optional.empty()), Optional.empty());
+    }
+
+    private static Optional<Boolean> classify(Optional<Boolean> classifierResult, Optional<Boolean> providedResult)
+    {
+        return new RetryPolicy(new DriftClientConfig(), classifier -> new ExceptionClassification(classifierResult, NORMAL))
+                .classifyException(new DriftApplicationException(new TestingUserException(), providedResult), true).isRetry();
+    }
+
     private static class TestingUserException
             extends Exception
     {}

drift-codec/src/main/java/io/airlift/drift/codec/metadata/ThriftMethodMetadata.java

@@ -27,6 +27,7 @@
 import io.airlift.drift.annotations.ThriftId;
 import io.airlift.drift.annotations.ThriftIdlAnnotation;
 import io.airlift.drift.annotations.ThriftMethod;
+import io.airlift.drift.annotations.ThriftRetryable;
 import io.airlift.drift.annotations.ThriftStruct;
 
 import javax.annotation.concurrent.Immutable;
@@ -62,7 +63,7 @@ public class ThriftMethodMetadata
     private final List<ThriftFieldMetadata> parameters;
     private final Set<ThriftHeaderParameter> headerParameters;
     private final Method method;
-    private final ImmutableMap<Short, ThriftType> exceptions;
+    private final Map<Short, ExceptionInfo> exceptions;
     private final boolean oneway;
     private final boolean idempotent;
     private final List<String> documentation;
@@ -218,7 +219,7 @@ public Set<ThriftHeaderParameter> getHeaderParameters()
         return headerParameters;
     }
 
-    public Map<Short, ThriftType> getExceptions()
+    public Map<Short, ExceptionInfo> getExceptions()
     {
         return exceptions;
     }
@@ -243,19 +244,21 @@ public List<String> getDocumentation()
         return documentation;
     }
 
-    private ImmutableMap<Short, ThriftType> buildExceptionMap(ThriftCatalog catalog, ThriftMethod thriftMethod)
+    private Map<Short, ExceptionInfo> buildExceptionMap(ThriftCatalog catalog, ThriftMethod thriftMethod)
     {
         boolean mixedStyle = (thriftMethod.exception().length > 0) &&
                 stream(method.getAnnotatedExceptionTypes()).anyMatch(type -> type.isAnnotationPresent(ThriftId.class));
         checkArgument(!mixedStyle, "ThriftMethod [%s] uses a mix of @ThriftException and @ThriftId", methodName(method));
 
-        Map<Short, ThriftType> exceptions = new HashMap<>();
+        Map<Short, ExceptionInfo> exceptions = new HashMap<>();
         Set<Type> exceptionTypes = new HashSet<>();
 
         for (ThriftException thriftException : thriftMethod.exception()) {
             checkArgument(!exceptions.containsKey(thriftException.id()), "ThriftMethod [%s] exception list contains multiple values for field ID [%s]", methodName(method), thriftException.id());
             checkArgument(!exceptionTypes.contains(thriftException.type()), "ThriftMethod [%s] exception list contains multiple values for type [%s]", methodName(method), thriftException.type().getSimpleName());
-            exceptions.put(thriftException.id(), catalog.getThriftType(thriftException.type()));
+            exceptions.put(thriftException.id(), new ExceptionInfo(
+                    catalog.getThriftType(thriftException.type()),
+                    retryable(thriftException.retryable())));
             exceptionTypes.add(thriftException.type());
         }
 
@@ -264,10 +267,16 @@ private ImmutableMap<Short, ThriftType> buildExceptionMap(ThriftCatalog catalog,
         for (int i = 0; i < allExceptionClasses.length; i++) {
             Class<?> exception = allExceptionClasses[i];
             ThriftId thriftId = exceptionAnnotations[i].getAnnotation(ThriftId.class);
+            ThriftRetryable thriftRetryable = exceptionAnnotations[i].getAnnotation(ThriftRetryable.class);
+            Optional<Boolean> retryable = Optional.empty();
+            if (thriftRetryable != null) {
+                checkArgument(thriftId != null, "ThriftMethod [%s] exception list contains @ThriftRetryable without @ThriftId", methodName((method)));
+                retryable = Optional.of(thriftRetryable.value());
+            }
             if (thriftId != null) {
                 checkArgument(!exceptions.containsKey(thriftId.value()), "ThriftMethod [%s] exception list contains multiple values for field ID [%s]", methodName(method), thriftId.value());
                 checkArgument(!exceptionTypes.contains(exception), "ThriftMethod [%s] exception list contains multiple values for type [%s]", methodName(method), exception.getSimpleName());
-                exceptions.put(thriftId.value(), catalog.getThriftType(exception));
+                exceptions.put(thriftId.value(), new ExceptionInfo(catalog.getThriftType(exception), retryable));
                 exceptionTypes.add(exception);
             }
         }
@@ -284,7 +293,7 @@ private ImmutableMap<Short, ThriftType> buildExceptionMap(ThriftCatalog catalog,
                 // there is no ordering guarantee for exception types,
                 // so we can only infer the id if there is a single custom exception
                 checkArgument(exceptionClasses.size() == 1, "ThriftMethod [%s] annotation must declare exception mapping when more than one custom exception is thrown", methodName(method));
-                exceptions.put((short) 1, catalog.getThriftType(exceptionClass));
+                exceptions.put((short) 1, new ExceptionInfo(catalog.getThriftType(exceptionClass), Optional.empty()));
             }
         }
 
@@ -327,4 +336,59 @@ private static String methodName(Method method)
     {
         return method.getDeclaringClass().getName() + "." + method.getName();
     }
+
+    private static Optional<Boolean> retryable(ThriftException.Retryable retryable)
+    {
+        switch (retryable) {
+            case UNKNOWN:
+                return Optional.empty();
+            case FALSE:
+                return Optional.of(false);
+            case TRUE:
+                return Optional.of(true);
+        }
+        throw new AssertionError("Unhandled value: " + retryable);
+    }
+
+    public static class ExceptionInfo
+    {
+        private final ThriftType thriftType;
+        private final Optional<Boolean> retryable;
+
+        public ExceptionInfo(ThriftType thriftType, Optional<Boolean> retryable)
+        {
+            this.thriftType = requireNonNull(thriftType, "thriftType is null");
+            this.retryable = requireNonNull(retryable, "retryable is null");
+        }
+
+        public ThriftType getThriftType()
+        {
+            return thriftType;
+        }
+
+        public Optional<Boolean> isRetryable()
+        {
+            return retryable;
+        }
+
+        @Override
+        public boolean equals(Object o)
+        {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            ExceptionInfo that = (ExceptionInfo) o;
+            return thriftType.equals(that.thriftType) &&
+                    retryable.equals(that.retryable);
+        }
+
+        @Override
+        public int hashCode()
+        {
+            return Objects.hash(thriftType, retryable);
+        }
+    }
 }