Add test and correct counting errors.

2 files changed, 47 insertions, 17 deletions

diff --git a/config/src/main/java/com/yahoo/config/subscription/impl/JRTManagedConnectionPools.java b/config/src/main/java/com/yahoo/config/subscription/impl/JRTManagedConnectionPools.java
index 6889bbdaf0c..0a606416827 100644
--- a/config/src/main/java/com/yahoo/config/subscription/impl/JRTManagedConnectionPools.java
+++ b/config/src/main/java/com/yahoo/config/subscription/impl/JRTManagedConnectionPools.java
@@ -26,31 +26,38 @@ public class JRTManagedConnectionPools {
         final ScheduledThreadPoolExecutor scheduler = new ScheduledThreadPoolExecutor(1, new JRTSourceThreadFactory());
         CountedPool(JRTConnectionPool requester) {
             this.pool = requester;
-            count = 1;
+            count = 0;
         }
     }
     private Map<ConfigSourceSet, CountedPool> pools = new HashMap<>();
 
-    public synchronized JRTConfigRequester acquire(ConfigSourceSet sourceSet, TimingValues timingValues) {
-        CountedPool countedPool = pools.get(sourceSet);
-        if (countedPool == null) {
-            countedPool = new CountedPool(new JRTConnectionPool(sourceSet));
-            pools.put(sourceSet, countedPool);
+    public JRTConfigRequester acquire(ConfigSourceSet sourceSet, TimingValues timingValues) {
+        CountedPool countedPool;
+        synchronized (pools) {
+            countedPool = pools.get(sourceSet);
+            if (countedPool == null) {
+                countedPool = new CountedPool(new JRTConnectionPool(sourceSet));
+                pools.put(sourceSet, countedPool);
+            }
+            countedPool.count++;
         }
-        countedPool.count++;
         return new JRTConfigRequester(sourceSet, countedPool.scheduler, countedPool.pool, timingValues);
     }
     public synchronized void release(ConfigSourceSet sourceSet) {
-        CountedPool countedPool = pools.get(sourceSet);
-        countedPool.count--;
-        if (countedPool.count == 0) {
-            countedPool.pool.close();
-            countedPool.scheduler.shutdown();
-           try {
-               countedPool.scheduler.awaitTermination(30, TimeUnit.SECONDS);
-           } catch (InterruptedException e) {
-               throw new RuntimeException("Failed shutting down scheduler:", e);
-           }
+        CountedPool countedPool;
+        synchronized (pools) {
+            countedPool = pools.get(sourceSet);
+            countedPool.count--;
+            if (countedPool.count > 0) return;
+            pools.remove(sourceSet);
+        }
+
+        countedPool.pool.close();
+        countedPool.scheduler.shutdown();
+        try {
+            countedPool.scheduler.awaitTermination(30, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            throw new RuntimeException("Failed shutting down scheduler:", e);
         }
     }
 }
diff --git a/config/src/test/java/com/yahoo/config/subscription/impl/JRTConfigRequesterTest.java b/config/src/test/java/com/yahoo/config/subscription/impl/JRTConfigRequesterTest.java
index 757dd99f43b..4211345dff7 100644
--- a/config/src/test/java/com/yahoo/config/subscription/impl/JRTConfigRequesterTest.java
+++ b/config/src/test/java/com/yahoo/config/subscription/impl/JRTConfigRequesterTest.java
@@ -1,10 +1,12 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.config.subscription.impl;
 
+import com.yahoo.config.subscription.ConfigSourceSet;
 import com.yahoo.foo.SimpletypesConfig;
 import com.yahoo.config.subscription.ConfigSubscriber;
 import com.yahoo.jrt.Request;
 import com.yahoo.vespa.config.ConfigKey;
+import com.yahoo.vespa.config.ConnectionPool;
 import com.yahoo.vespa.config.ErrorCode;
 import com.yahoo.vespa.config.ErrorType;
 import com.yahoo.vespa.config.TimingValues;
@@ -17,6 +19,8 @@ import java.util.Random;
 
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
@@ -349,4 +353,23 @@ public class JRTConfigRequesterTest {
         }
     }
 
+    @Test
+    public void testManagedPool() {
+        ConfigSourceSet sourceSet = ConfigSourceSet.createDefault();
+        TimingValues timingValues = new TimingValues();
+        JRTConfigRequester requester1 = JRTConfigRequester.create(sourceSet, timingValues);
+        JRTConfigRequester requester2 = JRTConfigRequester.create(sourceSet, timingValues);
+        assertNotSame(requester1, requester2);
+        assertSame(requester1.getConnectionPool(), requester2.getConnectionPool());
+        ConnectionPool firstPool = requester1.getConnectionPool();
+        requester1.close();
+        requester2.close();
+        requester1 = JRTConfigRequester.create(sourceSet, timingValues);
+        assertNotSame(firstPool, requester1.getConnectionPool());
+        requester2 = JRTConfigRequester.create(new ConfigSourceSet("test-managed-pool-2"), timingValues);
+        assertNotSame(requester1.getConnectionPool(), requester2.getConnectionPool());
+        requester1.close();
+        requester2.close();
+    }
+
 }