feat(pluginRegistry): add configuration to reduce runnable frequency (#…

…9990)
datahub-project · Mar 6, 2024 · eed0abe · eed0abe
1 parent d476f5f
commit eed0abe
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 11 deletions.
diff --git a/...istry/src/main/java/com/linkedin/metadata/models/registry/PluginEntityRegistryLoader.java b/...istry/src/main/java/com/linkedin/metadata/models/registry/PluginEntityRegistryLoader.java
@@ -28,6 +28,7 @@ public class PluginEntityRegistryLoader {
   private static int _MAXLOADFAILURES = 5;
   private final Boolean scanningEnabled;
   private final String pluginDirectory;
+  private final int loadDelaySeconds;
   // Registry Name -> Registry Version -> (Registry, LoadResult)
   private final Map<String, Map<ComparableVersion, Pair<EntityRegistry, EntityRegistryLoadResult>>>
       patchRegistries;
@@ -38,7 +39,7 @@ public class PluginEntityRegistryLoader {
   private boolean booted = false;
   private final ScheduledExecutorService executorService = Executors.newScheduledThreadPool(1);
 
-  public PluginEntityRegistryLoader(String pluginDirectory) {
+  public PluginEntityRegistryLoader(String pluginDirectory, int loadDelaySeconds) {
     File directory = new File(pluginDirectory);
     if (!directory.exists() || !directory.isDirectory()) {
       log.warn(
@@ -50,6 +51,7 @@ public PluginEntityRegistryLoader(String pluginDirectory) {
     }
     this.pluginDirectory = pluginDirectory;
     this.patchRegistries = new HashMap<>();
+    this.loadDelaySeconds = loadDelaySeconds;
   }
 
   public Map<String, Map<ComparableVersion, Pair<EntityRegistry, EntityRegistryLoadResult>>>
@@ -133,7 +135,7 @@ public PluginEntityRegistryLoader start(boolean waitForInitialization)
           }
         },
         0,
-        5,
+        loadDelaySeconds,
         TimeUnit.SECONDS);
     started = true;
     if (waitForInitialization) {

diff --git a/...y/src/test/java/com/linkedin/metadata/models/registry/PluginEntityRegistryLoaderTest.java b/...y/src/test/java/com/linkedin/metadata/models/registry/PluginEntityRegistryLoaderTest.java
@@ -76,7 +76,7 @@ public AspectTemplateEngine getAspectTemplateEngine() {
 
     MergedEntityRegistry configEntityRegistry = new MergedEntityRegistry(baseEntityRegistry);
     PluginEntityRegistryLoader pluginEntityRegistryLoader =
-        new PluginEntityRegistryLoader(TestConstants.BASE_DIRECTORY)
+        new PluginEntityRegistryLoader(TestConstants.BASE_DIRECTORY, 60)
             .withBaseRegistry(configEntityRegistry)
             .start(true);
     assertEquals(pluginEntityRegistryLoader.getPatchRegistries().size(), 1);
@@ -170,7 +170,7 @@ public void testEntityRegistryWithGoodBase() throws FileNotFoundException, Inter
 
     MergedEntityRegistry mergedEntityRegistry = new MergedEntityRegistry(getBaseEntityRegistry());
     PluginEntityRegistryLoader pluginEntityRegistryLoader =
-        new PluginEntityRegistryLoader(BASE_DIRECTORY)
+        new PluginEntityRegistryLoader(BASE_DIRECTORY, 60)
             .withBaseRegistry(mergedEntityRegistry)
             .start(true);
     assertEquals(pluginEntityRegistryLoader.getPatchRegistries().size(), 1);
@@ -215,7 +215,7 @@ public void testEntityRegistryVersioning() throws InterruptedException {
     String multiversionPluginDir = "src/test_plugins/";
 
     PluginEntityRegistryLoader pluginEntityRegistryLoader =
-        new PluginEntityRegistryLoader(multiversionPluginDir)
+        new PluginEntityRegistryLoader(multiversionPluginDir, 60)
             .withBaseRegistry(mergedEntityRegistry)
             .start(true);
     Map<String, Map<ComparableVersion, Pair<EntityRegistry, EntityRegistryLoadResult>>>

diff --git a/...uration/src/main/java/com/linkedin/metadata/config/EntityRegistryPluginConfiguration.java b/...uration/src/main/java/com/linkedin/metadata/config/EntityRegistryPluginConfiguration.java
@@ -5,4 +5,5 @@
 @Data
 public class EntityRegistryPluginConfiguration {
   String path;
+  int loadDelaySeconds;
 }
diff --git a/metadata-service/configuration/src/main/resources/application.yml b/metadata-service/configuration/src/main/resources/application.yml
@@ -84,6 +84,7 @@ datahub:
     pluginSecurityMode: ${PLUGIN_SECURITY_MODE:RESTRICTED} # Possible value RESTRICTED or LENIENT, default to RESTRICTED
     entityRegistry:
       path: ${ENTITY_REGISTRY_PLUGIN_PATH:/etc/datahub/plugins/models}
+      loadDelaySeconds: ${ENTITY_REGISTRY_PLUGIN_LOAD_DELAY_SECONDS:60} # Rate at which plugin runnable executes, will run every N seconds.
     retention:
       path: ${RETENTION_PLUGIN_PATH:/etc/datahub/plugins/retention}
     auth:

diff --git a/...es/src/main/java/com/linkedin/gms/factory/entityregistry/PluginEntityRegistryFactory.java b/...es/src/main/java/com/linkedin/gms/factory/entityregistry/PluginEntityRegistryFactory.java
@@ -1,26 +1,28 @@
 package com.linkedin.gms.factory.entityregistry;
 
+import com.linkedin.gms.factory.config.ConfigurationProvider;
+import com.linkedin.metadata.config.EntityRegistryPluginConfiguration;
 import com.linkedin.metadata.models.registry.PluginEntityRegistryLoader;
-import com.linkedin.metadata.spring.YamlPropertySourceFactory;
 import java.io.FileNotFoundException;
 import java.net.MalformedURLException;
 import javax.annotation.Nonnull;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.context.annotation.PropertySource;
 
 @Configuration
-@PropertySource(value = "classpath:/application.yml", factory = YamlPropertySourceFactory.class)
 public class PluginEntityRegistryFactory {
 
   @Value("${datahub.plugin.entityRegistry.path}")
   private String pluginRegistryPath;
 
   @Bean(name = "pluginEntityRegistry")
   @Nonnull
-  protected PluginEntityRegistryLoader getInstance()
+  protected PluginEntityRegistryLoader getInstance(ConfigurationProvider configurationProvider)
       throws FileNotFoundException, MalformedURLException {
-    return new PluginEntityRegistryLoader(pluginRegistryPath);
+    EntityRegistryPluginConfiguration pluginConfiguration =
+        configurationProvider.getDatahub().getPlugin().getEntityRegistry();
+    return new PluginEntityRegistryLoader(
+        pluginConfiguration.getPath(), pluginConfiguration.getLoadDelaySeconds());
   }
 }
diff --git a/...ervlet/src/test/java/io/datahubproject/openapi/config/OpenAPIEntityTestConfiguration.java b/...ervlet/src/test/java/io/datahubproject/openapi/config/OpenAPIEntityTestConfiguration.java
@@ -97,7 +97,7 @@ public EntityRegistry entityRegistry() throws EntityRegistryException, Interrupt
       dependency.
     */
     PluginEntityRegistryLoader custom =
-        new PluginEntityRegistryLoader(getClass().getResource("/custom-model").getFile());
+        new PluginEntityRegistryLoader(getClass().getResource("/custom-model").getFile(), 60);
 
     ConfigEntityRegistry standard =
         new ConfigEntityRegistry(