David leifker/elasticsearch optimization ext

docker/quickstart/docker-compose-without-neo4j-m1.quickstart.yml

@@ -74,6 +74,7 @@ services:
     depends_on:
     - mysql
     environment:
+    - BUILD_INDICES_HISTORY_KAFKA_CONSUMER_GROUP_ID=generic-bihe-consumer-job-client-gms
     - EBEAN_DATASOURCE_USERNAME=datahub
     - EBEAN_DATASOURCE_PASSWORD=datahub
     - EBEAN_DATASOURCE_HOST=mysql:3306

entity-registry/src/main/java/com/linkedin/metadata/models/SearchableFieldSpecExtractor.java

@@ -44,7 +44,7 @@ public class SearchableFieldSpecExtractor implements SchemaVisitor {
             Map.of(
                     "enableAutocomplete", "false",
                     "fieldType", "URN",
-                    "boostScore", "0.4"
+                    "boostScore", "0.1"
             )
     );
   }

metadata-io/src/main/java/com/linkedin/metadata/config/BuildIndicesConfiguration.java

@@ -11,5 +11,6 @@ public class BuildIndicesConfiguration {
   private String backOffFactor;
   private boolean waitForBuildIndices;
   private boolean cloneIndices;
+  private boolean allowDocCountMismatch;
 
 }

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java

@@ -11,6 +11,7 @@
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 
+import com.linkedin.metadata.config.ElasticSearchConfiguration;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -67,6 +68,8 @@ public class ESIndexBuilder {
   @Getter
   private final boolean enableIndexMappingsReindex;
 
+  private final ElasticSearchConfiguration elasticSearchConfiguration;
+
 
   public ReindexConfig buildReindexState(String indexName, Map<String, Object> mappings, Map<String, Object> settings) throws IOException {
     ReindexConfig.ReindexConfigBuilder builder = ReindexConfig.builder()
@@ -205,12 +208,19 @@ private void reindex(String indexName, ReindexConfig indexState) throws IOExcept
     }
 
     if (originalCount != reindexedCount) {
-      log.info("Post-reindex document count is different, source_doc_count: {} reindex_doc_count: {}", originalCount,
-          reindexedCount);
-      diff(indexName, tempIndexName, Math.max(originalCount, reindexedCount));
-      searchClient.indices().delete(new DeleteIndexRequest().indices(tempIndexName), RequestOptions.DEFAULT);
-      throw new RuntimeException(String.format("Reindex from %s to %s failed. Document count %s != %s", indexName, tempIndexName,
-              originalCount, reindexedCount));
+      if (elasticSearchConfiguration.getBuildIndices().isAllowDocCountMismatch()
+              && elasticSearchConfiguration.getBuildIndices().isCloneIndices()) {
+        log.warn("Index: {} - Post-reindex document count is different, source_doc_count: {} reindex_doc_count: {}\n"
+                + "This condition is explicitly ALLOWED, please refer to latest clone if original index is required.",
+                indexName, originalCount, reindexedCount);
+      } else {
+        log.error("Index: {} - Post-reindex document count is different, source_doc_count: {} reindex_doc_count: {}",
+                indexName, originalCount, reindexedCount);
+        diff(indexName, tempIndexName, Math.max(originalCount, reindexedCount));
+        searchClient.indices().delete(new DeleteIndexRequest().indices(tempIndexName), RequestOptions.DEFAULT);
+        throw new RuntimeException(String.format("Reindex from %s to %s failed. Document count %s != %s", indexName, tempIndexName,
+                originalCount, reindexedCount));
+      }
     }
 
     log.info("Reindex from {} to {} succeeded", indexName, tempIndexName);

metadata-io/src/test/java/com/linkedin/metadata/ESSampleDataFixture.java

@@ -2,6 +2,7 @@
 
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.client.JavaEntityClient;
+import com.linkedin.metadata.config.ElasticSearchConfiguration;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.search.SearchService;
@@ -67,7 +68,8 @@ protected EntityIndexBuilders entityIndexBuilders(
             @Qualifier("sampleDataIndexConvention") IndexConvention indexConvention
     ) {
         ESIndexBuilder indexBuilder = new ESIndexBuilder(_searchClient, 1, 0, 1,
-                1, Map.of(), true, false);
+                1, Map.of(), true, false,
+                new ElasticSearchConfiguration());
         SettingsBuilder settingsBuilder = new SettingsBuilder(null);
         return new EntityIndexBuilders(indexBuilder, entityRegistry, indexConvention, settingsBuilder);
     }

metadata-io/src/test/java/com/linkedin/metadata/ESSearchLineageFixture.java

@@ -2,6 +2,7 @@
 
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.client.JavaEntityClient;
+import com.linkedin.metadata.config.ElasticSearchConfiguration;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.graph.elastic.ESGraphQueryDAO;
 import com.linkedin.metadata.graph.elastic.ESGraphWriteDAO;
