Add exception handler for annotated http service (line#847)

Motivation:
Currently, there is no way to send a customized response when an exception is raised from an annotated http service method. We only have an HttpResponseException which sends an http status code without http body.

Modifications:
- Add `@ExceptionHandler` annotation in order to specify exception handler on annotated method.
- Add ExceptionHandlerFunction to handle exception.
- Rename HttpResponseException to HttpStatusException.
- Add HttpResponseException which holds an HttpResponse as its member.
- Rename `@Decorate` to `@Decorator` and make it repeatable.
- Remove ResourceNotFoundException and ServiceUnavailableException. Use HttpStatusException instead.

Result:
We can specify exception handlers on annotated method.

core/src/main/java/com/linecorp/armeria/client/encoding/HttpDecodedResponse.java

@@ -91,10 +91,10 @@ protected void beforeComplete(Subscriber<? super HttpObject> subscriber) {
     }
 
     @Override
-    protected void beforeError(Subscriber<? super HttpObject> subscriber, Throwable cause) {
-        if (responseDecoder == null) {
-            return;
+    protected Throwable beforeError(Subscriber<? super HttpObject> subscriber, Throwable cause) {
+        if (responseDecoder != null) {
+            responseDecoder.finish();
         }
-        responseDecoder.finish();
+        return cause;
     }
 }

core/src/main/java/com/linecorp/armeria/common/HttpMessageAggregator.java

@@ -116,8 +116,11 @@ public void accept(Void unused, Throwable cause) {
             content = HttpData.of(merged);
         }
 
-        final AggregatedHttpMessage aggregated = onSuccess(content);
-        future.complete(aggregated);
+        try {
+            future.complete(onSuccess(content));
+        } catch (Throwable e) {
+            future.completeExceptionally(e);
+        }
     }
 
     private void fail(Throwable cause) {

core/src/main/java/com/linecorp/armeria/common/stream/FilteredStreamMessage.java

@@ -23,6 +23,8 @@
 
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * A {@link StreamMessage} that filters objects as they are published. The filtering
@@ -31,6 +33,8 @@
  */
 public abstract class FilteredStreamMessage<T, U> implements StreamMessage<U> {
 
+    private static final Logger logger = LoggerFactory.getLogger(FilteredStreamMessage.class);
+
     private final StreamMessage<T> delegate;
 
     /**
@@ -64,9 +68,12 @@ protected void beforeComplete(Subscriber<? super U> subscriber) {}
     /**
      * A callback executed just before calling {@link Subscriber#onError(Throwable)} on {@code subscriber}.
      * Override this method to execute any cleanup logic that may be needed before failing the
-     * subscription.
+     * subscription. This method may rewrite the {@code cause} and then return a new one so that the new
+     * {@link Throwable} would be passed to {@link Subscriber#onError(Throwable)}.
      */
-    protected void beforeError(Subscriber<? super U> subscriber, Throwable cause) {}
+    protected Throwable beforeError(Subscriber<? super U> subscriber, Throwable cause) {
+        return cause;
+    }
 
     @Override
     public boolean isOpen() {
@@ -136,8 +143,14 @@ public void onNext(T t) {
 
         @Override
         public void onError(Throwable t) {
-            beforeError(delegate, t);
-            delegate.onError(t);
+            final Throwable filteredCause = beforeError(delegate, t);
+            if (filteredCause != null) {
+                delegate.onError(filteredCause);
+            } else {
+                logger.warn("{}#beforeError() returned null. Using the original exception:",
+                            FilteredStreamMessage.this.getClass().getName(), t.toString());
+                delegate.onError(t);
+            }
         }
 
         @Override

core/src/main/java/com/linecorp/armeria/server/AnnotatedHttpService.java

@@ -23,12 +23,9 @@
 import java.util.function.Function;
 
 import com.google.common.base.MoreObjects;
-import com.google.common.base.Throwables;
 
 import com.linecorp.armeria.common.HttpRequest;
 import com.linecorp.armeria.common.HttpResponse;
-import com.linecorp.armeria.common.HttpStatus;
-import com.linecorp.armeria.common.util.Exceptions;
 
 /**
  * {@link PathMapping} and their corresponding {@link BiFunction}.
@@ -43,7 +40,8 @@ final class AnnotatedHttpService implements HttpService {
     /**
      * The {@link BiFunction} that will handle the request actually.
      */
-    private final BiFunction<ServiceRequestContext, HttpRequest, Object> function;
+    private final BiFunction<ServiceRequestContext, HttpRequest,
+            CompletionStage<HttpResponse>> function;
 
     /**
      * A decorator of this service.
@@ -55,7 +53,8 @@ final class AnnotatedHttpService implements HttpService {
      * Creates a new instance.
      */
     AnnotatedHttpService(HttpHeaderPathMapping pathMapping,
-                         BiFunction<ServiceRequestContext, HttpRequest, Object> function,
+                         BiFunction<ServiceRequestContext, HttpRequest,
+                                 CompletionStage<HttpResponse>> function,
                          Function<Service<HttpRequest, HttpResponse>,
                                  ? extends Service<HttpRequest, HttpResponse>> decorator) {
         this.pathMapping = requireNonNull(pathMapping, "pathMapping");
@@ -90,36 +89,15 @@ boolean overlaps(AnnotatedHttpService entry) {
 
     @Override
     public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
-
-        final Object ret;
-        try {
-            ret = function.apply(ctx, req);
-        } catch (IllegalArgumentException ignore) {
-            throw new HttpResponseException(HttpStatus.BAD_REQUEST);
-        }
-
-        if (!(ret instanceof CompletionStage)) {
-            return HttpResponse.ofFailure(new IllegalStateException(
-                    "illegal return type: " + ret.getClass().getSimpleName()));
-        }
-
-        @SuppressWarnings("unchecked")
-        CompletionStage<HttpResponse> castStage = (CompletionStage<HttpResponse>) ret;
-        return HttpResponse.from(castStage.handle((httpResponse, throwable) -> {
-            if (throwable != null) {
-                if (Throwables.getRootCause(throwable) instanceof IllegalArgumentException) {
-                    return HttpResponse.of(HttpStatus.BAD_REQUEST);
-                }
-                return Exceptions.throwUnsafely(throwable);
-            }
-            return httpResponse;
-        }));
+        return HttpResponse.from(function.apply(ctx, req));
     }
 
     @Override
     public String toString() {
         return MoreObjects.toStringHelper(this)
                           .add("pathMapping", pathMapping)
-                          .add("function", function).toString();
+                          .add("function", function)
+                          .add("decorator", decorator)
+                          .toString();
     }
 }

core/src/main/java/com/linecorp/armeria/server/AnnotatedHttpServiceMethod.java

@@ -21,7 +21,6 @@
 import static com.linecorp.armeria.internal.DefaultValues.getSpecifiedValue;
 import static java.util.Objects.requireNonNull;
 
-import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.ParameterizedType;
 import java.util.List;
@@ -34,13 +33,17 @@
 
 import javax.annotation.Nullable;
 
+import org.reactivestreams.Subscriber;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 
 import com.linecorp.armeria.common.AggregatedHttpMessage;
+import com.linecorp.armeria.common.FilteredHttpResponse;
 import com.linecorp.armeria.common.HttpHeaderNames;
+import com.linecorp.armeria.common.HttpObject;
 import com.linecorp.armeria.common.HttpParameters;
 import com.linecorp.armeria.common.HttpRequest;
 import com.linecorp.armeria.common.HttpResponse;
@@ -50,6 +53,7 @@
 import com.linecorp.armeria.common.util.SafeCloseable;
 import com.linecorp.armeria.internal.Types;
 import com.linecorp.armeria.server.annotation.Default;
+import com.linecorp.armeria.server.annotation.ExceptionHandlerFunction;
 import com.linecorp.armeria.server.annotation.Header;
 import com.linecorp.armeria.server.annotation.Param;
 import com.linecorp.armeria.server.annotation.ResponseConverter;
@@ -71,11 +75,15 @@ final class AnnotatedHttpServiceMethod implements BiFunction<ServiceRequestConte
     private final List<Parameter> parameters;
     private final boolean isAsynchronous;
     private final AggregationStrategy aggregationStrategy;
+    private final List<ExceptionHandlerFunction> exceptionHandlers;
 
-    AnnotatedHttpServiceMethod(Object object, Method method, PathMapping pathMapping) {
+    AnnotatedHttpServiceMethod(Object object, Method method, PathMapping pathMapping,
+                               List<ExceptionHandlerFunction> exceptionHandlers) {
         this.object = requireNonNull(object, "object");
         this.method = requireNonNull(method, "method");
         requireNonNull(pathMapping, "pathMapping");
+        this.exceptionHandlers = ImmutableList.copyOf(
+                requireNonNull(exceptionHandlers, "exceptionHandlers"));
 
         parameters = parameters(method, pathMapping.paramNames());
         final Class<?> returnType = method.getReturnType();
@@ -118,7 +126,8 @@ public Object apply(ServiceRequestContext ctx, HttpRequest req) {
         return CompletableFuture.supplyAsync(() -> invoke(ctx, req, null), ctx.blockingTaskExecutor());
     }
 
-    BiFunction<ServiceRequestContext, HttpRequest, Object> withConverter(ResponseConverter converter) {
+    BiFunction<ServiceRequestContext, HttpRequest,
+            CompletionStage<HttpResponse>> withConverter(ResponseConverter converter) {
         return (ctx, req) ->
                 executeSyncOrAsync(ctx, req).thenApply(obj -> {
                     try (SafeCloseable ignored = RequestContext.push(ctx, false)) {
@@ -127,9 +136,8 @@ BiFunction<ServiceRequestContext, HttpRequest, Object> withConverter(ResponseCon
                 });
     }
 
-    BiFunction<ServiceRequestContext, HttpRequest, Object> withConverters(
-            Map<Class<?>, ResponseConverter> converters) {
-
+    BiFunction<ServiceRequestContext, HttpRequest,
+            CompletionStage<HttpResponse>> withConverters(Map<Class<?>, ResponseConverter> converters) {
         return (ctx, req) ->
                 executeSyncOrAsync(ctx, req).thenApply(obj -> {
                     try (SafeCloseable ignored = RequestContext.push(ctx, false)) {
@@ -140,22 +148,55 @@ BiFunction<ServiceRequestContext, HttpRequest, Object> withConverters(
 
     private CompletionStage<?> executeSyncOrAsync(ServiceRequestContext ctx, HttpRequest req) {
         final Object ret = apply(ctx, req);
-        return ret instanceof CompletionStage ?
-               (CompletionStage<?>) ret : CompletableFuture.completedFuture(ret);
+        if (ret instanceof CompletionStage) {
+            return ((CompletionStage<?>) ret).handle((obj, cause) -> {
+                if (cause == null) {
+                    return obj;
+                }
+                final HttpResponse response = getResponseForCause(ctx, req, cause);
+                if (response != null) {
+                    return response;
+                } else {
+                    return Exceptions.throwUnsafely(cause);
+                }
+            });
+        }
+        if (ret instanceof HttpResponse) {
+            return CompletableFuture.completedFuture(
+                    new ExceptionFilteredHttpResponse(ctx, req, (HttpResponse) ret));
+        }
+        return CompletableFuture.completedFuture(ret);
     }
 
     private Object invoke(ServiceRequestContext ctx, HttpRequest req, @Nullable AggregatedHttpMessage message) {
         try (SafeCloseable ignored = RequestContext.push(ctx, false)) {
             return method.invoke(object, parameterValues(ctx, req, parameters, message));
-        } catch (Exception e) {
-            if (e instanceof InvocationTargetException) {
-                final Throwable cause = e.getCause();
-                if (cause != null) {
-                    return Exceptions.throwUnsafely(cause);
+        } catch (Throwable cause) {
+            final HttpResponse response = getResponseForCause(ctx, req, cause);
+            if (response != null) {
+                return response;
+            }
+            return Exceptions.throwUnsafely(cause);
+        }
+    }
+
+    /**
+     * Returns a {@link HttpResponse} which is created by {@link ExceptionHandlerFunction}.
+     */
+    private HttpResponse getResponseForCause(ServiceRequestContext ctx, HttpRequest req,
+                                             Throwable cause) {
+        final Throwable rootCause = Throwables.getRootCause(cause);
+        for (ExceptionHandlerFunction func : exceptionHandlers) {
+            if (func.accept(rootCause)) {
+                try {
+                    return func.handle(ctx, req, rootCause);
+                } catch (Exception e) {
+                    logger.warn("Unexpected exception from an exception handler {}:",
+                                func.getClass().getName(), e);
                 }
             }
-            return Exceptions.throwUnsafely(e);
         }
+        return ExceptionHandlerFunction.DEFAULT.handle(ctx, req, rootCause);
     }
 
     /**
@@ -229,7 +270,7 @@ private static Parameter createOptionalSupportedParam(java.lang.reflect.Paramete
             type = parameterInfo.getType();
         }
 
-        return new Parameter(paramType, (!isOptionalType && aDefault == null), isOptionalType,
+        return new Parameter(paramType, !isOptionalType && aDefault == null, isOptionalType,
                              validateAndNormalizeSupportedType(type), paramValue, defaultValue);
     }
 
@@ -509,6 +550,36 @@ private static ResponseConverter findResponseConverter(
         throw new IllegalArgumentException("Converter not available for: " + type.getSimpleName());
     }
 
+    /**
+     * Intercepts a {@link Throwable} raised from {@link HttpResponse} and then rewrites it as an
+     * {@link HttpResponseException} by {@link ExceptionHandlerFunction}.
+     */
+    private class ExceptionFilteredHttpResponse extends FilteredHttpResponse {
+
+        private final ServiceRequestContext ctx;
+        private final HttpRequest req;
+
+        ExceptionFilteredHttpResponse(ServiceRequestContext ctx, HttpRequest req,
+                                      HttpResponse delegate) {
+            super(delegate);
+            this.ctx = ctx;
+            this.req = req;
+        }
+
+        @Override
+        protected HttpObject filter(HttpObject obj) {
+            return obj;
+        }
+
+        @Override
+        protected Throwable beforeError(Subscriber<? super HttpObject> subscriber,
+                                        Throwable cause) {
+            final HttpResponse response = getResponseForCause(ctx, req, cause);
+            return response != null ? HttpResponseException.of(response)
+                                    : cause;
+        }
+    }
+
     /**
      * Parameter entry, which will be used to invoke the {@link AnnotatedHttpService}.
      */