opensearch-project
diff --git a/‎data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/server/CreateServer.java‎
Lines changed: 104 additions & 23 deletions b/‎data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/server/CreateServer.java‎
Lines changed: 104 additions & 23 deletions
@@ -5,7 +5,6 @@
 
 package org.opensearch.dataprepper.plugins.server;
 
-
 import com.linecorp.armeria.common.grpc.GrpcExceptionHandlerFunction;
 import com.linecorp.armeria.common.util.BlockingTaskExecutor;
 import com.linecorp.armeria.server.HttpService;
@@ -40,17 +39,17 @@
 import java.io.ByteArrayInputStream;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.function.Function;
 
-
 public class CreateServer {
     private final ServerConfiguration serverConfiguration;
     private final Logger LOG;
-    private final PluginMetrics  pluginMetrics;
+    private final PluginMetrics pluginMetrics;
     private String sourceName;
     private String pipelineName;
 
@@ -68,7 +67,53 @@ public CreateServer(final ServerConfiguration serverConfiguration, final Logger
         this.pipelineName = pipelineName;
     }
 
-    public <K, V> Server createGRPCServer(final GrpcAuthenticationProvider authenticationProvider, final BindableService grpcService, final CertificateProvider certificateProvider, final MethodDescriptor<K, V> methodDescriptor) {
+    /**
+     * Service configuration class for GRPC services
+     * @param <K> Request type
+     * @param <V> Response type
+     */
+    public static class GRPCServiceConfig<K, V> {
+        private final BindableService service;
+        private final String path;
+        private final MethodDescriptor<K, V> methodDescriptor;
+
+        public GRPCServiceConfig(BindableService service, String path, MethodDescriptor<K, V> methodDescriptor) {
+            this.service = service;
+            this.path = path;
+            this.methodDescriptor = methodDescriptor;
+        }
+
+        public GRPCServiceConfig(BindableService service) {
+            this.service = service;
+            this.path = null;
+            this.methodDescriptor = null;
+        }
+
+        public BindableService getService() {
+            return service;
+        }
+
+        public String getPath() {
+            return path;
+        }
+
+        public MethodDescriptor<K, V> getMethodDescriptor() {
+            return methodDescriptor;
+        }
+    }
+
+    /**
+     * Creates a GRPC server with multiple services, each with its own optional path and method descriptor
+     * @param authenticationProvider Provider for authentication
+     * @param grpcServiceConfigs List of service configurations 
+     * @param certificateProvider Provider for SSL/TLS certificates
+     * @return Configured server
+     */
+    public Server createGRPCServer(
+            final GrpcAuthenticationProvider authenticationProvider,
+            final List<GRPCServiceConfig<?, ?>> grpcServiceConfigs,
+            final CertificateProvider certificateProvider) {
+            
         final List<ServerInterceptor> serverInterceptors = getAuthenticationInterceptor(authenticationProvider);
 
         final GrpcServiceBuilder grpcServiceBuilder = GrpcService
@@ -77,13 +122,20 @@ public <K, V> Server createGRPCServer(final GrpcAuthenticationProvider authentic
                 .useBlockingTaskExecutor(true)
                 .exceptionHandler(createGrpExceptionHandler());
 
-        final String sourcePath = serverConfiguration.getPath();
-        if (sourcePath != null) {
-            final String transformedSourcePath = sourcePath.replace(PIPELINE_NAME_PLACEHOLDER, pipelineName);
-            grpcServiceBuilder.addService(transformedSourcePath,
-                    ServerInterceptors.intercept(grpcService, serverInterceptors), methodDescriptor);
-        } else {
-            grpcServiceBuilder.addService(ServerInterceptors.intercept(grpcService, serverInterceptors));
+        // Add each service with its own path and method descriptor
+        for (GRPCServiceConfig<?, ?> serviceConfig : grpcServiceConfigs) {
+            if (serviceConfig.getPath() != null && serviceConfig.getMethodDescriptor() != null) {
+                final String transformedPath = serviceConfig.getPath().replace(PIPELINE_NAME_PLACEHOLDER, pipelineName);
+                grpcServiceBuilder.addService(
+                    transformedPath,
+                    ServerInterceptors.intercept(serviceConfig.getService(), serverInterceptors), 
+                    serviceConfig.getMethodDescriptor());
+                LOG.info("Adding service with path: {}", transformedPath);
+            } else {
+                grpcServiceBuilder.addService(
+                    ServerInterceptors.intercept(serviceConfig.getService(), serverInterceptors));
+                LOG.info("Adding service without specific path");
+            }
         }
 
         if (serverConfiguration.hasHealthCheck()) {
@@ -151,7 +203,35 @@ public <K, V> Server createGRPCServer(final GrpcAuthenticationProvider authentic
         return sb.build();
     }
 
-    public Server createHTTPServer(final Buffer<Record<Log>> buffer, final CertificateProviderFactory certificateProviderFactory, final ArmeriaHttpAuthenticationProvider authenticationProvider, final HttpRequestExceptionHandler httpRequestExceptionHandler, final Object logService) {
+    /**
+     * Creates a GRPC server with a single service
+     * @param authenticationProvider Provider for authentication
+     * @param grpcService Service to be added
+     * @param certificateProvider Provider for SSL/TLS certificates
+     * @param methodDescriptor Method descriptor for the service
+     * @return Configured server
+     */
+    public <K, V> Server createGRPCServer(
+            final GrpcAuthenticationProvider authenticationProvider,
+            final BindableService grpcService,
+            final CertificateProvider certificateProvider,
+            final MethodDescriptor<K, V> methodDescriptor) {
+        
+        List<GRPCServiceConfig<?, ?>> serviceConfigs = new ArrayList<>();
+        if (serverConfiguration.getPath() != null) {
+            serviceConfigs.add(new GRPCServiceConfig<>(grpcService, serverConfiguration.getPath(), methodDescriptor));
+        } else {
+            serviceConfigs.add(new GRPCServiceConfig<>(grpcService));
+        }
+        
+        return createGRPCServer(authenticationProvider, serviceConfigs, certificateProvider);
+    }
+
+
+    public Server createHTTPServer(final Buffer<Record<Log>> buffer,
+            final CertificateProviderFactory certificateProviderFactory,
+            final ArmeriaHttpAuthenticationProvider authenticationProvider,
+            final HttpRequestExceptionHandler httpRequestExceptionHandler, final Object logService) {
         final ServerBuilder sb = Server.builder();
 
         sb.disableServerHeader();
@@ -163,17 +243,16 @@ public Server createHTTPServer(final Buffer<Record<Log>> buffer, final Certifica
             // TODO: enable encrypted key with password
             sb.https(serverConfiguration.getPort()).tls(
                     new ByteArrayInputStream(certificate.getCertificate().getBytes(StandardCharsets.UTF_8)),
-                    new ByteArrayInputStream(certificate.getPrivateKey().getBytes(StandardCharsets.UTF_8)
-                    )
-            );
+                    new ByteArrayInputStream(certificate.getPrivateKey().getBytes(StandardCharsets.UTF_8)));
         } else {
             LOG.warn("Creating " + sourceName + " without SSL/TLS. This is not secure.");
             LOG.warn("In order to set up TLS for the " + sourceName + ", go here: https://github.com/opensearch-project/data-prepper/tree/main/data-prepper-plugins/http-source#ssl");
             sb.http(serverConfiguration.getPort());
         }
 
-        if(serverConfiguration.getAuthentication() != null) {
-            final Optional<Function<? super HttpService, ? extends HttpService>> optionalAuthDecorator = authenticationProvider.getAuthenticationDecorator();
+        if (serverConfiguration.getAuthentication() != null) {
+            final Optional<Function<? super HttpService, ? extends HttpService>> optionalAuthDecorator = authenticationProvider
+                    .getAuthenticationDecorator();
 
             if (serverConfiguration.isUnauthenticatedHealthCheck()) {
                 optionalAuthDecorator.ifPresent(authDecorator -> sb.decorator(REGEX_HEALTH, authDecorator));
@@ -184,7 +263,7 @@ public Server createHTTPServer(final Buffer<Record<Log>> buffer, final Certifica
 
         sb.maxNumConnections(serverConfiguration.getMaxConnectionCount());
         sb.requestTimeout(Duration.ofMillis(serverConfiguration.getRequestTimeoutInMillis()));
-        if(serverConfiguration.getMaxRequestLength() != null) {
+        if (serverConfiguration.getMaxRequestLength() != null) {
             sb.maxRequestLength(serverConfiguration.getMaxRequestLength().getBytes());
         }
         final int threads = serverConfiguration.getThreadCount();
@@ -193,15 +272,17 @@ public Server createHTTPServer(final Buffer<Record<Log>> buffer, final Certifica
         final int maxPendingRequests = serverConfiguration.getMaxPendingRequests();
         final LogThrottlingStrategy logThrottlingStrategy = new LogThrottlingStrategy(
                 maxPendingRequests, blockingTaskExecutor.getQueue());
-        final LogThrottlingRejectHandler logThrottlingRejectHandler = new LogThrottlingRejectHandler(maxPendingRequests, pluginMetrics);
+        final LogThrottlingRejectHandler logThrottlingRejectHandler = new LogThrottlingRejectHandler(maxPendingRequests,
+                pluginMetrics);
 
         final String httpSourcePath = serverConfiguration.getPath().replace(PIPELINE_NAME_PLACEHOLDER, pipelineName);
         sb.decorator(httpSourcePath, ThrottlingService.newDecorator(logThrottlingStrategy, logThrottlingRejectHandler));
 
         if (CompressionOption.NONE.equals(serverConfiguration.getCompression())) {
             sb.annotatedService(httpSourcePath, logService, httpRequestExceptionHandler);
         } else {
-            sb.annotatedService(httpSourcePath, logService, DecodingService.newDecorator(), httpRequestExceptionHandler);
+            sb.annotatedService(httpSourcePath, logService, DecodingService.newDecorator(),
+                    httpRequestExceptionHandler);
         }
 
         if (serverConfiguration.hasHealthCheck()) {
@@ -220,12 +301,12 @@ private GrpcExceptionHandlerFunction createGrpExceptionHandler() {
         return new GrpcRequestExceptionHandler(pluginMetrics, retryInfo.getMinDelay(), retryInfo.getMaxDelay());
     }
 
-    private List<ServerInterceptor> getAuthenticationInterceptor(final GrpcAuthenticationProvider authenticationProvider) {
+    private List<ServerInterceptor> getAuthenticationInterceptor(
+            final GrpcAuthenticationProvider authenticationProvider) {
         final ServerInterceptor authenticationInterceptor = authenticationProvider.getAuthenticationInterceptor();
         if (authenticationInterceptor == null) {
             return Collections.emptyList();
         }
         return Collections.singletonList(authenticationInterceptor);
     }
-}
-//
+}