@@ -72,7 +73,8 @@ protected EntityIndexBuilders entityIndexBuilders(
             @Qualifier("searchLineageIndexConvention") IndexConvention indexConvention
     ) {
         ESIndexBuilder indexBuilder = new ESIndexBuilder(_searchClient, 1, 0, 1,
-                1, Map.of(), true, false);
+                1, Map.of(), true, false,
+                new ElasticSearchConfiguration());
         SettingsBuilder settingsBuilder = new SettingsBuilder(null);
         return new EntityIndexBuilders(indexBuilder, entityRegistry, indexConvention, settingsBuilder);
     }
@@ -93,7 +95,8 @@ protected ElasticSearchService entitySearchService(
     @Nonnull
     protected ESIndexBuilder esIndexBuilder() {
         return new ESIndexBuilder(_searchClient, 1, 1, 1, 1, Map.of(),
-                true, true);
+                true, true,
+                new ElasticSearchConfiguration());
     }
 
     @Bean(name = "searchLineageGraphService")

metadata-io/src/test/java/com/linkedin/metadata/ESTestConfiguration.java

@@ -1,5 +1,6 @@
 package com.linkedin.metadata;
 
+import com.linkedin.metadata.config.ElasticSearchConfiguration;
 import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.models.registry.EntityRegistryException;
