Revert "Revert "Report metrics on athenz client errors""

6 files changed, 68 insertions, 32 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/aws/AwsCredentials.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/aws/AwsCredentials.java
index b027e7272ea..30ff63fb108 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/aws/AwsCredentials.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/aws/AwsCredentials.java
@@ -35,11 +35,11 @@ public class AwsCredentials {
     }
 
     public AwsCredentials(URI ztsUrl, ServiceIdentityProvider identityProvider, AthenzDomain athenzDomain, AwsRole awsRole) {
-        this(new DefaultZtsClient(ztsUrl, identityProvider), athenzDomain, awsRole);
+        this(new DefaultZtsClient.Builder(ztsUrl).withIdentityProvider(identityProvider).build(), athenzDomain, awsRole);
     }
 
     public AwsCredentials(URI ztsUrl, SSLContext sslContext, AthenzDomain athenzDomain, AwsRole awsRole) {
-        this(new DefaultZtsClient(ztsUrl, sslContext), athenzDomain, awsRole);
+        this(new DefaultZtsClient.Builder(ztsUrl).withSslContext(sslContext).build(), athenzDomain, awsRole);
     }
 
     /**
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/common/ClientBase.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/common/ClientBase.java
index c1ce45c35da..37ef513c786 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/common/ClientBase.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/common/ClientBase.java
@@ -39,12 +39,15 @@ public abstract class ClientBase implements AutoCloseable {
 
     private final CloseableHttpClient client;
     private final ClientExceptionFactory exceptionFactory;
+    private final ErrorHandler errorHandler;
 
     protected ClientBase(String userAgent,
                          Supplier<SSLContext> sslContextSupplier,
                          ClientExceptionFactory exceptionFactory,
-                         HostnameVerifier hostnameVerifier) {
+                         HostnameVerifier hostnameVerifier,
+                         ErrorHandler errorHandler) {
         this.exceptionFactory = exceptionFactory;
+        this.errorHandler = errorHandler;
         this.client = createHttpClient(userAgent, sslContextSupplier, hostnameVerifier);
     }
 
@@ -52,10 +55,17 @@ public abstract class ClientBase implements AutoCloseable {
         try {
             return client.execute(request, responseHandler);
         } catch (IOException e) {
+            try {
+                reportError(request, e);
+            } catch (Exception _ignored) {}
             throw new UncheckedIOException(e);
         }
     }
 
+    private void reportError(HttpUriRequest request, Exception e) {
+        errorHandler.reportError(request, e);
+    }
+
     protected StringEntity toJsonStringEntity(Object entity) {
         try {
             return new StringEntity(objectMapper.writeValueAsString(entity), ContentType.APPLICATION_JSON);
@@ -114,4 +124,11 @@ public abstract class ClientBase implements AutoCloseable {
     protected interface ClientExceptionFactory {
         RuntimeException createException(int errorCode, String description);
     }
+
+    public interface ErrorHandler {
+        static ErrorHandler empty() {
+            return (r,e)->{};
+        }
+        void reportError(HttpUriRequest request, Exception error);
+    }
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/DefaultZmsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/DefaultZmsClient.java
index 33cb6d7d5d4..3742996c274 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/DefaultZmsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/DefaultZmsClient.java
@@ -25,7 +25,6 @@ import javax.net.ssl.SSLContext;
 import java.net.URI;
 import java.util.Collections;
 import java.util.List;
-import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.Set;
 import java.util.function.Supplier;
@@ -40,16 +39,16 @@ public class DefaultZmsClient extends ClientBase implements ZmsClient {
     private final URI zmsUrl;
     private final AthenzIdentity identity;
 
-    public DefaultZmsClient(URI zmsUrl, AthenzIdentity identity, SSLContext sslContext) {
-        this(zmsUrl, identity, () -> sslContext);
+    public DefaultZmsClient(URI zmsUrl, AthenzIdentity identity, SSLContext sslContext, ErrorHandler errorHandler) {
+        this(zmsUrl, identity, () -> sslContext, errorHandler);
     }
 
-    public DefaultZmsClient(URI zmsUrl, ServiceIdentityProvider identityProvider) {
-        this(zmsUrl, identityProvider.identity(), identityProvider::getIdentitySslContext);
+    public DefaultZmsClient(URI zmsUrl, ServiceIdentityProvider identityProvider, ErrorHandler errorHandler) {
+        this(zmsUrl, identityProvider.identity(), identityProvider::getIdentitySslContext, errorHandler);
     }
 
-    private DefaultZmsClient(URI zmsUrl, AthenzIdentity identity, Supplier<SSLContext> sslContextSupplier) {
-        super("vespa-zms-client", sslContextSupplier, ZmsClientException::new, null);
+    private DefaultZmsClient(URI zmsUrl, AthenzIdentity identity, Supplier<SSLContext> sslContextSupplier, ErrorHandler errorHandler) {
+        super("vespa-zms-client", sslContextSupplier, ZmsClientException::new, null, errorHandler);
         this.zmsUrl = addTrailingSlash(zmsUrl);
         this.identity = identity;
     }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
index c05213c8008..28119dc1f5a 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
@@ -37,6 +37,7 @@ import java.security.KeyPair;
 import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -52,25 +53,8 @@ import static java.util.stream.Collectors.toList;
 public class DefaultZtsClient extends ClientBase implements ZtsClient {
 
     private final URI ztsUrl;
-
-    public DefaultZtsClient(URI ztsUrl, SSLContext sslContext) {
-        this(ztsUrl, sslContext, null);
-    }
-
-    public DefaultZtsClient(URI ztsUrl, SSLContext sslContext, HostnameVerifier hostnameVerifier) {
-        this(ztsUrl, () -> sslContext, hostnameVerifier);
-    }
-
-    public DefaultZtsClient(URI ztsUrl, ServiceIdentityProvider identityProvider) {
-        this(ztsUrl, identityProvider::getIdentitySslContext, null);
-    }
-
-    public DefaultZtsClient(URI ztsUrl, ServiceIdentityProvider identityProvider, HostnameVerifier hostnameVerifier) {
-        this(ztsUrl, identityProvider::getIdentitySslContext, hostnameVerifier);
-    }
-
-    private DefaultZtsClient(URI ztsUrl, Supplier<SSLContext> sslContextSupplier, HostnameVerifier hostnameVerifier) {
-        super("vespa-zts-client", sslContextSupplier, ZtsClientException::new, hostnameVerifier);
+    protected DefaultZtsClient(URI ztsUrl, Supplier<SSLContext> sslContextSupplier, HostnameVerifier hostnameVerifier, ErrorHandler errorHandler) {
+        super("vespa-zts-client", sslContextSupplier, ZtsClientException::new, hostnameVerifier, errorHandler);
         this.ztsUrl = addTrailingSlash(ztsUrl);
     }
 
@@ -239,5 +223,41 @@ public class DefaultZtsClient extends ClientBase implements ZtsClient {
         else
             return URI.create(ztsUrl.toString() + '/');
     }
+    public static class Builder {
+        private URI ztsUrl;
+        private ClientBase.ErrorHandler errorHandler = ErrorHandler.empty();
+        private HostnameVerifier hostnameVerifier = null;
+        private Supplier<SSLContext> sslContextSupplier = null;
+
+        public Builder(URI ztsUrl) {
+            this.ztsUrl = ztsUrl;
+        }
 
+        public Builder withErrorHandler(ClientBase.ErrorHandler errorHandler) {
+            this.errorHandler = errorHandler;
+            return this;
+        }
+
+        public Builder withHostnameVerifier(HostnameVerifier hostnameVerifier) {
+            this.hostnameVerifier = hostnameVerifier;
+            return this;
+        }
+
+        public Builder withSslContext(SSLContext sslContext) {
+            this.sslContextSupplier = () -> sslContext;
+            return this;
+        }
+
+        public Builder withIdentityProvider(ServiceIdentityProvider identityProvider) {
+            this.sslContextSupplier = identityProvider::getIdentitySslContext;
+            return this;
+        }
+
+        public DefaultZtsClient build() {
+            if (Objects.isNull(sslContextSupplier)) {
+                throw new IllegalArgumentException("No ssl context or identity provider available to set up zts client");
+            }
+            return new DefaultZtsClient(ztsUrl, sslContextSupplier, hostnameVerifier, errorHandler);
+        }
+    }
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
index 8e029906c30..612f9caa691 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
@@ -81,7 +81,7 @@ class AthenzCredentialsService {
                 document.ipAddresses(),
                 keyPair);
 
-        try (ZtsClient ztsClient = new DefaultZtsClient(ztsEndpoint, nodeIdentityProvider)) {
+        try (ZtsClient ztsClient = new DefaultZtsClient.Builder(ztsEndpoint).withIdentityProvider(nodeIdentityProvider).build()) {
             InstanceIdentity instanceIdentity =
                     ztsClient.registerInstance(
                             configserverIdentity,
@@ -102,7 +102,7 @@ class AthenzCredentialsService {
                 document.ipAddresses(),
                 newKeyPair);
 
-        try (ZtsClient ztsClient = new DefaultZtsClient(ztsEndpoint, sslContext)) {
+        try (ZtsClient ztsClient = new DefaultZtsClient.Builder(ztsEndpoint).withSslContext(sslContext).build()) {
             InstanceIdentity instanceIdentity =
                     ztsClient.refreshInstance(
                             configserverIdentity,
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
index 65574d7583e..724a3059f6d 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
@@ -301,7 +301,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     }
 
     private DefaultZtsClient createZtsClient() {
-        return new DefaultZtsClient(ztsEndpoint, getIdentitySslContext());
+        return new DefaultZtsClient.Builder(ztsEndpoint).withSslContext(getIdentitySslContext()).build();
     }
 
     @Override