Fix #468 - properly close clients when they are unused and no longer cached

Author: chadlwilson

build.gradle

@@ -20,7 +20,7 @@ apply from: "https://raw.githubusercontent.com/gocd/gocd-plugin-gradle-task-help
 
 gocdPlugin {
   id = 'cd.go.contrib.elasticagent.kubernetes'
-  pluginVersion = '4.1.0'
+  pluginVersion = '4.1.1'
   goCdVersion = '21.4.0'
   name = 'Kubernetes Elastic Agent Plugin'
   description = 'Kubernetes Based Elastic Agent Plugins for GoCD'

src/main/java/cd/go/contrib/elasticagent/KubernetesAgentInstances.java

@@ -87,13 +87,13 @@ private KubernetesInstance createKubernetesInstance(CreateAgentRequest request,
             return null;
         }
 
-        KubernetesClient client = factory.client(settings);
-        KubernetesInstance instance = kubernetesInstanceFactory.create(request, settings, client, pluginRequest);
-        consoleLogAppender.accept(String.format("Creating pod: %s", instance.name()));
-        register(instance);
-        consoleLogAppender.accept(String.format("Agent pod %s created. Waiting for it to register to the GoCD server.", instance.name()));
-
-        return instance;
+        try (KubernetesClientFactory.CachedClient client = factory.client(settings)) {
+            KubernetesInstance instance = kubernetesInstanceFactory.create(request, settings, client.get(), pluginRequest);
+            consoleLogAppender.accept(String.format("Creating pod: %s", instance.name()));
+            register(instance);
+            consoleLogAppender.accept(String.format("Agent pod %s created. Waiting for it to register to the GoCD server.", instance.name()));
+            return instance;
+        }
     }
 
     private boolean isAgentCreatedForJob(Long jobId) {
@@ -110,8 +110,9 @@ private boolean isAgentCreatedForJob(Long jobId) {
     public void terminate(String agentId, PluginSettings settings) {
         KubernetesInstance instance = instances.get(agentId);
         if (instance != null) {
-            KubernetesClient client = factory.client(settings);
-            instance.terminate(client);
+            try (KubernetesClientFactory.CachedClient client = factory.client(settings)) {
+                instance.terminate(client.get());
+            }
         } else {
             LOG.warn(format("Requested to terminate an instance that does not exist {0}.", agentId));
         }
@@ -152,16 +153,18 @@ public void refreshAll(PluginSettings properties) {
         LOG.debug("[Refresh Instances] Syncing k8s elastic agent pod information for cluster {}.", properties);
         PodList list = null;
         try {
-            KubernetesClient client = factory.client(properties);
-            list = client.pods().list();
+            try (KubernetesClientFactory.CachedClient client = factory.client(properties)) {
+                list = client.get().pods().list();
+            }
         } catch (Exception e) {
             LOG.error("Error occurred while trying to list kubernetes pods:", e);
 
             if (e.getCause() instanceof SocketTimeoutException) {
                 LOG.error("Error caused due to SocketTimeoutException. This generally happens due to stale kubernetes client. Clearing out existing kubernetes client and creating a new one!");
                 factory.clearOutExistingClient();
-                KubernetesClient client = factory.client(properties);
-                list = client.pods().list();
+                try (KubernetesClientFactory.CachedClient client = factory.client(properties)) {
+                    list = client.get().pods().list();
+                }
             }
         }
 
@@ -195,23 +198,23 @@ public void register(KubernetesInstance instance) {
     private KubernetesAgentInstances unregisteredAfterTimeout(PluginSettings settings, Agents knownAgents) throws Exception {
         Duration period = settings.getAutoRegisterPeriod();
         KubernetesAgentInstances unregisteredInstances = new KubernetesAgentInstances();
-        KubernetesClient client = factory.client(settings);
-
-        for (String instanceName : instances.keySet()) {
-            if (knownAgents.containsAgentWithId(instanceName)) {
-                continue;
-            }
+        try (KubernetesClientFactory.CachedClient client = factory.client(settings)) {
+            for (String instanceName : instances.keySet()) {
+                if (knownAgents.containsAgentWithId(instanceName)) {
+                    continue;
+                }
 
-            Pod pod = getPod(client, instanceName);
-            if (pod == null) {
-                LOG.debug(String.format("[server-ping] Pod with name %s is already deleted.", instanceName));
-                continue;
-            }
+                Pod pod = getPod(client.get(), instanceName);
+                if (pod == null) {
+                    LOG.debug(String.format("[server-ping] Pod with name %s is already deleted.", instanceName));
+                    continue;
+                }
 
-            Instant createdAt = Constants.KUBERNETES_POD_CREATION_TIME_FORMAT.parse(pod.getMetadata().getCreationTimestamp(), Instant::from);
+                Instant createdAt = Constants.KUBERNETES_POD_CREATION_TIME_FORMAT.parse(pod.getMetadata().getCreationTimestamp(), Instant::from);
 
-            if (clock.now().isAfter(createdAt.plus(period))) {
-                unregisteredInstances.register(kubernetesInstanceFactory.fromKubernetesPod(pod));
+                if (clock.now().isAfter(createdAt.plus(period))) {
+                    unregisteredInstances.register(kubernetesInstanceFactory.fromKubernetesPod(pod));
+                }
             }
         }
 

src/main/java/cd/go/contrib/elasticagent/KubernetesClientFactory.java

@@ -21,6 +21,7 @@
 import io.fabric8.kubernetes.client.KubernetesClientBuilder;
 
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static cd.go.contrib.elasticagent.KubernetesPlugin.LOG;
 import static cd.go.contrib.elasticagent.utils.Util.isBlank;
@@ -29,13 +30,13 @@
 
 public class KubernetesClientFactory {
     private static final KubernetesClientFactory KUBERNETES_CLIENT_FACTORY = new KubernetesClientFactory();
-    private final Clock clock;
-    private KubernetesClient client;
-    private PluginSettings clusterProfileConfigurations;
-    private long clientCreatedTime;
-    private long kubernetesClientRecycleIntervalInMinutes = -1;
     public static final String CLIENT_RECYCLE_SYSTEM_PROPERTY_KEY = "go.kubernetes.elastic-agent.plugin.client.recycle.interval.in.minutes";
 
+    private final Clock clock;
+
+    private volatile CachedClient client;
+    private volatile long kubernetesClientRecycleIntervalInMinutes = -1;
+
     KubernetesClientFactory() {
         this(Clock.DEFAULT);
     }
@@ -50,32 +51,30 @@ public static KubernetesClientFactory instance() {
         return KUBERNETES_CLIENT_FACTORY;
     }
 
-    public synchronized KubernetesClient client(PluginSettings clusterProfileConfigurations) {
+    public synchronized CachedClient client(PluginSettings clusterProfileConfigurations) {
         clearOutClientOnTimer();
-        if (clusterProfileConfigurations.equals(this.clusterProfileConfigurations) && this.client != null) {
+        if (this.client != null && clusterProfileConfigurations.equals(this.client.clusterProfileConfigurations)) {
             LOG.debug("Using previously created client.");
+            this.client.leases.incrementAndGet();
             return this.client;
         }
 
         LOG.debug(format("Creating a new client because {0}.", (client == null) ? "client is null" : "cluster profile configurations has changed"));
-        this.clusterProfileConfigurations = clusterProfileConfigurations;
+        clearOutExistingClient();
         this.client = createClientFor(clusterProfileConfigurations);
-        this.clientCreatedTime = this.clock.now().toEpochMilli();
         LOG.debug("New client is created.");
 
         return this.client;
     }
 
     private void clearOutClientOnTimer() {
-        long currentTime = this.clock.now().toEpochMilli();
-        long differenceInMinutes = TimeUnit.MILLISECONDS.toMinutes(currentTime - this.clientCreatedTime);
-        if (differenceInMinutes > getKubernetesClientRecycleInterval()) {
+        if (client != null && TimeUnit.MILLISECONDS.toMinutes(this.clock.now().toEpochMilli() - this.client.clientCreatedTime) > getKubernetesClientRecycleInterval()) {
             LOG.info("Recycling kubernetes client on timer...");
             clearOutExistingClient();
         }
     }
 
-    private KubernetesClient createClientFor(PluginSettings pluginSettings) {
+    private CachedClient createClientFor(PluginSettings pluginSettings) {
         Config config = Config.autoConfigure(null);
 
         setIfNotBlank(config::setMasterUrl, pluginSettings.getClusterUrl());
@@ -84,13 +83,12 @@ private KubernetesClient createClientFor(PluginSettings pluginSettings) {
         setIfNotBlank(config::setCaCertData, pluginSettings.getCaCertData());
         config.setRequestTimeout(pluginSettings.getClusterRequestTimeout());
 
-        return new KubernetesClientBuilder().withConfig(config).build();
+        return new CachedClient(new KubernetesClientBuilder().withConfig(config).build(), pluginSettings);
     }
 
-    public void clearOutExistingClient() {
+    public synchronized void clearOutExistingClient() {
         if (this.client != null) {
-            LOG.debug("Terminating existing kubernetes client...");
-            this.client.close();
+            this.client.closeIfUnused();
             this.client = null;
         }
     }
@@ -109,12 +107,63 @@ private long getKubernetesClientRecycleInterval() {
         }
 
         try {
-            this.kubernetesClientRecycleIntervalInMinutes = Integer.valueOf(property);
+            this.kubernetesClientRecycleIntervalInMinutes = Integer.parseInt(property);
         } catch (Exception e) {
             //set default value to 10 minutes when parsing user input fails
             this.kubernetesClientRecycleIntervalInMinutes = 10;
         }
 
         return this.kubernetesClientRecycleIntervalInMinutes;
     }
+
+    public class CachedClient implements AutoCloseable {
+
+        private final KubernetesClient client;
+        private final PluginSettings clusterProfileConfigurations;
+        private final AtomicInteger leases = new AtomicInteger(1);
+        private final long clientCreatedTime;
+        private volatile boolean closed;
+
+        CachedClient(KubernetesClient client, PluginSettings clusterProfileConfigurations) {
+            this.client = client;
+            this.clusterProfileConfigurations = clusterProfileConfigurations;
+            this.clientCreatedTime = KubernetesClientFactory.this.clock.now().toEpochMilli();
+        }
+
+        public KubernetesClient get() {
+            return client;
+        }
+
+        public int leases() {
+            return leases.get();
+        }
+
+        public boolean isClosed() {
+            return closed;
+        }
+
+        @Override
+        public void close() {
+            releaseLease();
+        }
+
+        private void releaseLease() {
+            // Close the client only if it is not the same as the one we have cached
+            if (leases.decrementAndGet() == 0 && this != KubernetesClientFactory.this.client && !closed) {
+                closeUnderlyingClient();
+            }
+        }
+
+        public void closeIfUnused() {
+            if (leases() == 0 && !closed) {
+                closeUnderlyingClient();
+            }
+        }
+
+        private void closeUnderlyingClient() {
+            LOG.debug("Terminating existing kubernetes client...");
+            client.close();
+            closed = true;
+        }
+    }
 }

src/main/java/cd/go/contrib/elasticagent/executors/AgentStatusReportExecutor.java

@@ -55,26 +55,26 @@ public GoPluginApiResponse execute() {
         String elasticAgentId = request.getElasticAgentId();
         JobIdentifier jobIdentifier = request.getJobIdentifier();
         LOG.info(format("[status-report] Generating status report for agent: {0} with job: {1}", elasticAgentId, jobIdentifier));
-        KubernetesClient client = factory.client(request.clusterProfileProperties());
+        try (KubernetesClientFactory.CachedClient client = factory.client(request.clusterProfileProperties())) {
+            try {
+                Pod pod;
+                if (!isBlank(elasticAgentId)) {
+                    pod = findPodUsingElasticAgentId(elasticAgentId, client.get());
+                } else {
+                    pod = findPodUsingJobIdentifier(jobIdentifier, client.get());
+                }
 
-        try {
-            Pod pod;
-            if (!isBlank(elasticAgentId)) {
-                pod = findPodUsingElasticAgentId(elasticAgentId, client);
-            } else {
-                pod = findPodUsingJobIdentifier(jobIdentifier, client);
-            }
-
-            KubernetesElasticAgent elasticAgent = KubernetesElasticAgent.fromPod(client, pod, jobIdentifier);
+                KubernetesElasticAgent elasticAgent = KubernetesElasticAgent.fromPod(client.get(), pod, jobIdentifier);
 
-            final String statusReportView = statusReportViewBuilder.build(statusReportViewBuilder.getTemplate("agent-status-report.template.ftlh"), elasticAgent);
+                final String statusReportView = statusReportViewBuilder.build(statusReportViewBuilder.getTemplate("agent-status-report.template.ftlh"), elasticAgent);
 
-            final JsonObject responseJSON = new JsonObject();
-            responseJSON.addProperty("view", statusReportView);
+                final JsonObject responseJSON = new JsonObject();
+                responseJSON.addProperty("view", statusReportView);
 
-            return DefaultGoPluginApiResponse.success(responseJSON.toString());
-        } catch (Exception e) {
-            return StatusReportGenerationErrorHandler.handle(statusReportViewBuilder, e);
+                return DefaultGoPluginApiResponse.success(responseJSON.toString());
+            } catch (Exception e) {
+                return StatusReportGenerationErrorHandler.handle(statusReportViewBuilder, e);
+            }
         }
     }
 

src/main/java/cd/go/contrib/elasticagent/executors/ClusterStatusReportExecutor.java

@@ -25,7 +25,6 @@
 import com.thoughtworks.go.plugin.api.response.DefaultGoPluginApiResponse;
 import com.thoughtworks.go.plugin.api.response.GoPluginApiResponse;
 import freemarker.template.Template;
-import io.fabric8.kubernetes.client.KubernetesClient;
 
 import static cd.go.contrib.elasticagent.KubernetesPlugin.LOG;
 
@@ -49,8 +48,10 @@ public ClusterStatusReportExecutor(ClusterStatusReportRequest request, PluginSta
     public GoPluginApiResponse execute() {
         try {
             LOG.info("[status-report] Generating status report.");
-            KubernetesClient client = factory.client(request.clusterProfileProperties());
-            final KubernetesCluster kubernetesCluster = new KubernetesCluster(client);
+            final KubernetesCluster kubernetesCluster;
+            try (KubernetesClientFactory.CachedClient client = factory.client(request.clusterProfileProperties())) {
+                kubernetesCluster = new KubernetesCluster(client.get());
+            }
             final Template template = statusReportViewBuilder.getTemplate("status-report.template.ftlh");
             final String statusReportView = statusReportViewBuilder.build(template, kubernetesCluster);
 

src/test/java/cd/go/contrib/elasticagent/KubernetesAgentInstancesIntegrationTest.java

@@ -24,9 +24,9 @@
 import io.fabric8.kubernetes.client.dsl.PodResource;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.Answers;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
-import org.mockito.Mockito;
 import org.mockito.stubbing.Answer;
 
 import java.util.ArrayList;
@@ -45,7 +45,7 @@
 
 public class KubernetesAgentInstancesIntegrationTest {
 
-    @Mock
+    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
     private KubernetesClientFactory mockedKubernetesClientFactory;
 
     @Mock
@@ -72,7 +72,7 @@ public class KubernetesAgentInstancesIntegrationTest {
     public void setUp() {
         openMocks(this);
         kubernetesAgentInstances = new KubernetesAgentInstances(mockedKubernetesClientFactory);
-        when(mockedKubernetesClientFactory.client(any())).thenReturn(mockKubernetesClient);
+        when(mockedKubernetesClientFactory.client(any()).get()).thenReturn(mockKubernetesClient);
 
         when(pods.resource(any(Pod.class))).thenAnswer((Answer<PodResource>) invocation -> {
             Object[] args = invocation.getArguments();

src/test/java/cd/go/contrib/elasticagent/KubernetesAgentInstancesTest.java

@@ -26,6 +26,7 @@
 import io.fabric8.kubernetes.client.dsl.PodResource;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.Answers;
 import org.mockito.InOrder;
 import org.mockito.Mock;
 
@@ -43,7 +44,7 @@
 import static org.mockito.MockitoAnnotations.openMocks;
 
 public class KubernetesAgentInstancesTest {
-    @Mock
+    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
     KubernetesClientFactory factory;
 
     @Mock
@@ -79,7 +80,7 @@ public void setUp() {
         when(mockCreateAgentRequest.properties()).thenReturn(testProperties);
         when(mockPluginSettings.getMaxPendingPods()).thenReturn(10);
         when(mockPluginSettings.getClusterRequestTimeout()).thenReturn(10000);
-        when(factory.client(mockPluginSettings)).thenReturn(mockKubernetesClient);
+        when(factory.client(mockPluginSettings).get()).thenReturn(mockKubernetesClient);
         JobIdentifier jobId = new JobIdentifier("test", 1L, "Test pipeline", "test name", "1", "test job", 100L);
         when(mockCreateAgentRequest.jobIdentifier()).thenReturn(jobId);