Allow SiaIdentityProvider to notify listeners after reload

1 files changed, 17 insertions, 2 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
index 2b4b267492a..437da05dfdd 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
@@ -14,10 +14,13 @@ import java.io.File;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.time.Duration;
+import java.util.Set;
+import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 import java.util.logging.Logger;
 
 /**
@@ -30,12 +33,13 @@ public class SiaIdentityProvider extends AbstractComponent implements AthenzIden
 
     private static final Duration REFRESH_INTERVAL = Duration.ofHours(1);
 
-    private final AtomicReference<SSLContext> sslContext = new AtomicReference<SSLContext>();
+    private final AtomicReference<SSLContext> sslContext = new AtomicReference<>();
     private final AthenzService service;
     private final File privateKeyFile;
     private final File certificateFile;
     private final File trustStoreFile;
     private final ScheduledExecutorService scheduler;
+    private final Set<Consumer<SSLContext>> listeners = new ConcurrentSkipListSet<>();
 
     @Inject
     public SiaIdentityProvider(SiaProviderConfig config) {
@@ -93,6 +97,14 @@ public class SiaIdentityProvider extends AbstractComponent implements AthenzIden
         return sslContext.get();
     }
 
+    public void addReloadListener(Consumer<SSLContext> listener) {
+        listeners.add(listener);
+    }
+
+    public void removeReloadListener(Consumer<SSLContext> listener) {
+        listeners.remove(listener);
+    }
+
     private SSLContext createIdentitySslContext() {
         return new SslContextBuilder()
                 .withTrustStore(trustStoreFile, KeyStoreType.JKS)
@@ -103,7 +115,9 @@ public class SiaIdentityProvider extends AbstractComponent implements AthenzIden
     private void reloadSslContext() {
         log.log(LogLevel.DEBUG, "Updating SSLContext for identity " + service.getFullName());
         try {
-            this.sslContext.set(createIdentitySslContext());
+            SSLContext sslContext = createIdentitySslContext();
+            this.sslContext.set(sslContext);
+            listeners.forEach(listener -> listener.accept(sslContext));
         } catch (Exception e) {
             log.log(LogLevel.SEVERE, "Failed to update SSLContext: " + e.getMessage(), e);
         }
@@ -122,6 +136,7 @@ public class SiaIdentityProvider extends AbstractComponent implements AthenzIden
         try {
             scheduler.shutdownNow();
             scheduler.awaitTermination(90, TimeUnit.SECONDS);
+            listeners.clear();
         } catch (InterruptedException e) {
             throw new RuntimeException(e);
         }