Addressing Comments

Signed-off-by: Ajay Kumar Movva <[email protected]>
opensearch-project · Oct 19, 2023 · fbf0934 · fbf0934
1 parent 0c6a63c
commit fbf0934
Show file tree

Hide file tree

Showing 6 changed files with 100 additions and 299 deletions.
diff --git a/...st/java/org/opensearch/ratelimitting/admissioncontrol/SimpleEnableAdmissionControlIT.java b/...st/java/org/opensearch/ratelimitting/admissioncontrol/SimpleEnableAdmissionControlIT.java
diff --git a/server/src/main/java/org/opensearch/common/network/NetworkModule.java b/server/src/main/java/org/opensearch/common/network/NetworkModule.java
@@ -61,7 +61,6 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.Transport;
 import org.opensearch.transport.TransportInterceptor;
-import org.opensearch.transport.TransportInterceptorRegistry;
 import org.opensearch.transport.TransportRequest;
 import org.opensearch.transport.TransportRequestHandler;
 
@@ -150,15 +149,9 @@ public NetworkModule(
         NetworkService networkService,
         HttpServerTransport.Dispatcher dispatcher,
         ClusterSettings clusterSettings,
-        Tracer tracer,
-        TransportInterceptorRegistry transportInterceptorRegistry
+        Tracer tracer
     ) {
         this.settings = settings;
-        // Adding core interceptors first and other plugin interceptor will follow
-        List<TransportInterceptor> coreTransportInterceptors = transportInterceptorRegistry.getTransportInterceptors();
-        for (TransportInterceptor interceptor : coreTransportInterceptors) {
-            registerTransportInterceptor(interceptor);
-        }
         for (NetworkPlugin plugin : plugins) {
             Map<String, Supplier<HttpServerTransport>> httpTransportFactory = plugin.getHttpTransports(
                 settings,
@@ -274,6 +267,14 @@ private void registerTransportInterceptor(TransportInterceptor interceptor) {
         this.transportInterceptors.add(Objects.requireNonNull(interceptor, "interceptor must not be null"));
     }
 
+    /**
+     * Registers a new {@link TransportInterceptor}
+     * This method used to register CoreInterceptors before the plugin interceptors
+     */
+    public void registerCoreTransportInterceptor(TransportInterceptor interceptor) {
+        this.transportInterceptors.add(0, Objects.requireNonNull(interceptor, "interceptor must not be null"));
+    }
+
     /**
      * Returns a composite {@link TransportInterceptor} containing all registered interceptors
      * @see #registerTransportInterceptor(TransportInterceptor)

diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
@@ -239,7 +239,6 @@
 import org.opensearch.transport.RemoteClusterService;
 import org.opensearch.transport.Transport;
 import org.opensearch.transport.TransportInterceptor;
-import org.opensearch.transport.TransportInterceptorRegistry;
 import org.opensearch.transport.TransportService;
 import org.opensearch.usage.UsageService;
 import org.opensearch.watcher.ResourceWatcherService;
@@ -822,18 +821,6 @@ protected Node(
                 remoteStoreStatsTrackerFactory,
                 recoverySettings
             );
-            TransportInterceptorRegistry transportInterceptorRegistry = new TransportInterceptorRegistry();
-
-            final AdmissionControlService admissionControlService = new AdmissionControlService(
-                settings,
-                clusterService.getClusterSettings(),
-                threadPool
-            );
-
-            AdmissionControlTransportInterceptor admissionControlTransportInterceptor = new AdmissionControlTransportInterceptor(
-                admissionControlService
-            );
-            transportInterceptorRegistry.addTransportInterceptor(admissionControlTransportInterceptor);
 
             final AliasValidator aliasValidator = new AliasValidator();
 
@@ -906,6 +893,16 @@ protected Node(
 
             final RestController restController = actionModule.getRestController();
 
+            final AdmissionControlService admissionControlService = new AdmissionControlService(
+                settings,
+                clusterService.getClusterSettings(),
+                threadPool
+            );
+
+            AdmissionControlTransportInterceptor admissionControlTransportInterceptor = new AdmissionControlTransportInterceptor(
+                admissionControlService
+            );
+
             final NetworkModule networkModule = new NetworkModule(
                 settings,
                 pluginsService.filterPlugins(NetworkPlugin.class),
@@ -918,9 +915,10 @@ protected Node(
                 networkService,
                 restController,
                 clusterService.getClusterSettings(),
-                tracer,
-                transportInterceptorRegistry
+                tracer
             );
+            networkModule.registerCoreTransportInterceptor(admissionControlTransportInterceptor);
+
             Collection<UnaryOperator<Map<String, IndexTemplateMetadata>>> indexTemplateMetadataUpgraders = pluginsService.filterPlugins(
                 Plugin.class
             ).stream().map(Plugin::getIndexTemplateMetadataUpgrader).collect(Collectors.toList());

diff --git a/server/src/main/java/org/opensearch/transport/TransportInterceptorRegistry.java b/server/src/main/java/org/opensearch/transport/TransportInterceptorRegistry.java