Provide feed client reference

2 files changed, 48 insertions, 28 deletions

diff --git a/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java b/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java
index 05bc608df27..94d5a0ba2d1 100644
--- a/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java
+++ b/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java
@@ -3,18 +3,13 @@ package ai.vespa.feed.client;
 
 import javax.net.ssl.HostnameVerifier;
 import javax.net.ssl.SSLContext;
-import java.lang.reflect.Constructor;
-import java.lang.reflect.InvocationTargetException;
 import java.net.URI;
 import java.nio.file.Path;
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.List;
-import java.util.ServiceLoader;
 import java.util.function.Supplier;
 
 /**
@@ -32,31 +27,12 @@ public interface FeedClientBuilder {
 
     /** Creates a builder for multiple container endpoints **/
     static FeedClientBuilder create(List<URI> endpoints) {
-        String defaultImplementation = "ai.vespa.feed.client.impl.FeedClientBuilderImpl";
-        String preferredImplementation = System.getProperty(PREFERRED_IMPLEMENTATION_PROPERTY, defaultImplementation);
-        Iterator<FeedClientBuilder> iterator = ServiceLoader.load(FeedClientBuilder.class).iterator();
-        if (iterator.hasNext()) {
-            List<FeedClientBuilder> builders = new ArrayList<>();
-            iterator.forEachRemaining(builders::add);
-            return builders.stream()
-                    .filter(builder -> preferredImplementation.equals(builder.getClass().getName()))
-                    .findFirst()
-                    .orElse(builders.get(0));
-        } else {
-            try {
-                Class<?> aClass = Class.forName(preferredImplementation);
-                for (Constructor<?> constructor : aClass.getConstructors()) {
-                    if (constructor.getParameterTypes().length==0) {
-                        return ((FeedClientBuilder)constructor.newInstance()).setEndpointUris(endpoints);
-                    }
-                }
-                throw new RuntimeException("Could not find Feed client builder implementation");
-            } catch (ClassNotFoundException | InvocationTargetException | InstantiationException | IllegalAccessException e) {
-                throw new RuntimeException(e);
-            }
-        }
+        return Helper.getFeedClientBuilder().setEndpointUris(endpoints);
     }
 
+    static void setFeedClientBuilderSupplier(Supplier<FeedClientBuilder> supplier) {
+        Helper.setFeedClientBuilderReference(supplier);
+    }
     /**
      * Sets the number of connections this client will use per endpoint.
      *
diff --git a/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/Helper.java b/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/Helper.java
index 59c12077bef..edf9457aaa1 100644
--- a/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/Helper.java
+++ b/vespa-feed-client-api/src/main/java/ai/vespa/feed/client/Helper.java
@@ -1,18 +1,62 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.feed.client;
 
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
+import java.util.ServiceLoader;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
+import static ai.vespa.feed.client.FeedClientBuilder.PREFERRED_IMPLEMENTATION_PROPERTY;
+
 /**
  * @author bjorncs
  */
 class Helper {
 
+    static final AtomicReference<Supplier<FeedClientBuilder>> feedClientBuilderSupplier = new AtomicReference<>();
+
+    static final void setFeedClientBuilderReference(Supplier<FeedClientBuilder> supplier) {
+        feedClientBuilderSupplier.set(supplier);
+    }
+
+    static FeedClientBuilder getFeedClientBuilder() {
+        if (Helper.feedClientBuilderSupplier.get()!=null) {
+            return Helper.feedClientBuilderSupplier.get().get();
+        } else {
+            String defaultImplementation = "ai.vespa.feed.client.impl.FeedClientBuilderImpl";
+            String preferredImplementation = System.getProperty(PREFERRED_IMPLEMENTATION_PROPERTY, defaultImplementation);
+            Iterator<FeedClientBuilder> iterator = ServiceLoader.load(FeedClientBuilder.class).iterator();
+            if (iterator.hasNext()) {
+                List<FeedClientBuilder> builders = new ArrayList<>();
+                iterator.forEachRemaining(builders::add);
+                return builders.stream()
+                        .filter(builder -> preferredImplementation.equals(builder.getClass().getName()))
+                        .findFirst()
+                        .orElse(builders.get(0));
+            } else {
+                try {
+                    Class<?> aClass = Class.forName(preferredImplementation);
+                    for (Constructor<?> constructor : aClass.getConstructors()) {
+                        if (constructor.getParameterTypes().length == 0) {
+                            return ((FeedClientBuilder) constructor.newInstance());
+                        }
+                    }
+                    throw new RuntimeException("Could not find Feed client builder implementation");
+                } catch (ClassNotFoundException | InvocationTargetException | InstantiationException | IllegalAccessException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        }
+    }
+
     private Helper() {}
 
     @SafeVarargs