@@ -90,7 +91,8 @@ public ESBulkProcessor getBulkProcessor(@Qualifier("elasticSearchRestHighLevelCl
     @Nonnull
     protected ESIndexBuilder getIndexBuilder(@Qualifier("elasticSearchRestHighLevelClient") RestHighLevelClient searchClient) {
         return new ESIndexBuilder(searchClient, 1, 1, 3, 1, Map.of(),
-                false, false);
+                false, false,
+                new ElasticSearchConfiguration());
     }
 
     @Bean(name = "entityRegistry")

metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilderTest.java

@@ -2,6 +2,7 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.linkedin.metadata.ESTestConfiguration;
+import com.linkedin.metadata.config.ElasticSearchConfiguration;
 import com.linkedin.metadata.systemmetadata.SystemMetadataMappingsBuilder;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.action.admin.indices.alias.get.GetAliasesRequest;
@@ -44,7 +45,8 @@ public class ESIndexBuilderTest extends AbstractTestNGSpringContextTests {
     public void setup() {
         _indexClient = _searchClient.indices();
         testDefaultBuilder = new ESIndexBuilder(_searchClient, 1, 0, 0,
-                0, Map.of(), false, false);
+                0, Map.of(), false, false,
+                new ElasticSearchConfiguration());
     }
 
     @BeforeMethod
@@ -74,7 +76,8 @@ public static GetIndexResponse getTestIndex() throws IOException {
     @Test
     public void testESIndexBuilderCreation() throws Exception {
         ESIndexBuilder customIndexBuilder = new ESIndexBuilder(_searchClient, 2, 0, 1,
-                0, Map.of(), false, false);
+                0, Map.of(), false, false,
+                new ElasticSearchConfiguration());
         customIndexBuilder.buildIndex(TEST_INDEX_NAME, Map.of(), Map.of());
         GetIndexResponse resp = getTestIndex();
 
@@ -86,7 +89,8 @@ public void testESIndexBuilderCreation() throws Exception {
     @Test
     public void testMappingReindex() throws Exception {
         ESIndexBuilder enabledMappingReindex = new ESIndexBuilder(_searchClient, 1, 0, 0,
-                0, Map.of(), false, true);
+                0, Map.of(), false, true,
+                new ElasticSearchConfiguration());
 
         // No mappings
         enabledMappingReindex.buildIndex(TEST_INDEX_NAME, Map.of(), Map.of());
@@ -130,7 +134,8 @@ public void testSettingsNumberOfShardsReindex() throws Exception {
                 testDefaultBuilder.getNumRetries(),
                 testDefaultBuilder.getRefreshIntervalSeconds(),
                 Map.of(),
-                true, false);
+                true, false,
+                new ElasticSearchConfiguration());
 
         // add new shard setting
         changedShardBuilder.buildIndex(TEST_INDEX_NAME, Map.of(), Map.of());
@@ -156,28 +161,32 @@ public void testSettingsNoReindex() throws Exception {
                         testDefaultBuilder.getNumRetries(),
                         testDefaultBuilder.getRefreshIntervalSeconds(),
                         Map.of(),
-                        true, false),
+                        true, false,
+                        new ElasticSearchConfiguration()),
                 new ESIndexBuilder(_searchClient,
                         testDefaultBuilder.getNumShards(),
                         testDefaultBuilder.getNumReplicas(),
                         testDefaultBuilder.getNumRetries(),
                         testDefaultBuilder.getRefreshIntervalSeconds() + 10,
                         Map.of(),
-                        true, false),
+                        true, false,
+                        new ElasticSearchConfiguration()),
                new ESIndexBuilder(_searchClient,
                                 testDefaultBuilder.getNumShards() + 1,
                                 testDefaultBuilder.getNumReplicas(),
                                 testDefaultBuilder.getNumRetries(),
                                 testDefaultBuilder.getRefreshIntervalSeconds(),
                                 Map.of(),
-                                false, false),
+                                false, false,
+                       new ElasticSearchConfiguration()),
                 new ESIndexBuilder(_searchClient,
                         testDefaultBuilder.getNumShards(),
                         testDefaultBuilder.getNumReplicas() + 1,
                         testDefaultBuilder.getNumRetries(),
                         testDefaultBuilder.getRefreshIntervalSeconds(),
                         Map.of(),
-                        false, false)
+                        false, false,
+                        new ElasticSearchConfiguration())
         );
 
         for (ESIndexBuilder builder : noReindexBuilders) {

metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/elasticsearch/indices/BuildIndicesKafkaListener.java

@@ -11,6 +11,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
 
 import joptsimple.internal.Strings;
 import lombok.RequiredArgsConstructor;
@@ -26,6 +27,7 @@
 import org.springframework.kafka.config.KafkaListenerEndpointRegistry;
 import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
 import org.springframework.kafka.listener.ConsumerSeekAware;
+import org.springframework.kafka.listener.MessageListenerContainer;
 import org.springframework.stereotype.Component;
 
 // We don't disable this on GMS since we want GMS to also wait until the indices are ready to read in case of
@@ -99,11 +101,13 @@ public void waitForUpdate() {
     for (int i = 0; i < maxBackOffs; i++) {
       if (isUpdated.get()) {
         log.debug("Finished waiting for updated indices.");
-        String consumerId = Strings.join(List.of(CONSUMER_GROUP, SUFFIX), "");
         try {
-          registry.getListenerContainer(consumerId).stop();
+          log.info("Containers: {}", registry.getListenerContainers().stream()
+                  .map(MessageListenerContainer::getListenerId)
+                  .collect(Collectors.toList()));
+          registry.getListenerContainer(consumerGroup).stop();
         } catch (NullPointerException e) {
-          log.error("Expected consumer `{}` to shutdown.", consumerId);
+          log.error("Expected consumer `{}` to shutdown.", consumerGroup);
         }
         return;
       }

metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java

@@ -5,6 +5,7 @@
 import com.linkedin.gms.factory.common.IndexConventionFactory;
 import com.linkedin.gms.factory.common.RestHighLevelClientFactory;
 import com.linkedin.gms.factory.spring.YamlPropertySourceFactory;
+import com.linkedin.metadata.config.ElasticSearchConfiguration;
 import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -28,7 +29,7 @@
 
 
 @Configuration
-@Import({RestHighLevelClientFactory.class, IndexConventionFactory.class})
+@Import({RestHighLevelClientFactory.class, IndexConventionFactory.class, ElasticSearchConfiguration.class})
 @PropertySource(value = "classpath:/application.yml", factory = YamlPropertySourceFactory.class)
 public class ElasticSearchIndexBuilderFactory {
 
@@ -76,9 +77,10 @@ protected Map<String, Map<String, String>> getIndexSettingsOverrides(
   @Bean(name = "elasticSearchIndexBuilder")
   @Nonnull
   protected ESIndexBuilder getInstance(
-          @Qualifier("elasticSearchIndexSettingsOverrides") Map<String, Map<String, String>> overrides) {
+          @Qualifier("elasticSearchIndexSettingsOverrides") Map<String, Map<String, String>> overrides,
+          final ElasticSearchConfiguration elasticSearchConfiguration) {
     return new ESIndexBuilder(searchClient, numShards, numReplicas, numRetries, refreshIntervalSeconds, overrides,
-            enableSettingsReindex, enableMappingsReindex);
+            enableSettingsReindex, enableMappingsReindex, elasticSearchConfiguration);
   }
 
   @Nonnull

metadata-service/factories/src/main/resources/application.yml

@@ -165,6 +165,7 @@ elasticsearch:
     settingsOverrides: ${ELASTICSEARCH_INDEX_BUILDER_SETTINGS_OVERRIDES:#{null}}
     entitySettingsOverrides: ${ELASTICSEARCH_INDEX_BUILDER_ENTITY_SETTINGS_OVERRIDES:#{null}}
   buildIndices:
+    allowDocCountMismatch: ${ELASTICSEARCH_BUILD_INDICES_ALLOW_DOC_COUNT_MISMATCH:false} # when cloneIndices is also enabled
     cloneIndices: ${ELASTICSEARCH_BUILD_INDICES_CLONE_INDICES:true}
     initialBackOffMs: ${ELASTICSEARCH_BUILD_INDICES_INITIAL_BACK_OFF_MILLIS:5000}
     maxBackOffs: ${ELASTICSEARCH_BUILD_INDICES_MAX_BACK_OFFS:5}