add hook in StreamManager plugin to get its instance on initialization

Signed-off-by: Rishabh Maurya <[email protected]>
rishabhmaurya · Feb 24, 2025 · 5af10e0 · 5af10e0
1 parent 7c098a5
commit 5af10e0
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 2 deletions.
diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
@@ -56,6 +56,7 @@
 import org.opensearch.action.search.SearchTransportService;
 import org.opensearch.action.support.TransportAction;
 import org.opensearch.action.update.UpdateHelper;
+import org.opensearch.arrow.spi.StreamManager;
 import org.opensearch.bootstrap.BootstrapCheck;
 import org.opensearch.bootstrap.BootstrapContext;
 import org.opensearch.cluster.ClusterInfoService;
@@ -218,6 +219,7 @@
 import org.opensearch.plugins.SearchPipelinePlugin;
 import org.opensearch.plugins.SearchPlugin;
 import org.opensearch.plugins.SecureSettingsFactory;
+import org.opensearch.plugins.StreamManagerPlugin;
 import org.opensearch.plugins.SystemIndexPlugin;
 import org.opensearch.plugins.TaskManagerClientPlugin;
 import org.opensearch.plugins.TelemetryAwarePlugin;
@@ -309,11 +311,13 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static java.util.stream.Collectors.toList;
+import static org.opensearch.common.util.FeatureFlags.ARROW_STREAMS_SETTING;
 import static org.opensearch.common.util.FeatureFlags.BACKGROUND_TASK_EXECUTION_EXPERIMENTAL;
 import static org.opensearch.common.util.FeatureFlags.TELEMETRY;
 import static org.opensearch.env.NodeEnvironment.collectFileCacheDataPath;
@@ -1386,6 +1390,33 @@ protected Node(
                 cacheService
             );
 
+            Supplier<StreamManager> streamManager = null;
+            if (FeatureFlags.isEnabled(ARROW_STREAMS_SETTING)) {
+                List<StreamManagerPlugin> streamManagerPlugins = pluginsService.filterPlugins(StreamManagerPlugin.class);
+                if (!streamManagerPlugins.isEmpty()) {
+                    for (StreamManagerPlugin smPlugin : streamManagerPlugins) {
+                        Supplier<StreamManager> baseStreamManager = smPlugin.getStreamManager();
+                        if (streamManager == null && baseStreamManager != null) {
+                            streamManager = baseStreamManager;
+                            logger.info("StreamManager initialized");
+                        } else if (streamManager != null && baseStreamManager != null) {
+                            throw new IllegalStateException(
+                                String.format(
+                                    Locale.ROOT,
+                                    "Only one StreamManagerPlugin can be installed. Found: %d",
+                                    streamManagerPlugins.size()
+                                )
+                            );
+                        }
+                    }
+                    if (streamManager != null) {
+                        for (StreamManagerPlugin plugin : streamManagerPlugins) {
+                            plugin.onStreamManagerInitialized(streamManager);
+                        }
+                    }
+                }
+            }
+
             final SearchService searchService = newSearchService(
                 clusterService,
                 indicesService,

diff --git a/server/src/main/java/org/opensearch/plugins/StreamManagerPlugin.java b/server/src/main/java/org/opensearch/plugins/StreamManagerPlugin.java
@@ -14,7 +14,9 @@
 
 /**
  * An interface for OpenSearch plugins to implement to provide a StreamManager.
- * Plugins can implement this interface to provide custom StreamManager implementation.
+ * Plugins can implement this interface to provide custom StreamManager implementation
+ * or get a reference to the StreamManager instance provided by OpenSearch.
+ *
  * @see StreamManager
  */
 public interface StreamManagerPlugin {
@@ -23,5 +25,13 @@ public interface StreamManagerPlugin {
      *
      * @return The StreamManager instance
      */
-    Supplier<StreamManager> getStreamManager();
+    default Supplier<StreamManager> getStreamManager() {
+        return null;
+    }
+
+    /**
+     * Called when the StreamManager is initialized.
+     * @param streamManager Supplier of the StreamManager instance
+     */
+    default void onStreamManagerInitialized(Supplier<StreamManager> streamManager) {}
 }