Merge pull request #20648 from vespa-engine/mortent/feedclientbuilder-reference

Provide feed client reference

5 files changed, 51 insertions, 31 deletions

diff --git a/cloud-tenant-cd/src/main/java/ai/vespa/hosted/cd/cloud/impl/VespaTestRuntime.java b/cloud-tenant-cd/src/main/java/ai/vespa/hosted/cd/cloud/impl/VespaTestRuntime.java
index 201ddcb3908..1e6065c63b9 100644
--- a/cloud-tenant-cd/src/main/java/ai/vespa/hosted/cd/cloud/impl/VespaTestRuntime.java
+++ b/cloud-tenant-cd/src/main/java/ai/vespa/hosted/cd/cloud/impl/VespaTestRuntime.java
@@ -43,7 +43,7 @@ public class VespaTestRuntime implements TestRuntime {
         DefaultEndpointAuthenticator authenticator = new DefaultEndpointAuthenticator(config.system());
         this.deploymentToTest = new HttpDeployment(config.deployments().get(config.zone()), authenticator);
         FeedClientBuilder.setEndpointAuthenticator(authenticator);
-        System.setProperty(ai.vespa.feed.client.FeedClientBuilder.PREFERRED_IMPLEMENTATION_PROPERTY, FeedClientBuilder.class.getName());
+        ai.vespa.feed.client.FeedClientBuilder.setFeedClientBuilderSupplier(FeedClientBuilder::new);
     }
 
     @Override
diff --git a/cloud-tenant-cd/src/main/resources/META-INF/services/ai.vespa.hosted.cd.TestRuntime b/cloud-tenant-cd/src/main/resources/META-INF/services/ai.vespa.hosted.cd.TestRuntime
deleted file mode 100644
index 9d318b87fca..00000000000
--- a/cloud-tenant-cd/src/main/resources/META-INF/services/ai.vespa.hosted.cd.TestRuntime
+++ /dev/null
@@ -1,2 +0,0 @@
-# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-ai.vespa.hosted.cd.cloud.impl.VespaTestRuntime
\ No newline at end of file
diff --git a/vespa-feed-client-api/abi-spec.json b/vespa-feed-client-api/abi-spec.json
index 8af7798984f..16e532a2c9a 100644
--- a/vespa-feed-client-api/abi-spec.json
+++ b/vespa-feed-client-api/abi-spec.json
@@ -123,6 +123,7 @@
     "methods": [
       "public static ai.vespa.feed.client.FeedClientBuilder create(java.net.URI)",
       "public static ai.vespa.feed.client.FeedClientBuilder create(java.util.List)",
+      "public static void setFeedClientBuilderSupplier(java.util.function.Supplier)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setConnectionsPerEndpoint(int)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setMaxStreamPerConnection(int)",
       "public abstract ai.vespa.feed.client.FeedClientBuilder setSslContext(javax.net.ssl.SSLContext)",
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..95c9b2c95fe 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,13 @@ 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.getFeedClientBuilderSupplier().get().setEndpointUris(endpoints);
     }
 
+    /** Override FeedClientBuilder. This will be preferred in {@link #create} */
+    static void setFeedClientBuilderSupplier(Supplier<FeedClientBuilder> supplier) {
+        Helper.setFeedClientBuilderSupplier(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..6971b2ea8f5 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 {
 
+    private static final AtomicReference<Supplier<FeedClientBuilder>> feedClientBuilderSupplier = new AtomicReference<>(Helper::getFeedClientBuilder);
+
+    static final void setFeedClientBuilderSupplier(Supplier<FeedClientBuilder> supplier) {
+        feedClientBuilderSupplier.set(supplier);
+    }
+
+    static Supplier<FeedClientBuilder> getFeedClientBuilderSupplier() {
+        return feedClientBuilderSupplier.get();
+    }
+
+    static FeedClientBuilder getFeedClientBuilder() {
+        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