Datanode sniffer implementation for opensearch client (#22184)

* Datanode sniffer implementation for opensearch client

* added license

* code cleanup

* fixed import

* Added changelog

Author: todvora

changelog/unreleased/pr-22184.toml

@@ -0,0 +1,5 @@
+type = "added"
+message = "Added datanode sniffer to opensearch clients where applicable"
+
+issues = []
+pulls = ["22155"]

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/OpenSearch2Module.java

@@ -18,6 +18,7 @@
 
 import com.google.inject.assistedinject.FactoryModuleBuilder;
 import com.google.inject.binder.LinkedBindingBuilder;
+import com.google.inject.multibindings.Multibinder;
 import org.graylog.events.search.MoreSearchAdapter;
 import org.graylog.plugins.datanode.DatanodeUpgradeServiceAdapter;
 import org.graylog.plugins.views.migrations.V20200730000000_AddGl2MessageIdFieldAliasForEvents;
@@ -28,6 +29,12 @@
 import org.graylog.storage.opensearch2.client.OSCredentialsProvider;
 import org.graylog.storage.opensearch2.fieldtypes.streams.StreamsForFieldRetrieverOS2;
 import org.graylog.storage.opensearch2.migrations.V20170607164210_MigrateReopenedIndicesToAliasesClusterStateOS2;
+import org.graylog.storage.opensearch2.sniffer.SnifferBuilder;
+import org.graylog.storage.opensearch2.sniffer.SnifferFilter;
+import org.graylog.storage.opensearch2.sniffer.impl.DatanodesSniffer;
+import org.graylog.storage.opensearch2.sniffer.impl.NodeAttributesFilter;
+import org.graylog.storage.opensearch2.sniffer.impl.NodeLoggingFilter;
+import org.graylog.storage.opensearch2.sniffer.impl.OpensearchClusterSniffer;
 import org.graylog.storage.opensearch2.views.migrations.V20200730000000_AddGl2MessageIdFieldAliasForEventsOS2;
 import org.graylog2.indexer.IndexToolsAdapter;
 import org.graylog2.indexer.client.IndexerHostsAdapter;
@@ -93,6 +100,14 @@ protected void configure() {
         bind(CredentialsProvider.class).toProvider(OSCredentialsProvider.class);
         bindForSupportedVersion(DatanodeUpgradeServiceAdapter.class).to(DatanodeUpgradeServiceAdapterOS2.class);
 
+        Multibinder<SnifferBuilder> snifferBuilders = Multibinder.newSetBinder(binder(), SnifferBuilder.class);
+        snifferBuilders.addBinding().to(OpensearchClusterSniffer.class);
+        snifferBuilders.addBinding().to(DatanodesSniffer.class);
+
+        Multibinder<SnifferFilter> snifferFilters = Multibinder.newSetBinder(binder(), SnifferFilter.class);
+        snifferFilters.addBinding().to(NodeAttributesFilter.class);
+        snifferFilters.addBinding().to(NodeLoggingFilter.class);
+
         bindForSupportedVersion(IndexerHostsAdapter.class).to(IndexerHostsAdapterOS2.class);
     }
 

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/RestClientProvider.java

@@ -17,6 +17,7 @@
 package org.graylog.storage.opensearch2;
 
 import com.google.common.base.Suppliers;
+import jakarta.annotation.Nonnull;
 import jakarta.inject.Inject;
 import jakarta.inject.Provider;
 import jakarta.inject.Singleton;
@@ -25,24 +26,28 @@
 import org.graylog.shaded.opensearch2.org.apache.http.client.CredentialsProvider;
 import org.graylog.shaded.opensearch2.org.opensearch.client.RestClient;
 import org.graylog.shaded.opensearch2.org.opensearch.client.RestHighLevelClient;
-import org.graylog.shaded.opensearch2.org.opensearch.client.sniff.OpenSearchNodesSniffer;
+import org.graylog.shaded.opensearch2.org.opensearch.client.sniff.NodesSniffer;
+import org.graylog.shaded.opensearch2.org.opensearch.client.sniff.Sniffer;
+import org.graylog.storage.opensearch2.sniffer.SnifferAggregator;
+import org.graylog.storage.opensearch2.sniffer.SnifferBuilder;
+import org.graylog.storage.opensearch2.sniffer.SnifferFilter;
 import org.graylog2.configuration.ElasticsearchClientConfiguration;
 import org.graylog2.configuration.IndexerHosts;
 import org.graylog2.configuration.RunsWithDataNode;
 import org.graylog2.security.IndexerJwtAuthTokenProvider;
 import org.graylog2.security.TrustManagerAndSocketFactoryProvider;
 import org.graylog2.system.shutdown.GracefulShutdownService;
-import jakarta.annotation.Nonnull;
 
 import java.net.URI;
 import java.util.List;
-import java.util.Locale;
-import java.util.concurrent.TimeUnit;
+import java.util.Set;
 import java.util.function.Supplier;
 
 @Singleton
 public class RestClientProvider implements Provider<RestHighLevelClient> {
     private final Supplier<RestHighLevelClient> clientSupplier;
+    private final Set<SnifferBuilder> snifferBuilders;
+    private final Set<SnifferFilter> snifferFilters;
     private final GracefulShutdownService shutdownService;
     private final ElasticsearchClientConfiguration configuration;
     private final CredentialsProvider credentialsProvider;
@@ -58,46 +63,46 @@ public RestClientProvider(
             CredentialsProvider credentialsProvider,
             TrustManagerAndSocketFactoryProvider trustManagerAndSocketFactoryProvider,
             @RunsWithDataNode Boolean runsWithDataNode,
-            IndexerJwtAuthTokenProvider indexerJwtAuthTokenProvider) {
+            IndexerJwtAuthTokenProvider indexerJwtAuthTokenProvider,
+            Set<SnifferBuilder> snifferBuilders,
+            Set<SnifferFilter> snifferFilters
+    ) {
         this.shutdownService = shutdownService;
         this.configuration = configuration;
         this.credentialsProvider = credentialsProvider;
         this.trustManagerAndSocketFactoryProvider = trustManagerAndSocketFactoryProvider;
         this.runsWithDataNode = runsWithDataNode;
         this.indexerJwtAuthTokenProvider = indexerJwtAuthTokenProvider;
         this.clientSupplier = Suppliers.memoize(() -> createClient(hosts));
+        this.snifferBuilders = snifferBuilders;
+        this.snifferFilters = snifferFilters;
     }
 
 
     @Nonnull
     private RestHighLevelClient createClient(List<URI> hosts) {
         final RestHighLevelClient client = buildBasicRestClient(hosts);
+        registerSniffers(client);
+        return client;
+    }
 
-        var sniffer = SnifferWrapper.create(
-                client,
-                TimeUnit.SECONDS.toMillis(5),
-                configuration.discoveryFrequency(),
-                mapDefaultScheme(configuration.defaultSchemeForDiscoveredNodes())
-        );
+    private void registerSniffers(RestHighLevelClient client) {
+        final List<NodesSniffer> sniffers = snifferBuilders.stream()
+                .filter(SnifferBuilder::enabled)
+                .map(b -> b.create(client.getLowLevelClient()))
+                .toList();
 
-        if (configuration.discoveryEnabled()) {
-            sniffer.add(FilteredOpenSearchNodesSniffer.create(configuration.discoveryFilter()));
-        }
-        if (configuration.isNodeActivityLogger()) {
-            sniffer.add(NodeListSniffer.create());
-        }
+        if (!sniffers.isEmpty()) {
+            final List<SnifferFilter> filters = snifferFilters.stream().filter(SnifferFilter::enabled).toList();
+            final SnifferAggregator snifferAggregator = new SnifferAggregator(sniffers, filters);
 
-        sniffer.build().ifPresent(s -> shutdownService.register(s::close));
-        return client;
-    }
+            final Sniffer sniffer = Sniffer.builder(client.getLowLevelClient())
+                    .setSniffIntervalMillis(Math.toIntExact(configuration.discoveryFrequency().toMilliseconds()))
+                    .setNodesSniffer(snifferAggregator)
+                    .build();
 
-    private OpenSearchNodesSniffer.Scheme mapDefaultScheme(String defaultSchemeForDiscoveredNodes) {
-        return switch (defaultSchemeForDiscoveredNodes.toUpperCase(Locale.ENGLISH)) {
-            case "HTTP" -> OpenSearchNodesSniffer.Scheme.HTTP;
-            case "HTTPS" -> OpenSearchNodesSniffer.Scheme.HTTPS;
-            default ->
-                    throw new IllegalArgumentException("Invalid default scheme for discovered OS nodes: " + defaultSchemeForDiscoveredNodes);
-        };
+            shutdownService.register(sniffer::close);
+        }
     }
 
     @Override

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/SnifferWrapper.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog.storage.opensearch2.sniffer;
+
+import jakarta.annotation.Nonnull;
+import org.graylog.shaded.opensearch2.org.opensearch.client.Node;
+import org.graylog.shaded.opensearch2.org.opensearch.client.sniff.NodesSniffer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.function.Predicate;
+import java.util.function.Function;
+
+/**
+ * This aggregator has two functions. It allows more sniffers to work together, each providing its own list of available nodes,
+ * even if other sniffers fail. Then it applies filters to the result. Each filter can modify the nodes list.
+ */
+public class SnifferAggregator implements org.graylog.shaded.opensearch2.org.opensearch.client.sniff.NodesSniffer {
+
+    private static final Logger LOG = LoggerFactory.getLogger(SnifferAggregator.class);
+
+    private final List<NodesSniffer> sniffers;
+    private final List<SnifferFilter> filters;
+
+    public SnifferAggregator(List<NodesSniffer> sniffers, List<SnifferFilter> filters) {
+        this.sniffers = sniffers;
+        this.filters = filters;
+    }
+
+    @Override
+    public List<Node> sniff() throws IOException {
+        List<Node> discoveredNodes = discoverNodes().stream()
+                .filter(distinctByKey(n -> n.getHost().toURI()))
+                .collect(Collectors.toCollection(ArrayList::new));
+
+        for (SnifferFilter sniffer : filters) {
+            discoveredNodes = sniffer.filterNodes(discoveredNodes);
+        }
+        return discoveredNodes;
+    }
+
+    @Nonnull
+    private List<Node> discoverNodes() {
+        return sniffers.stream().flatMap(SnifferAggregator::sniff).toList();
+    }
+
+    @Nonnull
+    private static Stream<Node> sniff(NodesSniffer sniffer) {
+        try {
+            return sniffer.sniff().stream();
+        } catch (IOException e) {
+            LOG.warn("Sniffer {} failed to sniff nodes: {}", sniffer.getClass(), e.getMessage());
+            return Stream.empty();
+        }
+    }
+
+    private static <T> Predicate<T> distinctByKey(Function<? super T, ?> keyExtractor) {
+        Set<Object> seen = ConcurrentHashMap.newKeySet();
+        return t -> seen.add(keyExtractor.apply(t));
+    }
+}

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/sniffer/SnifferAggregator.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog.storage.opensearch2.sniffer;
+
+import org.graylog.shaded.opensearch2.org.opensearch.client.RestClient;
+import org.graylog.shaded.opensearch2.org.opensearch.client.sniff.NodesSniffer;
+
+public interface SnifferBuilder {
+    /**
+     * @return true if the configuration of this node allows that type of sniffer
+     */
+    boolean enabled();
+
+    /**
+     * @return Always a new instance of a sniffer.
+     */
+    NodesSniffer create(RestClient restClient);
+}

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/sniffer/SnifferBuilder.java

@@ -14,13 +14,14 @@
  * along with this program. If not, see
  * <http://www.mongodb.com/licensing/server-side-public-license>.
  */
-package org.graylog.storage.opensearch2;
+package org.graylog.storage.opensearch2.sniffer;
 
 import org.graylog.shaded.opensearch2.org.opensearch.client.Node;
 
 import java.io.IOException;
 import java.util.List;
 
-public interface NodesSniffer {
-    List<Node> sniff(List<Node> nodes) throws IOException;
+public interface SnifferFilter {
+    boolean enabled();
+    List<Node> filterNodes(List<Node> nodes) throws IOException;
 }