Implement @Attribute Injection.

core/src/main/java/com/linecorp/armeria/internal/server/annotation/AnnotatedElementNameUtil.java

@@ -21,52 +21,29 @@
 import java.lang.reflect.Field;
 import java.lang.reflect.Parameter;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Ascii;
 import com.google.common.base.CaseFormat;
 
+import com.linecorp.armeria.server.annotation.Attribute;
 import com.linecorp.armeria.server.annotation.Header;
 import com.linecorp.armeria.server.annotation.Param;
 
 final class AnnotatedElementNameUtil {
 
     /**
-     * Returns the value of the {@link Param} annotation which is specified on the {@code element} if
-     * the value is not blank. If the value is blank, it returns the name of the specified
-     * {@code nameRetrievalTarget} object which is an instance of {@link Parameter} or {@link Field}.
+     * Returns the value of {@link Header}, {@link Param}, {@link Attribute} if the value is not blank.
+     * If the value is blank, it returns the name of the specified {@code nameRetrievalTarget} object
+     * which is an instance of {@link Header}, {@link Param}, {@link Attribute} or {@link Field}.
      */
-    static String findName(Param param, Object nameRetrievalTarget) {
+    static String findName(Object nameRetrievalTarget, String value) {
         requireNonNull(nameRetrievalTarget, "nameRetrievalTarget");
-
-        final String value = param.value();
         if (DefaultValues.isSpecified(value)) {
             checkArgument(!value.isEmpty(), "value is empty.");
             return value;
         }
         return getName(nameRetrievalTarget);
     }
 
-    /**
-     * Returns the value of the {@link Header} annotation which is specified on the {@code element} if
-     * the value is not blank. If the value is blank, it returns the name of the specified
-     * {@code nameRetrievalTarget} object which is an instance of {@link Parameter} or {@link Field}.
-     *
-     * <p>Note that the name of the specified {@code nameRetrievalTarget} will be converted as
-     * {@link CaseFormat#LOWER_HYPHEN} that the string elements are separated with one hyphen({@code -})
-     * character. The value of the {@link Header} annotation will not be converted because it is clearly
-     * specified by a user.
-     */
-    static String findName(Header header, Object nameRetrievalTarget) {
-        requireNonNull(nameRetrievalTarget, "nameRetrievalTarget");
-
-        final String value = header.value();
-        if (DefaultValues.isSpecified(value)) {
-            checkArgument(!value.isEmpty(), "value is empty.");
-            return value;
-        }
-        return toHeaderName(getName(nameRetrievalTarget));
-    }
-
     /**
      * Returns the name of the specified element or the default name if it can't get.
      */
@@ -89,8 +66,7 @@ static String getName(Object element) {
                 throw new IllegalArgumentException(
                         "cannot obtain the name of the parameter or field automatically. " +
                         "Please make sure you compiled your code with '-parameters' option. " +
-                        "If not, you need to specify parameter and header names with @" +
-                        Param.class.getSimpleName() + " and @" + Header.class.getSimpleName() + '.');
+                        "Alternatively, you could specify the name explicitly in the annotation.");
             }
             return parameter.getName();
         }
@@ -100,7 +76,6 @@ static String getName(Object element) {
         throw new IllegalArgumentException("cannot find the name: " + element.getClass().getName());
     }
 
