From b2539b2b26c18952aba2b1af76491e86fbdf7110 Mon Sep 17 00:00:00 2001
From: Gregor Zeitlinger <gregor.zeitlinger@grafana.com>
Date: Fri, 9 Feb 2024 22:33:34 +0100
Subject: [PATCH] add ComponentLoader to support more auto configuration
 scenarios, e.g. spring boot

---
 .../sdk/autoconfigure/ComponentLoader.java    |  6 ++++++
 .../sdk/autoconfigure/internal/SpiHelper.java | 21 +++++--------------
 2 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/sdk-extensions/autoconfigure/src/main/java/io/opentelemetry/sdk/autoconfigure/ComponentLoader.java b/sdk-extensions/autoconfigure/src/main/java/io/opentelemetry/sdk/autoconfigure/ComponentLoader.java
index 712b7e20d54..6bbb92c2abf 100644
--- a/sdk-extensions/autoconfigure/src/main/java/io/opentelemetry/sdk/autoconfigure/ComponentLoader.java
+++ b/sdk-extensions/autoconfigure/src/main/java/io/opentelemetry/sdk/autoconfigure/ComponentLoader.java
@@ -11,6 +11,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 
@@ -18,6 +19,11 @@
 public interface ComponentLoader {
   <T> Iterable<T> load(Class<T> spiClass);
 
+  default <T> Optional<T> loadOptional(Class<T> spiClass) {
+    Iterable<T> iterable = load(spiClass);
+    return StreamSupport.stream(iterable.spliterator(), false).findFirst();
+  }
+
   /**
    * Load implementations of an ordered SPI (i.e. implements {@link Ordered}).
    *
diff --git a/sdk-extensions/autoconfigure/src/main/java/io/opentelemetry/sdk/autoconfigure/internal/SpiHelper.java b/sdk-extensions/autoconfigure/src/main/java/io/opentelemetry/sdk/autoconfigure/internal/SpiHelper.java
index aa6abd7cc0d..9d28aba18cd 100644
--- a/sdk-extensions/autoconfigure/src/main/java/io/opentelemetry/sdk/autoconfigure/internal/SpiHelper.java
+++ b/sdk-extensions/autoconfigure/src/main/java/io/opentelemetry/sdk/autoconfigure/internal/SpiHelper.java
@@ -68,13 +68,7 @@ public <T, S extends ConfigurableProvider> NamedSpiManager<T> loadConfigurable(
       // both the provider and the result may have a listener
       maybeAddListener(provider);
 
-      nameToProvider.put(
-          name,
-          () -> {
-            T result = getConfigurable.apply(provider, config);
-            maybeAddListener(result);
-            return result;
-          });
+      nameToProvider.put(name, () -> maybeAddListener(getConfigurable.apply(provider, config)));
     }
     return NamedSpiManager.create(nameToProvider);
   }
@@ -111,20 +105,20 @@ public <T> List<T> load(Class<T> spiClass) {
   private <T> List<T> init(Iterable<T> components) {
     List<T> result = new ArrayList<>();
     for (T service : components) {
-      maybeAddListener(service);
-      result.add(service);
+      result.add(maybeAddListener(service));
     }
     return result;
   }
 
   public <T> Optional<T> loadOptional(Class<T> spiClass) {
-    return load(spiClass).stream().findFirst();
+    return componentLoader.loadOptional(spiClass).map(this::maybeAddListener);
   }
 
-  private void maybeAddListener(Object object) {
+  private <T> T maybeAddListener(T object) {
     if (object instanceof AutoConfigureListener) {
       listeners.add((AutoConfigureListener) object);
     }
+    return object;
   }
 
   /** Return the set of SPIs loaded which implement {@link AutoConfigureListener}. */
@@ -132,11 +126,6 @@ public Set<AutoConfigureListener> getListeners() {
     return Collections.unmodifiableSet(listeners);
   }
 
-  // Visible for testing
-  interface SpiFinder {
-    <T> Iterable<T> load(Class<T> spiClass, ClassLoader classLoader);
-  }
-
   private static class ServiceLoaderComponentLoader implements ComponentLoader {
     private final ClassLoader classLoader;