From 0c44a7aa649b6465b284b9f584e6d1e914a8c1d0 Mon Sep 17 00:00:00 2001
From: Senni Tan <sennitan@geotab.com>
Date: Tue, 11 Jan 2022 12:16:42 -0500
Subject: [PATCH 1/6] refactor(search): refactor NUM_RETRIES in esindexbuilder
 to be configurable

---
 docker/datahub-gms/Dockerfile                                   | 1 +
 docker/datahub-gms/env/docker.env                               | 2 +-
 docker/docker-compose.consumers.yml                             | 1 +
 docker/monitoring/docker-compose.monitoring.yml                 | 1 +
 docker/quickstart/docker-compose-without-neo4j.quickstart.yml   | 1 +
 docker/quickstart/docker-compose.monitoring.quickstart.yml      | 1 +
 .../search/elasticsearch/indexbuilder/ESIndexBuilder.java       | 2 +-
 7 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/docker/datahub-gms/Dockerfile b/docker/datahub-gms/Dockerfile
index b1ac5fcbb41f97..f0978e1ae2f87d 100644
--- a/docker/datahub-gms/Dockerfile
+++ b/docker/datahub-gms/Dockerfile
@@ -50,6 +50,7 @@ USER datahub
 
 ENV JMX_OPTS=""
 ENV JAVA_OPTS=""
+ENV ES_INDEXBUILDER_NUM_RETRIES="10"
 
 EXPOSE 8080
 
diff --git a/docker/datahub-gms/env/docker.env b/docker/datahub-gms/env/docker.env
index d6bddc9a2775e8..30580078144434 100644
--- a/docker/datahub-gms/env/docker.env
+++ b/docker/datahub-gms/env/docker.env
@@ -16,7 +16,7 @@ JAVA_OPTS=-Xms1g -Xmx1g
 GRAPH_SERVICE_IMPL=neo4j
 ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
 MAE_CONSUMER_ENABLED=true
-MCE_CONSUMER_ENABLED=true
+MCE_CONSUMER_ENABLED=trueES_INDEXBUILDER_NUM_RETRIES=10
 
 # Uncomment to enable Metadata Service Authentication
 # METADATA_SERVICE_AUTH_ENABLED=true
diff --git a/docker/docker-compose.consumers.yml b/docker/docker-compose.consumers.yml
index 3eb8580a61a42f..61537efa4a9bcb 100644
--- a/docker/docker-compose.consumers.yml
+++ b/docker/docker-compose.consumers.yml
@@ -5,6 +5,7 @@ services:
     environment:
       - MAE_CONSUMER_ENABLED=false
       - MCE_CONSUMER_ENABLED=false
+      - ES_INDEXBUILDER_NUM_RETRIES=10
 
   datahub-mae-consumer:
     build:
diff --git a/docker/monitoring/docker-compose.monitoring.yml b/docker/monitoring/docker-compose.monitoring.yml
index 73cc3ec39a2593..dbaf191e375c8c 100644
--- a/docker/monitoring/docker-compose.monitoring.yml
+++ b/docker/monitoring/docker-compose.monitoring.yml
@@ -9,6 +9,7 @@ services:
       - OTEL_EXPORTER_JAEGER_ENDPOINT=http://jaeger-all-in-one:14250
       - OTEL_METRICS_EXPORTER=none
       - OTEL_SERVICE_NAME=datahub-gms
+      - ES_INDEXBUILDER_NUM_RETRIES=10
     ports:
       - "4318"
 
diff --git a/docker/quickstart/docker-compose-without-neo4j.quickstart.yml b/docker/quickstart/docker-compose-without-neo4j.quickstart.yml
index 88ca38c63a6df3..ed3e6e9fad2995 100644
--- a/docker/quickstart/docker-compose-without-neo4j.quickstart.yml
+++ b/docker/quickstart/docker-compose-without-neo4j.quickstart.yml
@@ -60,6 +60,7 @@ services:
     - ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
     - MAE_CONSUMER_ENABLED=true
     - MCE_CONSUMER_ENABLED=true
+    - ES_INDEXBUILDER_NUM_RETRIES=10
     hostname: datahub-gms
     image: linkedin/datahub-gms:${DATAHUB_VERSION:-head}
     ports:
diff --git a/docker/quickstart/docker-compose.monitoring.quickstart.yml b/docker/quickstart/docker-compose.monitoring.quickstart.yml
index 0c9e822e2dfb31..30577f01cae5ff 100644
--- a/docker/quickstart/docker-compose.monitoring.quickstart.yml
+++ b/docker/quickstart/docker-compose.monitoring.quickstart.yml
@@ -7,6 +7,7 @@ services:
     - OTEL_EXPORTER_JAEGER_ENDPOINT=http://jaeger-all-in-one:14250
     - OTEL_METRICS_EXPORTER=none
     - OTEL_SERVICE_NAME=datahub-gms
