From 68c8949d2732f416a49a7c07514dc2b057e66734 Mon Sep 17 00:00:00 2001
From: John Zhuge <jzhuge@apache.org>
Date: Mon, 29 Jul 2024 20:37:02 -0700
Subject: [PATCH] spark: load config resolver

---
 .../netflix/spectator/spark/SparkSink.java    | 33 +++++++++++++++++--
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/spectator-ext-spark/src/main/java/com/netflix/spectator/spark/SparkSink.java b/spectator-ext-spark/src/main/java/com/netflix/spectator/spark/SparkSink.java
index bf966db5e..c6c4f2880 100644
--- a/spectator-ext-spark/src/main/java/com/netflix/spectator/spark/SparkSink.java
+++ b/spectator-ext-spark/src/main/java/com/netflix/spectator/spark/SparkSink.java
@@ -23,6 +23,8 @@
 import com.netflix.spectator.sidecar.SidecarRegistry;
 import com.typesafe.config.Config;
 import com.typesafe.config.ConfigFactory;
+import com.typesafe.config.ConfigResolveOptions;
+import com.typesafe.config.ConfigResolver;
 import org.apache.spark.metrics.sink.Sink;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -56,7 +58,7 @@ public SparkSink(
       Properties properties,
       MetricRegistry registry,
       org.apache.spark.SecurityManager manager) {
-    final Config config = loadConfig();
+    final Config config = loadConfig(properties);
     sidecarRegistry = new SidecarRegistry(
         Clock.SYSTEM, new SpectatorConfig(config.getConfig("spectator.spark.sidecar")));
     reporter = SpectatorReporter.forRegistry(registry)
@@ -76,8 +78,33 @@ public SparkSink(
     }
   }
 
-  private Config loadConfig() {
-    return ConfigFactory.load(pickClassLoader());
+  private Config loadConfig(Properties properties) {
+    ClassLoader classLoader = pickClassLoader();
+    ConfigResolver configResolver = getConfigResolver(properties);
+    if (configResolver != null) {
+      ConfigResolveOptions resolveOptions = ConfigResolveOptions.defaults().appendResolver(configResolver);
+      return ConfigFactory.load(classLoader, resolveOptions);
+    } else {
+      return ConfigFactory.load(classLoader);
+    }
+  }
+
+  private static ConfigResolver getConfigResolver(Properties properties) {
+    final String resolverClassName = properties.getProperty("resolver.class");
+    if (resolverClassName == null) {
+      return null;
+    }
+    try {
+      Class<?> clazz = Class.forName(resolverClassName);
+      ConfigResolver resolver = (ConfigResolver) clazz.getDeclaredConstructor().newInstance();
+      LOGGER.info("Created ConfigResolver instance for {}", resolverClassName);
+      return resolver;
+    } catch (ClassNotFoundException e) {
+      LOGGER.error("Class {} not found", resolverClassName, e);
+    } catch (Exception e) {
+      LOGGER.error("Error creating instance for {}", resolverClassName, e);
+    }
+    return null;
   }
 
   @SuppressWarnings("PMD.UseProperClassLoader")