Support adding gRPC interceptors using annotation '@GrpcInterceptor'

grpc/src/main/java/com/linecorp/armeria/server/grpc/GrpcInterceptor.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2024 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.grpc;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Repeatable;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import io.grpc.ServerInterceptor;
+
+/**
+ * Specifies a {@link ServerInterceptor} class which intercepts requests and responses of a gRPC service or its
+ * methods.
+ */
+@Repeatable(GrpcInterceptors.class)
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ ElementType.TYPE, ElementType.METHOD })
+public @interface GrpcInterceptor {
+
+    /**
+     * {@link ServerInterceptor} implementation type.
+     */
+    Class<? extends ServerInterceptor> value();
+}

grpc/src/main/java/com/linecorp/armeria/server/grpc/GrpcInterceptors.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2024 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.grpc;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * The containing annotation type for {@link GrpcInterceptor}.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ ElementType.TYPE, ElementType.METHOD })
+public @interface GrpcInterceptors {
+
+    /**
+     * An array of {@link GrpcInterceptor}s.
+     */
+    GrpcInterceptor[] value();
+}

grpc/src/main/java/com/linecorp/armeria/server/grpc/GrpcServiceBuilder.java

@@ -58,7 +58,6 @@
 import com.linecorp.armeria.server.VirtualHost;
 import com.linecorp.armeria.server.VirtualHostBuilder;
 import com.linecorp.armeria.server.encoding.EncodingService;
-import com.linecorp.armeria.server.grpc.HandlerRegistry.Entry;
 import com.linecorp.armeria.unsafe.grpc.GrpcUnsafeBufferUtil;
 
 import io.grpc.BindableService;
