Use reference counting to avoid relying on GC to drop threads.

2 files changed, 36 insertions, 8 deletions

diff --git a/security-utils/src/main/java/com/yahoo/security/tls/ConfigFileBasedTlsContext.java b/security-utils/src/main/java/com/yahoo/security/tls/ConfigFileBasedTlsContext.java
index 26dfbf9fd9f..9527d50f339 100644
--- a/security-utils/src/main/java/com/yahoo/security/tls/ConfigFileBasedTlsContext.java
+++ b/security-utils/src/main/java/com/yahoo/security/tls/ConfigFileBasedTlsContext.java
@@ -31,13 +31,17 @@ public class ConfigFileBasedTlsContext implements TlsContext {
 
     private static TlsManager getOrCreateTrustManager(Path tlsOptionsConfigFile) {
         synchronized (trustManagers) {
-            WeakReference<TlsManager> tlsManager = trustManagers.get(tlsOptionsConfigFile);
-            if (tlsManager == null || tlsManager.get() == null) {
-                TlsManager manager = new TlsManager(tlsOptionsConfigFile);
-                trustManagers.put(tlsOptionsConfigFile, new WeakReference<>(manager));
-                return manager;
+            WeakReference<TlsManager> tlsRef = trustManagers.get(tlsOptionsConfigFile);
+            TlsManager tlsManager = null;
+            if (tlsRef != null) {
+                tlsManager = tlsRef.get();
             }
-            return tlsManager.get();
+            if (tlsManager == null) {
+                tlsManager = new TlsManager(tlsOptionsConfigFile);
+                trustManagers.put(tlsOptionsConfigFile, new WeakReference<>(tlsManager));
+            }
+            tlsManager.addRef();
+            return tlsManager;
         }
     }
 
@@ -59,6 +63,15 @@ public class ConfigFileBasedTlsContext implements TlsContext {
     @Override public SSLParameters parameters() { return tlsContext.parameters(); }
     @Override public SSLEngine createSslEngine() { return tlsContext.createSslEngine(); }
     @Override public SSLEngine createSslEngine(String peerHost, int peerPort) { return tlsContext.createSslEngine(peerHost, peerPort); }
+    @Override public void close() {
+        synchronized (trustManagers) {
+            int references = tlsManager.subRef();
+            if (references == 0) {
+                tlsManager.close();
+                trustManagers.remove(tlsManager.getTlsConfigFile());
+            }
+        }
+    }
 
     private static DefaultTlsContext createDefaultTlsContext(TransportSecurityOptions options,
                                                              AuthorizationMode mode,
diff --git a/security-utils/src/main/java/com/yahoo/security/tls/TlsManager.java b/security-utils/src/main/java/com/yahoo/security/tls/TlsManager.java
index bade1993982..c2286c3e8ac 100644
--- a/security-utils/src/main/java/com/yahoo/security/tls/TlsManager.java
+++ b/security-utils/src/main/java/com/yahoo/security/tls/TlsManager.java
@@ -16,6 +16,7 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -23,10 +24,12 @@ class TlsManager {
     private static final Logger log = Logger.getLogger(TlsManager.class.getName());
     private static final Duration UPDATE_PERIOD = Duration.ofHours(1);
 
+    private final Path tlsConfigFile;
     private final MutableX509TrustManager trustManager;
     private final MutableX509KeyManager keyManager;
     private final ScheduledExecutorService scheduler;
     private TransportSecurityOptions options;
+    private final AtomicInteger references = new AtomicInteger(0);
 
     private static void reloadTrustManager(TransportSecurityOptions options, MutableX509TrustManager trustManager) {
         if (options.getCaCertificatesFile().isPresent()) {
@@ -86,13 +89,13 @@ class TlsManager {
             try {
                 TlsManager tlsManager = this.tlsManager.get();
                 if (tlsManager == null) {
+                    // If reference count is done correctly this should not be necessary.
                     scheduler.shutdown();
                     return;
                 }
                 TransportSecurityOptions options = TransportSecurityOptions.fromJsonFile(tlsOptionsConfigFile);
                 reloadTrustManager(options, tlsManager.getTrustManager());
-                MutableX509KeyManager keyManager = tlsManager.getKeyManager();
-                reloadKeyManager(options, keyManager);
+                reloadKeyManager(options, tlsManager.getKeyManager());
             } catch (Throwable t) {
                 log.log(Level.SEVERE, String.format("Failed to reload crypto material (path='%s'): %s", tlsOptionsConfigFile, t.getMessage()), t);
             }
@@ -114,6 +117,7 @@ class TlsManager {
     }
 
     TlsManager(Path tlsOptionsConfigFile) {
+        tlsConfigFile = tlsOptionsConfigFile;
         trustManager = new MutableX509TrustManager();
         keyManager = new MutableX509KeyManager();
         options = TransportSecurityOptions.fromJsonFile(tlsOptionsConfigFile);
@@ -133,7 +137,18 @@ class TlsManager {
         return keyManager;
     }
 
+    Path getTlsConfigFile() {
+        return tlsConfigFile;
+    }
+
     TransportSecurityOptions getOptions() {
         return options;
     }
+
+    void close() {
+        scheduler.shutdown();
+    }
+
+    int addRef() { return references.incrementAndGet(); }
+    int subRef() { return references.decrementAndGet(); }
 }