From fa9fca9cbbd73cbb8971bc20e3b1e71902db0159 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    | 24 ++++++++++++++++---
 1 file changed, 21 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..1490f4b9f 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,24 @@ public SparkSink(
     }
   }
 
-  private Config loadConfig() {
-    return ConfigFactory.load(pickClassLoader());
+  private Config loadConfig(Properties properties) {
+    ClassLoader classLoader = pickClassLoader();
+    String resolverClassName = properties.getProperty("resolver");
+    if (resolverClassName == null) {
+      return ConfigFactory.load(classLoader);
+    }
+    ConfigResolver configResolver = getConfigResolver(resolverClassName, classLoader);
+    ConfigResolveOptions resolveOptions = ConfigResolveOptions.defaults().appendResolver(configResolver);
+    return ConfigFactory.load(classLoader, resolveOptions);
+  }
+
+  private static ConfigResolver getConfigResolver(String resolverClassName, ClassLoader classLoader) {
+    try {
+      LOGGER.info("Creating ConfigResolver instance for {}", resolverClassName);
+      return (ConfigResolver) classLoader.loadClass(resolverClassName).getDeclaredConstructor().newInstance();
+    } catch (Exception e) {
+      throw new IllegalArgumentException(String.format("Error creating instance for %s", resolverClassName), e);
+    }
   }
 
   @SuppressWarnings("PMD.UseProperClassLoader")