From e3406debf7ef4f658583ecaa19d8c4203461d311 Mon Sep 17 00:00:00 2001
From: David Leifker <david.leifker@acryl.io>
Date: Mon, 24 Jul 2023 16:56:05 -0500
Subject: [PATCH] config(ebean): add ebean retry configuration

---
 .../metadata/entity/EntityServiceImpl.java     | 18 +++++++++++++++++-
 .../factory/entity/EntityServiceFactory.java   |  6 +++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java
index 32c77b66679f9b..c7a9895992d90b 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java
@@ -146,7 +146,7 @@ public class EntityServiceImpl implements EntityService {
     OBJECT_MAPPER.getFactory().setStreamReadConstraints(StreamReadConstraints.builder().maxStringLength(maxSize).build());
   }
 
-  private static final int DEFAULT_MAX_TRANSACTION_RETRY = 3;
+  private static final int DEFAULT_MAX_TRANSACTION_RETRY = 4;
 
   protected final AspectDao _aspectDao;
   private final EventProducer _producer;
@@ -163,6 +163,8 @@ public class EntityServiceImpl implements EntityService {
   // TODO(iprentic): Move this to a common utils location once used in other places
   private static final String DELIMITER_SEPARATOR = "␟";
 
+  private final Integer ebeanMaxTransactionRetry;
+
   public EntityServiceImpl(
       @Nonnull final AspectDao aspectDao,
       @Nonnull final EventProducer producer,
@@ -170,6 +172,17 @@ public EntityServiceImpl(
       final boolean alwaysEmitChangeLog,
       final UpdateIndicesService updateIndicesService,
       final PreProcessHooks preProcessHooks) {
+    this(aspectDao, producer, entityRegistry, alwaysEmitChangeLog, updateIndicesService, preProcessHooks, DEFAULT_MAX_TRANSACTION_RETRY);
+  }
+
+  public EntityServiceImpl(
+          @Nonnull final AspectDao aspectDao,
+          @Nonnull final EventProducer producer,
+          @Nonnull final EntityRegistry entityRegistry,
+          final boolean alwaysEmitChangeLog,
+          final UpdateIndicesService updateIndicesService,
+          final PreProcessHooks preProcessHooks,
+          final Integer retry) {
 
     _aspectDao = aspectDao;
     _producer = producer;
@@ -178,8 +191,11 @@ public EntityServiceImpl(
     _alwaysEmitChangeLog = alwaysEmitChangeLog;
     _updateIndicesService = updateIndicesService;
     _preProcessHooks = preProcessHooks;
+    ebeanMaxTransactionRetry = retry != null ? retry : DEFAULT_MAX_TRANSACTION_RETRY;
   }
 
+
+
   /**
    * Retrieves the latest aspects corresponding to a batch of {@link Urn}s based on a provided
    * set of aspect names.
diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityServiceFactory.java
index b8edd6918dc6c7..63ac26f124fef5 100644
--- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityServiceFactory.java
+++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityServiceFactory.java
@@ -14,6 +14,7 @@
 import org.apache.avro.generic.IndexedRecord;
 import org.apache.kafka.clients.producer.Producer;
 import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.DependsOn;
@@ -24,6 +25,9 @@
 @Configuration
 public class EntityServiceFactory {
 
+  @Value("${EBEAN_MAX_TRANSACTION_RETRY:4}")
+  private Integer _ebeanMaxTransactionRetry;
+
   @Bean(name = "entityService")
   @DependsOn({"entityAspectDao", "kafkaEventProducer", "kafkaHealthChecker",
           TopicConventionFactory.TOPIC_CONVENTION_BEAN, "entityRegistry"})
@@ -40,6 +44,6 @@ protected EntityService createInstance(
     final KafkaEventProducer eventProducer = new KafkaEventProducer(producer, convention, kafkaHealthChecker);
     FeatureFlags featureFlags = configurationProvider.getFeatureFlags();
     return new EntityServiceImpl(aspectDao, eventProducer, entityRegistry,
-        featureFlags.isAlwaysEmitChangeLog(), updateIndicesService, featureFlags.getPreProcessHooks());
+        featureFlags.isAlwaysEmitChangeLog(), updateIndicesService, featureFlags.getPreProcessHooks(), _ebeanMaxTransactionRetry);
   }
 }