fix: Fix lifecycle of shared clients obtained from EtcdClusterConfig (#58)

Motivation

A bug was reported where an invalid client was returned from EtcdClusterConfig#getClient() after a client had previously been obtained and closed. This turned out to be due to some incorrect logic in how "released" ref-counted clients are recycled.

Modifications

- Fix the non-blocking logic in EtcdClientConfig#getClient() for obtaining a shared ref-counted client - in particular ensure that retain() is called on preexisting client instances from the cache.
- Add unit test

Result

Client caching/sharing in EtcdClusterConfig works as intended

Author: njhill

src/main/java/com/ibm/etcd/client/config/EtcdClusterConfig.java

@@ -25,11 +25,7 @@
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
 import java.security.cert.CertificateException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Objects;
-import java.util.Properties;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.ExecutionException;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -276,6 +272,7 @@ public static EtcdClusterConfig fromJsonFileOrSimple(String fileOrSimpleString)
     }
 
     private static final Cache<CacheKey, EtcdClient> clientCache = CacheBuilder.newBuilder().weakValues()
+            // Not a problem to call close() more than once on the underlying clients
             .<CacheKey, EtcdClient>removalListener(rn -> EtcdClient.Internal.cleanup(rn.getValue())).build();
 
     public static EtcdClient getClient(EtcdClusterConfig config) throws IOException, CertificateException {
@@ -285,15 +282,24 @@ public static EtcdClient getClient(EtcdClusterConfig config) throws IOException,
             while (true) {
                 EtcdClient client = clientCache.getIfPresent(key);
                 if (client == null) {
-                    EtcdClient[] maybeNew = new EtcdClient[1];
-                    client  = clientCache.get(key, () -> maybeNew[0] = config.newClient());
-                    // if client != maybeNew[0] then we got a pre-existing client
-                    if (client != maybeNew[0] && !EtcdClient.Internal.retain(client)) {
-                        clientCache.invalidate(key);
-                        continue;
+                    final Deque<EtcdClient> maybeNew = new ArrayDeque<>(1);
+                    try {
+                        client = clientCache.get(key, () -> {
+                            EtcdClient newClient = config.newClient();
+                            maybeNew.add(newClient);
+                            return newClient;
+                        });
+                        if (client == maybeNew.peekLast()) { // else we got a preexisting client
+                            return maybeNew.pollLast();
+                        }
+                    } finally {
+                        maybeNew.stream().forEach(EtcdClient::close);
                     }
                 }
-                return client;
+                if (EtcdClient.Internal.retain(client)) {
+                    return client;
+                }
+                clientCache.invalidate(key); // Client is closed, discard and try again
             }
         } catch (ExecutionException ee) {
             Throwables.throwIfInstanceOf(ee.getCause(), IOException.class);

src/test/java/com/ibm/etcd/client/JsonConfigTest.java

@@ -16,8 +16,7 @@
 package com.ibm.etcd.client;
 
 import static com.ibm.etcd.client.KeyUtils.bs;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.*;
 
 import java.nio.charset.StandardCharsets;
 
@@ -85,6 +84,33 @@ public void testSslCaCertConfig_noClientCert() throws Exception { // should fail
         runBasicTests(EtcdClusterConfig.fromJson(json));
     }
 
+    @Test
+    public void testLifecycle() throws Exception {
+        // Test shared client reuse/recreation
+        ByteSource json = makeJson("http://localhost:2379", null, null, null);
+        EtcdClusterConfig config = EtcdClusterConfig.fromJson(json);
+        runBasicTests(config);
+        runBasicTests(config);
+        EtcdClient anotherClient = config.getClient();
+        runBasicTests(config);
+        assert !anotherClient.isClosed();
+        anotherClient.close();
+        assert anotherClient.isClosed();
+
+        EtcdClient c1 = config.getClient(), c2 = config.getClient();
+        assertSame(c1, c2);
+        c2.close();
+        EtcdClient c3 = config.getClient();
+        assertSame(c1, c3);
+        c1.close();
+        assert !c1.isClosed();
+        c3.close();
+        assert c1.isClosed();
+
+        EtcdClient c4 = config.getClient();
+        assertNotSame(c1, c4);
+    }
+
     void runBasicTests(EtcdClusterConfig config) throws Exception {
         try (EtcdClient client = config.getClient()) {
             runBasicTests(client);