+    - ES_INDEXBUILDER_NUM_RETRIES=10
     ports:
     - '4318'
   grafana:
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java
index 648e61b6a93c82..fac46c5b1c0771 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java
@@ -40,7 +40,7 @@ public class ESIndexBuilder {
   private final int numShards;
   private final int numReplicas;
 
-  private static final int NUM_RETRIES = 3;
+  private static final int NUM_RETRIES = Integer.parseInt(System.getenv("ES_INDEXBUILDER_NUM_RETRIES"));
   private static final List<String> SETTINGS_TO_COMPARE = ImmutableList.of("number_of_shards", "number_of_replicas");
 
   public void buildIndex(String indexName, Map<String, Object> mappings, Map<String, Object> settings)

From 0c1f27d9711ee0be34d37fab339eda3d26df7f52 Mon Sep 17 00:00:00 2001
From: Senni Tan <sennitan@geotab.com>
Date: Tue, 11 Jan 2022 12:20:06 -0500
Subject: [PATCH 2/6] refactor(search): refactor NUM_RETRIES in esindexbuilder
 to be configurable

---
 docker/datahub-gms/env/docker.env | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docker/datahub-gms/env/docker.env b/docker/datahub-gms/env/docker.env
index 30580078144434..2a732853c5ed88 100644
--- a/docker/datahub-gms/env/docker.env
+++ b/docker/datahub-gms/env/docker.env
@@ -16,7 +16,8 @@ JAVA_OPTS=-Xms1g -Xmx1g
 GRAPH_SERVICE_IMPL=neo4j
 ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
 MAE_CONSUMER_ENABLED=true
-MCE_CONSUMER_ENABLED=trueES_INDEXBUILDER_NUM_RETRIES=10
+MCE_CONSUMER_ENABLED=true
+ES_INDEXBUILDER_NUM_RETRIES=10
 
 # Uncomment to enable Metadata Service Authentication
 # METADATA_SERVICE_AUTH_ENABLED=true

From c6f949481102e07b1157088b4e1ed90d716987f7 Mon Sep 17 00:00:00 2001
From: Senni Tan <sennitan@geotab.com>
Date: Wed, 12 Jan 2022 11:20:10 -0500
Subject: [PATCH 3/6] change the way of refactor

---
 docker/datahub-gms/Dockerfile                                | 1 -
 docker/datahub-gms/env/docker.env                            | 1 -
 docker/docker-compose.consumers.yml                          | 1 -
 docker/monitoring/docker-compose.monitoring.yml              | 1 -
 .../quickstart/docker-compose-without-neo4j.quickstart.yml   | 1 -
 docker/quickstart/docker-compose.monitoring.quickstart.yml   | 1 -
 .../search/elasticsearch/indexbuilder/ESIndexBuilder.java    | 4 ++--
 .../gms/factory/search/ElasticSearchIndexBuilderFactory.java | 5 ++++-
 .../factories/src/main/resources/application.yml             | 1 +
 9 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/docker/datahub-gms/Dockerfile b/docker/datahub-gms/Dockerfile
index f0978e1ae2f87d..b1ac5fcbb41f97 100644
--- a/docker/datahub-gms/Dockerfile
+++ b/docker/datahub-gms/Dockerfile
@@ -50,7 +50,6 @@ USER datahub
 
 ENV JMX_OPTS=""
 ENV JAVA_OPTS=""
-ENV ES_INDEXBUILDER_NUM_RETRIES="10"
 
 EXPOSE 8080
 
diff --git a/docker/datahub-gms/env/docker.env b/docker/datahub-gms/env/docker.env
index 2a732853c5ed88..d6bddc9a2775e8 100644
--- a/docker/datahub-gms/env/docker.env
+++ b/docker/datahub-gms/env/docker.env
@@ -17,7 +17,6 @@ GRAPH_SERVICE_IMPL=neo4j
 ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
 MAE_CONSUMER_ENABLED=true
 MCE_CONSUMER_ENABLED=true
-ES_INDEXBUILDER_NUM_RETRIES=10
 
 # Uncomment to enable Metadata Service Authentication
 # METADATA_SERVICE_AUTH_ENABLED=true
diff --git a/docker/docker-compose.consumers.yml b/docker/docker-compose.consumers.yml
index 61537efa4a9bcb..3eb8580a61a42f 100644
--- a/docker/docker-compose.consumers.yml
+++ b/docker/docker-compose.consumers.yml
@@ -5,7 +5,6 @@ services:
     environment:
       - MAE_CONSUMER_ENABLED=false
       - MCE_CONSUMER_ENABLED=false
-      - ES_INDEXBUILDER_NUM_RETRIES=10
 
   datahub-mae-consumer:
     build:
diff --git a/docker/monitoring/docker-compose.monitoring.yml b/docker/monitoring/docker-compose.monitoring.yml
index dbaf191e375c8c..73cc3ec39a2593 100644
--- a/docker/monitoring/docker-compose.monitoring.yml
+++ b/docker/monitoring/docker-compose.monitoring.yml
@@ -9,7 +9,6 @@ services:
       - OTEL_EXPORTER_JAEGER_ENDPOINT=http://jaeger-all-in-one:14250
       - OTEL_METRICS_EXPORTER=none
       - OTEL_SERVICE_NAME=datahub-gms
-      - ES_INDEXBUILDER_NUM_RETRIES=10
     ports:
       - "4318"
 
diff --git a/docker/quickstart/docker-compose-without-neo4j.quickstart.yml b/docker/quickstart/docker-compose-without-neo4j.quickstart.yml
index ed3e6e9fad2995..88ca38c63a6df3 100644
--- a/docker/quickstart/docker-compose-without-neo4j.quickstart.yml
+++ b/docker/quickstart/docker-compose-without-neo4j.quickstart.yml
@@ -60,7 +60,6 @@ services:
     - ENTITY_REGISTRY_CONFIG_PATH=/datahub/datahub-gms/resources/entity-registry.yml
     - MAE_CONSUMER_ENABLED=true
     - MCE_CONSUMER_ENABLED=true
-    - ES_INDEXBUILDER_NUM_RETRIES=10
     hostname: datahub-gms
     image: linkedin/datahub-gms:${DATAHUB_VERSION:-head}
     ports:
diff --git a/docker/quickstart/docker-compose.monitoring.quickstart.yml b/docker/quickstart/docker-compose.monitoring.quickstart.yml
index 30577f01cae5ff..0c9e822e2dfb31 100644
--- a/docker/quickstart/docker-compose.monitoring.quickstart.yml
+++ b/docker/quickstart/docker-compose.monitoring.quickstart.yml
@@ -7,7 +7,6 @@ services:
     - OTEL_EXPORTER_JAEGER_ENDPOINT=http://jaeger-all-in-one:14250
     - OTEL_METRICS_EXPORTER=none
     - OTEL_SERVICE_NAME=datahub-gms
-    - ES_INDEXBUILDER_NUM_RETRIES=10
     ports:
     - '4318'
   grafana:
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java
index fac46c5b1c0771..4eb1f5312cd215 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/indexbuilder/ESIndexBuilder.java
@@ -39,8 +39,8 @@ public class ESIndexBuilder {
   private final RestHighLevelClient searchClient;
   private final int numShards;
   private final int numReplicas;
+  private final int numRetries;
 
-  private static final int NUM_RETRIES = Integer.parseInt(System.getenv("ES_INDEXBUILDER_NUM_RETRIES"));
   private static final List<String> SETTINGS_TO_COMPARE = ImmutableList.of("number_of_shards", "number_of_replicas");
 
   public void buildIndex(String indexName, Map<String, Object> mappings, Map<String, Object> settings)
@@ -128,7 +128,7 @@ public void buildIndex(String indexName, Map<String, Object> mappings, Map<Strin
     // There can be some delay between the reindex finishing and count being fully up to date, so try multiple times
     long originalCount = 0;
     long reindexedCount = 0;
-    for (int i = 0; i < NUM_RETRIES; i++) {
+    for (int i = 0; i < this.numRetries; i++) {
       // Check if reindex succeeded by comparing document counts
       originalCount = getCount(indexName);
       reindexedCount = getCount(tempIndexName);
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java
index 46562cb54e338e..d43603b328c17d 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java
@@ -28,9 +28,12 @@ public class ElasticSearchIndexBuilderFactory {
   @Value("${elasticsearch.index.numReplicas}")
   private Integer numReplicas;
 
+  @Value("${elasticsearch.index.numRetries}")
+  private Integer numRetries;
+
   @Bean(name = "elasticSearchIndexBuilder")
   @Nonnull
   protected ESIndexBuilder getInstance() {
-    return new ESIndexBuilder(searchClient, numShards, numReplicas);
+    return new ESIndexBuilder(searchClient, numShards, numReplicas, numRetries);
   }
 }
\ No newline at end of file
diff --git a/metadata-service/factories/src/main/resources/application.yml b/metadata-service/factories/src/main/resources/application.yml
index 62faf6ad2bc1c1..2001906ebffaf7 100644
--- a/metadata-service/factories/src/main/resources/application.yml
+++ b/metadata-service/factories/src/main/resources/application.yml
@@ -100,6 +100,7 @@ elasticsearch:
     prefix: ${INDEX_PREFIX:}
     numShards: ${ELASTICSEARCH_NUM_SHARDS_PER_INDEX:1}
     numReplicas: ${ELASTICSEARCH_NUM_REPLICAS_PER_INDEX:1}
+    numRetries: ${ELASTICSEARCH_NUM_REPLICAS_PER_INDEX:3}
     maxArrayLength: ${SEARCH_DOCUMENT_MAX_ARRAY_LENGTH:1000}
 
 # TODO: Kafka topic convention

From 6a3d75fd78c90369a2dfc4abf0fb68b5a597e007 Mon Sep 17 00:00:00 2001
From: Senni Tan <sennitan@geotab.com>
Date: Wed, 12 Jan 2022 11:29:26 -0500
Subject: [PATCH 4/6] add env var in yml file

---
 metadata-service/factories/src/main/resources/application.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metadata-service/factories/src/main/resources/application.yml b/metadata-service/factories/src/main/resources/application.yml
index 2001906ebffaf7..44f1190438b932 100644
--- a/metadata-service/factories/src/main/resources/application.yml
+++ b/metadata-service/factories/src/main/resources/application.yml
@@ -100,7 +100,7 @@ elasticsearch:
     prefix: ${INDEX_PREFIX:}
     numShards: ${ELASTICSEARCH_NUM_SHARDS_PER_INDEX:1}
     numReplicas: ${ELASTICSEARCH_NUM_REPLICAS_PER_INDEX:1}
-    numRetries: ${ELASTICSEARCH_NUM_REPLICAS_PER_INDEX:3}
+    numRetries: ${ELASTICSEARCH_NUM_RETRIES_PER_INDEX:3}
     maxArrayLength: ${SEARCH_DOCUMENT_MAX_ARRAY_LENGTH:1000}
 
 # TODO: Kafka topic convention

From 39ce4339c1405a23b3a9c3f7323d126b81d44f96 Mon Sep 17 00:00:00 2001
From: Senni Tan <sennitan@geotab.com>
Date: Thu, 13 Jan 2022 10:08:06 -0500
Subject: [PATCH 5/6] rename the env variable

---
 metadata-service/factories/src/main/resources/application.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metadata-service/factories/src/main/resources/application.yml b/metadata-service/factories/src/main/resources/application.yml
index 44f1190438b932..c42333d573bb9d 100644
--- a/metadata-service/factories/src/main/resources/application.yml
+++ b/metadata-service/factories/src/main/resources/application.yml
@@ -100,7 +100,7 @@ elasticsearch:
     prefix: ${INDEX_PREFIX:}
     numShards: ${ELASTICSEARCH_NUM_SHARDS_PER_INDEX:1}
     numReplicas: ${ELASTICSEARCH_NUM_REPLICAS_PER_INDEX:1}
-    numRetries: ${ELASTICSEARCH_NUM_RETRIES_PER_INDEX:3}
+    numRetries: ${ELASTICSEARCH_INDEX_BUILDER_NUM_RETRIES :3}
     maxArrayLength: ${SEARCH_DOCUMENT_MAX_ARRAY_LENGTH:1000}
 
 # TODO: Kafka topic convention

From d2d993d1bb5727832ab226bb650b82807f9ce412 Mon Sep 17 00:00:00 2001
From: Senni Tan <sennitan@geotab.com>
Date: Fri, 14 Jan 2022 10:12:36 -0500
Subject: [PATCH 6/6] fix test fail on line 100 in
 ElasticSearchServiceTest.java

---
 .../metadata/search/elasticsearch/ElasticSearchServiceTest.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchServiceTest.java
index f5813fd10cf02c..105742b434216b 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchServiceTest.java
@@ -97,7 +97,7 @@ public static BulkProcessor getBulkProcessor(RestHighLevelClient searchClient) {
   }
 
   public static ESIndexBuilder getIndexBuilder(RestHighLevelClient searchClient) {
-    return new ESIndexBuilder(searchClient, 1, 1);
+    return new ESIndexBuilder(searchClient, 1, 1, 3);
   }
 
   @Nonnull