-    @VisibleForTesting
     static String toHeaderName(String name) {
         requireNonNull(name, "name");
         checkArgument(!name.isEmpty(), "name is empty.");

core/src/main/java/com/linecorp/armeria/internal/server/annotation/AnnotatedValueResolver.java

@@ -22,6 +22,7 @@
 import static com.linecorp.armeria.internal.server.annotation.AnnotatedElementNameUtil.findName;
 import static com.linecorp.armeria.internal.server.annotation.AnnotatedElementNameUtil.getName;
 import static com.linecorp.armeria.internal.server.annotation.AnnotatedElementNameUtil.getNameOrDefault;
+import static com.linecorp.armeria.internal.server.annotation.AnnotatedElementNameUtil.toHeaderName;
 import static com.linecorp.armeria.internal.server.annotation.AnnotatedServiceFactory.findDescription;
 import static com.linecorp.armeria.internal.server.annotation.AnnotatedServiceTypeUtil.stringToType;
 import static com.linecorp.armeria.internal.server.annotation.DefaultValues.getSpecifiedValue;
@@ -64,6 +65,7 @@
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.common.primitives.Primitives;
 
 import com.linecorp.armeria.common.AggregatedHttpRequest;
 import com.linecorp.armeria.common.Cookie;
@@ -84,6 +86,7 @@
 import com.linecorp.armeria.internal.server.FileAggregatedMultipart;
 import com.linecorp.armeria.internal.server.annotation.AnnotatedBeanFactoryRegistry.BeanFactoryId;
 import com.linecorp.armeria.server.ServiceRequestContext;
+import com.linecorp.armeria.server.annotation.Attribute;
 import com.linecorp.armeria.server.annotation.ByteArrayRequestConverterFunction;
 import com.linecorp.armeria.server.annotation.Default;
 import com.linecorp.armeria.server.annotation.Delimiter;
@@ -99,6 +102,7 @@
 import com.linecorp.armeria.server.docs.DescriptionInfo;
 
 import io.netty.handler.codec.http.HttpConstants;
+import io.netty.util.AttributeKey;
 import scala.concurrent.ExecutionContext;
 
 final class AnnotatedValueResolver {
@@ -444,9 +448,16 @@ private static AnnotatedValueResolver of(AnnotatedElement annotatedElement,
         requireNonNull(dependencyInjector, "dependencyInjector");
 
         final DescriptionInfo description = findDescription(annotatedElement);
+
+        final Attribute attr = annotatedElement.getAnnotation(Attribute.class);
+        if (attr != null) {
+            final String name = findName(typeElement, attr.value());
+            return ofAttribute(name, attr, annotatedElement, typeElement, type, description);
+        }
+
         final Param param = annotatedElement.getAnnotation(Param.class);
         if (param != null) {
-            final String name = findName(param, typeElement);
+            final String name = findName(typeElement, param.value());
             if (type == File.class || type == Path.class || type == MultipartFile.class) {
                 return ofFileParam(name, annotatedElement, typeElement, type, description);
             }
@@ -459,7 +470,7 @@ private static AnnotatedValueResolver of(AnnotatedElement annotatedElement,
 
         final Header header = annotatedElement.getAnnotation(Header.class);
         if (header != null) {
-            final String name = findName(header, typeElement);
+            final String name = toHeaderName(findName(typeElement, header.value()));
             return ofHeader(name, annotatedElement, typeElement, type, description);
         }
 
@@ -520,6 +531,7 @@ static List<RequestObjectResolver> addToFirstIfExists(List<RequestObjectResolver
 
     private static boolean isAnnotationPresent(AnnotatedElement element) {
         return element.isAnnotationPresent(Param.class) ||
+               element.isAnnotationPresent(Attribute.class) ||
                element.isAnnotationPresent(Header.class) ||
                element.isAnnotationPresent(RequestObject.class);
     }
@@ -621,6 +633,34 @@ private static AnnotatedValueResolver ofRequestObject(String name, AnnotatedElem
                 .build();
     }
 
+    private static AnnotatedValueResolver ofAttribute(String name,
+                                                      Attribute attr,
+                                                      AnnotatedElement annotatedElement,
+                                                      AnnotatedElement typeElement, Class<?> type,
+                                                      DescriptionInfo description) {
+
+        final ImmutableList.Builder<AttributeKey<?>> builder = ImmutableList.builder();
+
+        if (attr.prefix() != Attribute.class) {
+            builder.add(AttributeKey.valueOf(attr.prefix(), name));
+        } else {
+            final Class<?> serviceClass = ((Parameter) annotatedElement).getDeclaringExecutable()
+                                                                        .getDeclaringClass();
+            builder.add(AttributeKey.valueOf(serviceClass, name));
+            builder.add(AttributeKey.valueOf(name));
+        }
+
+        final ImmutableList<AttributeKey<?>> attrKeys = builder.build();
+        return new Builder(annotatedElement, type, name)
+                .annotationType(Attribute.class)
+                .typeElement(typeElement)
+                .supportDefault(true)
+                .supportContainer(true)
+                .description(description)
+                .resolver(attributeResolver(attrKeys))
+                .build();
+    }
+
     @Nullable
     private static AnnotatedValueResolver ofInjectableTypes(String name, AnnotatedElement annotatedElement,
                                                             Class<?> type, boolean useBlockingExecutor) {
@@ -818,6 +858,45 @@ private static AnnotatedValueResolver ofInjectableTypes0(String name, AnnotatedE
         };
     }
 
+    /**
+     * Returns an attribute resolver which retrieves a value specified by {@code attrKeys}
+     * from the {@link RequestContext}.
+     */
+    private static BiFunction<AnnotatedValueResolver, ResolverContext, Object>
+    attributeResolver(Iterable<AttributeKey<?>> attrKeys) {
+        return (resolver, ctx) -> {
+            Object lastInvalidValue = null;
+            String lastAttrKeyName = null;
+            Class<?> targetType = resolver.rawType();
+            if (targetType.isPrimitive()) {
+                targetType = Primitives.wrap(targetType);
+            }
+
+            for (AttributeKey<?> attrKey : attrKeys) {
+                final Object value = ctx.context.attr(attrKey);
+                if (value != null) {
+                    final boolean isValidType = targetType.isInstance(value);
+                    if (isValidType) {
+                        return value;
+                    } else {
+                        lastInvalidValue = value;
+                        lastAttrKeyName = attrKey.name();
+                    }
+                }
+            }
+
+            if (lastInvalidValue != null && lastAttrKeyName != null) {
+                throw new IllegalStateException(
+                        String.format("'%s' which is from AttributeKey(%s) is not an instance of '%s'.",
+                                      targetType.getName(),
+                                      lastAttrKeyName,
+                                      lastInvalidValue.getClass().getName()));
+            } else {
+                return resolver.defaultOrException();
+            }
+        };
+    }
+
     private static BiFunction<AnnotatedValueResolver, ResolverContext, Object> fileResolver() {
         return (resolver, ctx) -> {
             final FileAggregatedMultipart fileAggregatedMultipart = ctx.aggregatedMultipart();
@@ -896,6 +975,7 @@ static boolean isAnnotatedNullable(AnnotatedElement annotatedElement) {
 
     @Nullable
     private final Class<?> containerType;
+    private final Class<?> rawType;
     private final Class<?> elementType;
     @Nullable
     private final ParameterizedType parameterizedElementType;
@@ -920,6 +1000,7 @@ private AnnotatedValueResolver(@Nullable Class<? extends Annotation> annotationT
                                    boolean isPathVariable, boolean shouldExist,
                                    boolean shouldWrapValueAsOptional,
                                    @Nullable Class<?> containerType, Class<?> elementType,
+                                   Class<?> rawType,
                                    @Nullable ParameterizedType parameterizedElementType,
                                    @Nullable String defaultValue,
                                    DescriptionInfo description,
@@ -935,6 +1016,7 @@ private AnnotatedValueResolver(@Nullable Class<? extends Annotation> annotationT
         this.parameterizedElementType = parameterizedElementType;
         this.description = requireNonNull(description, "description");
         this.containerType = containerType;
+        this.rawType = rawType;
         this.resolver = requireNonNull(resolver, "resolver");
         this.beanFactoryId = beanFactoryId;
         this.aggregationStrategy = requireNonNull(aggregationStrategy, "aggregationStrategy");
@@ -980,6 +1062,10 @@ Class<?> containerType() {
         return containerType;
     }
 
+    Class<?> rawType() {
+        return rawType;
+    }
+
     Class<?> elementType() {
         return elementType;
     }
@@ -1052,7 +1138,8 @@ private Object defaultOrException() {
             }
             return defaultValue;
         }
-        throw new IllegalArgumentException("Mandatory parameter/header is missing: " + httpElementName);
+        throw new IllegalArgumentException("Mandatory parameter/header/attribute is missing: " +
+                                           httpElementName);
     }
 
     @Override
@@ -1106,6 +1193,7 @@ private Builder(AnnotatedElement annotatedElement, Type type, String name) {
          */
         private Builder annotationType(Class<? extends Annotation> annotationType) {
             assert annotationType == Param.class ||
+                   annotationType == Attribute.class ||
                    annotationType == Header.class ||
                    annotationType == RequestObject.class : annotationType.getSimpleName();
             this.annotationType = annotationType;
@@ -1240,6 +1328,7 @@ private AnnotatedValueResolver build() {
             }
 
             final Class<?> containerType = getContainerType(unwrappedParameterizedType);
+            final Class<?> rawType = toRawType(unwrappedParameterizedType);
             final Class<?> elementType;
             final ParameterizedType parameterizedElementType;
 
@@ -1268,7 +1357,7 @@ private AnnotatedValueResolver build() {
             }
 
             return new AnnotatedValueResolver(annotationType, httpElementName, pathVariable, shouldExist,
-                                              isOptional, containerType, elementType,
+                                              isOptional, containerType, elementType, rawType,
                                               parameterizedElementType, defaultValue, description, resolver,
                                               beanFactoryId, aggregation);
         }

core/src/main/java/com/linecorp/armeria/server/annotation/Attribute.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2016 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you 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:
+ *
+ *   https://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 com.linecorp.armeria.server.annotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import com.linecorp.armeria.common.RequestContext;
+
+import io.netty.util.AttributeKey;
+
+/**
+ * Annotation for mapping an attribute of the given {@link AttributeKey}, retrieved
+ * from a {@link RequestContext}, onto the following elements.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.PARAMETER)
+public @interface Attribute {
+
+    /**
+     * The name of the {@link AttributeKey} to bind to.
+     * You might also want to specify the {@link #prefix()}.
+     */
+    String value();
+
+    /**
+     * The class of the {@link AttributeKey} to bind to. If you created an {@link AttributeKey} with
+     * {@code AttributeKey.valueOf(MyAttributeKeys.class, "INT_ATTR")},
+     * the {@link #prefix()} should be {@code MyAttributeKeys.class}.
+     * See <a href="https://armeria.dev/docs/advanced-custom-attributes/">advanced-customer-attributes</a>.
+     */
+    Class<?> prefix() default Attribute.class;
+}