@@ -971,7 +970,6 @@ private ImmutableList.Builder<ServerInterceptor> interceptors() {
      * without interfering with other services.
      */
     public GrpcService build() {
-        final HandlerRegistry handlerRegistry;
         if (USE_COROUTINE_CONTEXT_INTERCEPTOR) {
             final ServerInterceptor coroutineContextInterceptor =
                     new ArmeriaCoroutineContextInterceptor(useBlockingTaskExecutor);
@@ -1009,24 +1007,13 @@ public GrpcService build() {
         }
         registryBuilder.setDefaultExceptionHandler(grpcExceptionHandler);
 
-        if (interceptors != null) {
-            final HandlerRegistry.Builder newRegistryBuilder = new HandlerRegistry.Builder();
-            final ImmutableList<ServerInterceptor> interceptors = this.interceptors.build();
-            for (Entry entry : registryBuilder.entries()) {
-                final MethodDescriptor<?, ?> methodDescriptor = entry.method();
-                final ServerServiceDefinition intercepted =
-                        ServerInterceptors.intercept(entry.service(), interceptors);
-                newRegistryBuilder.addService(entry.path(), intercepted, methodDescriptor, entry.type(),
-                                              entry.additionalDecorators());
-            }
-            if (grpcExceptionHandler != null) {
-                newRegistryBuilder.setDefaultExceptionHandler(grpcExceptionHandler);
-            }
-            handlerRegistry = newRegistryBuilder.build();
-        } else {
-            handlerRegistry = registryBuilder.build();
+        // Interceptors passed via the grpc service builder.
+        if (this.interceptors != null) {
+            registryBuilder.addInterceptors(this.interceptors.build());
         }
 
+        final HandlerRegistry handlerRegistry = registryBuilder.build();
+
         GrpcService grpcService = new FramedGrpcService(
                 handlerRegistry,
                 firstNonNull(decompressorRegistry, DecompressorRegistry.getDefaultInstance()),

grpc/src/main/java/com/linecorp/armeria/server/grpc/HandlerRegistry.java

@@ -60,6 +60,8 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.CaseFormat;
@@ -81,6 +83,8 @@
 import com.linecorp.armeria.server.annotation.Blocking;
 
 import io.grpc.MethodDescriptor;
+import io.grpc.ServerInterceptor;
+import io.grpc.ServerInterceptors;
 import io.grpc.ServerMethodDefinition;
 import io.grpc.ServerServiceDefinition;
 
@@ -207,6 +211,7 @@ private static HttpService applyDecorators(
 
     static final class Builder {
         private final List<Entry> entries = new ArrayList<>();
+        private final List<ServerInterceptor> interceptors = new ArrayList<>();
 
         @Nullable
         private GrpcExceptionHandlerFunction defaultExceptionHandler;
@@ -240,6 +245,16 @@ Builder setDefaultExceptionHandler(GrpcExceptionHandlerFunction defaultException
             return this;
         }
 
+        /**
+         * Find the applied {@link GrpcInterceptor} annotations and apply them with the given
+         * {@link ServerInterceptor}s to all services.
+         */
+        Builder addInterceptors(List<ServerInterceptor> globalInterceptors) {
+            requireNonNull(globalInterceptors, "globalInterceptors");
+            this.interceptors.addAll(globalInterceptors);
+            return this;
+        }
+
         private static String normalizePath(String path, boolean isServicePath) {
             if (path.isEmpty()) {
                 return path;
@@ -287,8 +302,157 @@ private static void putGrpcExceptionHandlerIfPresent(
             });
         }
 
-        List<Entry> entries() {
-            return entries;
+        /**
+         * Get the Method from the MethodDescriptor for the given gRPC service. This method can be used to
+         * get annotations applied to the method.
+         *
+         * @param clazz The class of the service.
+         * @param methodDescriptor The method descriptor to get the method for.
+         * @return The method for the given method descriptor.
+         */
+        private static Optional<Method> getMethodFromMethodDescriptor(Class<?> clazz,
+                                                               MethodDescriptor<?, ?> methodDescriptor) {
+            final String methodName = methodDescriptor.getBareMethodName();
+
+            if (methodName == null) {
+                return Optional.empty();
+            }
+
+            final String matchingMethodName = CaseFormat.UPPER_CAMEL
+                    .converterTo(CaseFormat.LOWER_CAMEL)
+                    .convert(methodName);
+
+            if (matchingMethodName == null) {
+                return Optional.empty();
+            }
+
+            return InternalReflectionUtils.getAllSortedMethods(clazz, withModifier(Modifier.PUBLIC))
+                                          .stream()
+                                          .filter(m -> matchingMethodName.equals(m.getName()))
+                                          .findFirst();
+        }
+
+        /**
+         * Get the interceptors for the given method created by using annotations.
+         * @param clazz The class of the service.
+         * @param method The method to get interceptors for.
+         * @param dependencyInjector The dependency injector to use.
+         * @param globalInterceptors The global interceptors to use. This comes from the builder.
+         * @return The list of interceptors for the given method in order.
+         */
+        private static List<ServerInterceptor>
+        getInterceptorsFromAnnotations(Class<?> clazz, Method method, DependencyInjector dependencyInjector,
+                                       List<ServerInterceptor> globalInterceptors) {
+            final List<ServerInterceptor> methodAndClassInterceptors =
+                    AnnotationUtil.getAnnotatedInstances(method, clazz,
+                                                         GrpcInterceptor.class,
+                                                         ServerInterceptor.class,
+                                                         dependencyInjector).build();
+
+            return Stream.concat(methodAndClassInterceptors.stream(), globalInterceptors.stream())
+                         .collect(Collectors.toList());
+        }
+
+        private static String calculateServicePath(Entry entry, MethodDescriptor<?, ?> methodDescriptor) {
+            // Use the path of method descriptor instead of the path of service. We are adding a
+            // single method to the registry as a service opposed to adding the entire service
+            // to the registry. The reason is that we can't intercept individual methods if we
+            // add the service as a whole to the registry.
+            return entry.path() + '/' + methodDescriptor.getBareMethodName();
+        }
+
+        /**
+         * Intercepts the entries using {@link GrpcInterceptor} annotations and globally added interceptors.
+         * This method should only be called once. Otherwise the interceptors will be added multiple times.
+         * @param dependencyInjector The dependency injector to find the annotations.
+         */
+        private void interceptEntries(DependencyInjector dependencyInjector) {
+            // Copy services and methods to intercept them.
+            final List<Entry> initialListOfEntries = ImmutableList.copyOf(this.entries);
+
+            for (Entry entry : initialListOfEntries) {
+                final MethodDescriptor<?, ?> methodDescriptor = entry.method();
+
+                if (entry.type() != null && methodDescriptor == null) {
+                    // A "Service" entry thus there is no method descriptor.
+
+                    final List<MethodDescriptor<?, ?>> serverMethodDescriptors =
+                            entry.service().getMethods().stream()
+                                 .map(ServerMethodDefinition::getMethodDescriptor)
+                                 .collect(Collectors.toList());
+
+                    final boolean shouldSplitServiceToMethod = serverMethodDescriptors.stream().anyMatch(
+                            methodDescriptor1 -> {
+                                final Optional<Method> methodOption =
+                                        getMethodFromMethodDescriptor(entry.type(), methodDescriptor1);
+
+                                if (methodOption.isPresent()) {
+                                    final List<ServerInterceptor> allInterceptors =
+                                            getInterceptorsFromAnnotations(entry.type(), methodOption.get(),
+                                                                           dependencyInjector,
+                                                                           ImmutableList.of());
+
+                                    return !allInterceptors.isEmpty();
+                                }
+
+                                return false;
+                            }
+                    );
+
+                    if (shouldSplitServiceToMethod) {
+                        // Add all methods of the service to the new registry builder one by one and intercept.
+                        for (MethodDescriptor<?, ?> serverMethodDescriptor : serverMethodDescriptors) {
+                            final Optional<Method> methodOption =
+                                    getMethodFromMethodDescriptor(entry.type(), serverMethodDescriptor);
+
+                            if (methodOption.isPresent()) {
+                                final List<ServerInterceptor> allInterceptors =
+                                        getInterceptorsFromAnnotations(entry.type(), methodOption.get(),
+                                                                       dependencyInjector, this.interceptors);
+
+                                final ServerServiceDefinition intercepted =
+                                        ServerInterceptors.intercept(entry.service(), allInterceptors);
+
+                                final String path = calculateServicePath(entry, serverMethodDescriptor);
+                                addService(path, intercepted,
+                                           serverMethodDescriptor, entry.type(),
+                                           ImmutableList.copyOf(entry.additionalDecorators()));
+                            }
+                        }
+                    } else {
+                        // No need to split service into individual methods if there are no interceptors.
+                        final ServerServiceDefinition intercepted =
+                                ServerInterceptors.intercept(entry.service(), this.interceptors);
+                        addService(entry.path(), intercepted, methodDescriptor,
+                                   entry.type(), entry.additionalDecorators());
+                    }
+                } else if (entry.type() != null) {
+                    // A "Method" entry
+                    final Optional<Method> methodOption =
+                            getMethodFromMethodDescriptor(entry.type(), methodDescriptor);
+
+                    if (methodOption.isPresent()) {
+                        final List<ServerInterceptor> allInterceptors =
+                                getInterceptorsFromAnnotations(entry.type(), methodOption.get(),
+                                                               dependencyInjector, this.interceptors);
+
+                        final ServerServiceDefinition intercepted =
+                                ServerInterceptors.intercept(entry.service(), allInterceptors);
+                        addService(entry.path(), intercepted, methodDescriptor,
+                                   entry.type(), entry.additionalDecorators());
+                    }
+                } else {
+                    // Others
+                    // Only intercept the service with global interceptors.
+                    final ServerServiceDefinition intercepted = ServerInterceptors.intercept(entry.service(),
+                                                                                             this.interceptors);
+                    addService(entry.path(), intercepted, methodDescriptor,
+                               entry.type(), entry.additionalDecorators());
+                }
+            }
+
+            // Remove the original entries as they are now re-added as intercepted entries.
+            entries.removeAll(initialListOfEntries);
         }
 
         HandlerRegistry build() {
@@ -311,6 +475,12 @@ HandlerRegistry build() {
             final ImmutableMap.Builder<ServerMethodDefinition<?, ?>, GrpcExceptionHandlerFunction>
                     grpcExceptionHandlersBuilder = ImmutableMap.builder();
             final DependencyInjector dependencyInjector = new ReflectiveDependencyInjector();
+
+            // Intercept entries using {@link GrpcInterceptor} annotations and globally added interceptors.
+            // TODO(ikhoon): Use ServerConfig.dependencyInjector() instead of creating
+            //  ReflectiveDependencyInjector directly.
+            interceptEntries(dependencyInjector);
+
             for (Entry entry : entries) {
                 final ServerServiceDefinition service = entry.service();
                 final String path = entry.path();