Remove scheduler wrapper

2 files changed, 14 insertions, 48 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImpl.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImpl.java
index 12f2ce5f074..78ad95f84f3 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImpl.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImpl.java
@@ -42,7 +42,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     private volatile AthenzCredentials credentials;
     private final Metric metric;
     private final AthenzCredentialsService athenzCredentialsService;
-    private final Scheduler scheduler;
+    private final ScheduledExecutorService scheduler;
     private final Clock clock;
     private final String domain;
     private final String service;
@@ -55,7 +55,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
                                           new IdentityDocumentService(config.loadBalancerAddress()),
                                           new AthenzService(),
                                           Clock.systemUTC()),
-             new ThreadPoolScheduler(),
+             new ScheduledThreadPoolExecutor(1),
              Clock.systemUTC());
     }
 
@@ -63,7 +63,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     AthenzIdentityProviderImpl(IdentityConfig config,
                                Metric metric,
                                AthenzCredentialsService athenzCredentialsService,
-                               Scheduler scheduler,
+                               ScheduledExecutorService scheduler,
                                Clock clock) {
         this.metric = metric;
         this.athenzCredentialsService = athenzCredentialsService;
@@ -77,8 +77,8 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     private void registerInstance() {
         try {
             credentials = athenzCredentialsService.registerInstance();
-            scheduler.scheduleAtFixedRate(new CertificateExpiryMetricUpdater(), Duration.ofMinutes(0), Duration.ofMinutes(5));
-            scheduler.scheduleAtFixedRate(new UpdateCredentialsTask(), UPDATE_PERIOD, UPDATE_PERIOD);
+            scheduler.scheduleAtFixedRate(this::refreshCertificate, 0, 5, TimeUnit.MINUTES);
+            scheduler.scheduleAtFixedRate(this::reportMetrics, UPDATE_PERIOD.toMinutes(), UPDATE_PERIOD.toMinutes(), TimeUnit.MINUTES);
         } catch (Throwable t) {
             throw new AthenzIdentityProviderException("Could not retrieve Athenz credentials", t);
         }
@@ -106,7 +106,12 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     @Override
     public void deconstruct() {
-        scheduler.shutdown(AWAIT_TERMINTATION_TIMEOUT);
+        try {
+            scheduler.shutdownNow();
+            scheduler.awaitTermination(AWAIT_TERMINTATION_TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
     }
 
     private boolean isExpired(AthenzCredentials credentials) {
@@ -137,43 +142,5 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
             log.log(LogLevel.WARNING, "Failed to update metrics: " + t.getMessage(), t);
         }
     }
-
-    private class UpdateCredentialsTask implements Runnable {
-        @Override
-        public void run() {
-            refreshCertificate();
-        }
-    }
-
-    private class CertificateExpiryMetricUpdater implements Runnable {
-        @Override
-        public void run() {
-            reportMetrics();
-        }
-    }
-
-    private static class ThreadPoolScheduler implements Scheduler {
-        private final ScheduledExecutorService executor = Executors.newScheduledThreadPool(1);
-
-        @Override
-        public void scheduleAtFixedRate(Runnable runnable, Duration initialDelay, Duration period) {
-            executor.scheduleAtFixedRate(runnable, initialDelay.getSeconds(), period.getSeconds(), TimeUnit.SECONDS);
-        }
-
-        @Override
-        public void shutdown(Duration timeout) {
-            try {
-                executor.shutdownNow();
-                executor.awaitTermination(timeout.getSeconds(), TimeUnit.SECONDS);
-            } catch (InterruptedException e) {
-                throw new RuntimeException(e);
-            }
-        }
-    }
-
-    public interface Scheduler {
-        void scheduleAtFixedRate(Runnable runnable, Duration initialDelay, Duration period);
-        void shutdown(Duration timeout);
-    }
 }
 
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImplTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImplTest.java
index 29b122b09e5..3bb074cc1dc 100644
--- a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImplTest.java
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/AthenzIdentityProviderImplTest.java
@@ -2,11 +2,9 @@
 package com.yahoo.vespa.athenz.identityprovider;
 
 import com.yahoo.container.core.identity.IdentityConfig;
-import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
 import com.yahoo.container.jdisc.athenz.AthenzIdentityProviderException;
 import com.yahoo.jdisc.Metric;
 import com.yahoo.test.ManualClock;
-import com.yahoo.vespa.athenz.identityprovider.AthenzIdentityProviderImpl.Scheduler;
 import org.junit.Test;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
@@ -15,6 +13,7 @@ import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Date;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Supplier;
 
 import static org.mockito.Matchers.any;
@@ -42,7 +41,7 @@ public class AthenzIdentityProviderImplTest {
         when(credentialService.registerInstance())
                 .thenThrow(new RuntimeException("athenz unavailable"));
 
-        new AthenzIdentityProviderImpl(IDENTITY_CONFIG, mock(Metric.class), credentialService, mock(Scheduler.class), new ManualClock(Instant.EPOCH));
+        new AthenzIdentityProviderImpl(IDENTITY_CONFIG, mock(Metric.class), credentialService, mock(ScheduledExecutorService.class), new ManualClock(Instant.EPOCH));
     }
 
     @Test
@@ -70,7 +69,7 @@ public class AthenzIdentityProviderImplTest {
                 new AthenzCredentialsService(IDENTITY_CONFIG, identityDocumentService, athenzService, clock);
 
         AthenzIdentityProviderImpl identityProvider =
-                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, metric, credentialService, mock(Scheduler.class), clock);
+                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, metric, credentialService, mock(ScheduledExecutorService.class), clock);
 
         identityProvider.reportMetrics();
         verify(metric).set(eq(AthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.getSeconds()), any());