Update clean lost node without requiring a instanceId

src/main/java/com/google/jenkins/plugins/computeengine/CleanLostNodesWork.java

@@ -16,17 +16,23 @@
 
 package com.google.jenkins.plugins.computeengine;
 
-import static com.google.jenkins.plugins.computeengine.ComputeEngineCloud.CLOUD_ID_LABEL_KEY;
 import static java.util.Collections.emptyList;
 
 import com.google.api.services.compute.model.Instance;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
+import com.google.jenkins.plugins.computeengine.client.ComputeClientV2;
 import hudson.Extension;
 import hudson.model.PeriodicWork;
 import hudson.model.Slave;
 import java.io.IOException;
+import java.security.GeneralSecurityException;
+import java.time.LocalDateTime;
+import java.time.OffsetDateTime;
+import java.time.ZoneOffset;
+import java.time.format.DateTimeFormatter;
+import java.time.temporal.ChronoUnit;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -39,11 +45,27 @@
 @Symbol("cleanLostNodesWork")
 public class CleanLostNodesWork extends PeriodicWork {
     protected final Logger logger = Logger.getLogger(getClass().getName());
+    public static final String NODE_IN_USE_LABEL_KEY = "jenkins_node_last_refresh";
+    public static final long RECURRENCE_PERIOD = Long.parseLong(
+            System.getProperty(CleanLostNodesWork.class.getName() + ".recurrencePeriod", String.valueOf(HOUR)));
+
+    @VisibleForTesting
+    public static final int LOST_MULTIPLIER = 3;
+    /**
+     * The formatter for the label timestamp value as per google label format,
+     * "The value can only contain lowercase letters, numeric characters, underscores and dashes.
+     * The value can be at most 63 characters long. International characters are allowed".
+     */
+    private static final DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy_MM_dd't'HH_mm_ss_SSS'z'");
 
     /** {@inheritDoc} */
     @Override
     public long getRecurrencePeriod() {
-        return HOUR;
+        return RECURRENCE_PERIOD;
+    }
+
+    public static String getLastRefreshLabelVal() {
+        return formatter.format(OffsetDateTime.now(ZoneOffset.UTC));
     }
 
     /** {@inheritDoc} */
@@ -55,58 +77,106 @@
 
     private void cleanCloud(ComputeEngineCloud cloud) {
         logger.log(Level.FINEST, "Cleaning cloud " + cloud.getCloudName());
-        List<Instance> remoteInstances = findRemoteInstances(cloud);
+        ComputeClientV2 clientV2;
+        try {
+            clientV2 = cloud.getClientV2();
+        } catch (GeneralSecurityException | IOException ex) {
+            logger.log(Level.WARNING, "Error getting clientV2 for cloud " + cloud.getCloudName(), ex);
+            return;
+        }
+        List<Instance> remoteInstances = findRunningRemoteInstances(clientV2);
         Set<String> localInstances = findLocalInstances(cloud);
+        if (!(localInstances.isEmpty() || remoteInstances.isEmpty())) {
+            updateLocalInstancesLabel(clientV2, localInstances, remoteInstances);
+        }
         remoteInstances.stream()
                 .filter(remote -> isOrphaned(remote, localInstances))
                 .forEach(remote -> terminateInstance(remote, cloud));
     }
 
     private boolean isOrphaned(Instance remote, Set<String> localInstances) {
-        String instanceName = remote.getName();
-        logger.log(Level.FINEST, "Checking instance " + instanceName);
-        return !localInstances.contains(instanceName);
+        /* It is necessary to check if the remote instance is present in localInstances.
+           The `remote` instance has an old timestamp because it hasn't been fetched again
+           after the `updateLocalInstancesLabel` call, to avoid extra network calls.
+        */
+        if (localInstances.contains(remote.getName())) {
+            return false;
+        }
+        String nodeLastRefresh = remote.getLabels().get(NODE_IN_USE_LABEL_KEY);
+        if (nodeLastRefresh == null) {
+            return false;
+        }
+        OffsetDateTime lastRefresh =
+                LocalDateTime.parse(nodeLastRefresh, formatter).atOffset(ZoneOffset.UTC);
+        boolean isOrphan = lastRefresh
+                .plus(RECURRENCE_PERIOD * LOST_MULTIPLIER, ChronoUnit.MILLIS)
+                .isBefore(OffsetDateTime.now(ZoneOffset.UTC));
+        logger.log(
+                Level.FINEST,
+                "Instance " + remote.getName() + " last_refresh label value: " + nodeLastRefresh + ", isOrphan: "
+                        + isOrphan);
+        return isOrphan;
     }
 
     private void terminateInstance(Instance remote, ComputeEngineCloud cloud) {
         String instanceName = remote.getName();
-        logger.log(Level.INFO, "Remote instance " + instanceName + " not found locally, removing it");
+        logger.log(Level.INFO, "Removing orphaned instance: " + instanceName);
         try {
             cloud.getClient().terminateInstanceAsync(cloud.getProjectId(), remote.getZone(), instanceName);
         } catch (IOException ex) {
             logger.log(Level.WARNING, "Error terminating remote instance " + instanceName, ex);
         }
     }
 
     private List<ComputeEngineCloud> getClouds() {
         return Jenkins.get().clouds.stream()
                 .filter(cloud -> cloud instanceof ComputeEngineCloud)
                 .map(cloud -> (ComputeEngineCloud) cloud)
                 .collect(Collectors.toList());
     }
 
     private Set<String> findLocalInstances(ComputeEngineCloud cloud) {
-        return Jenkins.get().getNodes().stream()
+        var localInstances = Jenkins.get().getNodes().stream()
                 .filter(node -> node instanceof ComputeEngineInstance)
                 .map(node -> (ComputeEngineInstance) node)
                 .filter(node -> node.getCloud().equals(cloud))
                 .map(Slave::getNodeName)
                 .collect(Collectors.toSet());
+        logger.log(Level.FINEST, "Found " + localInstances.size() + " local instances");
+        return localInstances;
     }
 
-    private List<Instance> findRemoteInstances(ComputeEngineCloud cloud) {
-        Map<String, String> filterLabel = ImmutableMap.of(CLOUD_ID_LABEL_KEY, cloud.getInstanceId());
+    private List<Instance> findRunningRemoteInstances(ComputeClientV2 clientV2) {
         try {
-            return cloud.getClient().listInstancesWithLabel(cloud.getProjectId(), filterLabel).stream()
-                    .filter(instance -> shouldTerminateStatus(instance.getStatus()))
-                    .collect(Collectors.toList());
+            var remoteInstances = clientV2.retrieveInstanceByLabelKeyAndStatus(NODE_IN_USE_LABEL_KEY, "RUNNING");
+            logger.log(Level.FINEST, "Found " + remoteInstances.size() + " running remote instances");
+            return remoteInstances;
         } catch (IOException ex) {
             logger.log(Level.WARNING, "Error finding remote instances", ex);
             return emptyList();
         }
     }
 
-    private boolean shouldTerminateStatus(String status) {
-        return !status.equals("STOPPING");
+    /**
+     * Updates the label of the local instances to indicate they are still in use. The method makes N network calls
+     * for N local instances, couldn't find any bulk update apis.
+     */
+    private void updateLocalInstancesLabel(
+            ComputeClientV2 clientV2, Set<String> localInstances, List<Instance> remoteInstances) {
+        var remoteInstancesByName =
+                remoteInstances.stream().collect(Collectors.toMap(Instance::getName, instance -> instance));
+        var labelToUpdate = ImmutableMap.of(NODE_IN_USE_LABEL_KEY, getLastRefreshLabelVal());
+        for (String instanceName : localInstances) {
+            var remoteInstance = remoteInstancesByName.get(instanceName);
+            if (remoteInstance == null) {
+                continue;
+            }
+            try {
+                clientV2.updateInstanceLabels(remoteInstance, labelToUpdate);
+                logger.log(Level.FINEST, "Updated label for instance " + instanceName);
+            } catch (IOException e) {
+                logger.log(Level.WARNING, "Error updating label for instance " + instanceName, e);
+            }
+        }
     }
 }

src/main/java/com/google/jenkins/plugins/computeengine/ComputeEngineCloud.java

@@ -29,6 +29,7 @@
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMap;
 import com.google.jenkins.plugins.computeengine.client.ClientUtil;
+import com.google.jenkins.plugins.computeengine.client.ComputeClientV2;
 import com.google.jenkins.plugins.credentials.oauth.GoogleOAuth2Credentials;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import hudson.Extension;
@@ -48,6 +49,7 @@
 import hudson.util.ListBoxModel;
 import java.io.IOException;
 import java.io.PrintStream;
+import java.security.GeneralSecurityException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Iterator;
@@ -90,6 +92,7 @@
     private List<InstanceConfiguration> configurations;
 
     private transient volatile ComputeClient client;
+    private transient volatile ComputeClientV2 clientV2;
     private boolean noDelayProvisioning;
 
     @DataBoundConstructor
@@ -161,6 +164,8 @@
 
                 // Apply a label that identifies the name of this instance configuration
                 configuration.appendLabel(CONFIG_LABEL_KEY, configuration.getNamePrefix());
+                configuration.appendLabel(
+                        CleanLostNodesWork.NODE_IN_USE_LABEL_KEY, CleanLostNodesWork.getLastRefreshLabelVal());
             }
         }
         setInstanceId(instanceId);
@@ -209,6 +214,17 @@
         return client;
     }
 
+    public ComputeClientV2 getClientV2() throws IOException, GeneralSecurityException {
+        if (clientV2 == null) {
+            synchronized (this) {
+                if (clientV2 == null) {
+                    clientV2 = ClientUtil.createComputeClientV2(projectId, credentialsId);
+                }
+            }
+        }
+        return clientV2;
+    }
+
     /**
      * Set configurations for this cloud.
      *

src/main/java/com/google/jenkins/plugins/computeengine/client/ClientUtil.java

@@ -4,7 +4,10 @@
 import com.cloudbees.plugins.credentials.CredentialsProvider;
 import com.cloudbees.plugins.credentials.domains.DomainRequirement;
 import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport;
 import com.google.api.client.http.HttpTransport;
+import com.google.api.client.json.jackson2.JacksonFactory;
+import com.google.api.services.compute.Compute;
 import com.google.cloud.graphite.platforms.plugin.client.ClientFactory;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
@@ -19,6 +22,7 @@
 import java.security.GeneralSecurityException;
 import java.util.List;
 import java.util.Optional;
+import jenkins.model.Jenkins;
 
 /** Utilities for using the gcp-plugin-core clients. */
 public class ClientUtil {
@@ -88,4 +92,17 @@
     private static Credential getGoogleCredential(GoogleRobotCredentials credentials) throws GeneralSecurityException {
         return credentials.getGoogleCredential(new ComputeEngineScopeRequirement());
     }
+
+    public static ComputeClientV2 createComputeClientV2(String projectId, String credentialsId)
+            throws GeneralSecurityException, IOException {
+        Credential httpRequestInitializer = ClientUtil.getGoogleCredential(
+                ClientUtil.getRobotCredentials(Jenkins.get(), ImmutableList.of(), credentialsId));
+        Compute compute = new Compute.Builder(
+                        GoogleNetHttpTransport.newTrustedTransport(), new JacksonFactory(), httpRequestInitializer)
+                .setGoogleClientRequestInitializer(request ->
+                        request.setRequestHeaders(request.getRequestHeaders().setUserAgent(APPLICATION_NAME)))
+                .setApplicationName(ClientUtil.APPLICATION_NAME)
+                .build();
+        return new ComputeClientV2(projectId, compute);
+    }
 }

src/main/java/com/google/jenkins/plugins/computeengine/client/ComputeClientV2.java

@@ -0,0 +1,84 @@
+package com.google.jenkins.plugins.computeengine.client;
+
+import com.google.api.services.compute.Compute;
+import com.google.api.services.compute.model.Instance;
+import com.google.api.services.compute.model.InstancesScopedList;
+import com.google.api.services.compute.model.InstancesSetLabelsRequest;
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import lombok.Getter;
+
+/**
+ * Extends {@link com.google.cloud.graphite.platforms.plugin.client.ComputeClient} with additional functionalities.
+ * <p>This class serves as a venue for implementing features not available in the archived Graphite Java library
+ * (<a href="https://github.com/GoogleCloudPlatform/gcp-plugin-core-java">gcp-plugin-core-java</a>, last updated in December 2019).
+ * Consideration for the gradual evolution of this class is suggested, including the re-implementation of methods
+ * currently utilized from the Graphite library, to ensure dependency solely on the Google API Java Client Services
+ * library (<a href="https://github.com/googleapis/google-api-java-client-services">google-api-java-client-services</a>).
+ * This approach aims to eventually eliminate the reliance on the Graphite library.
+ */
+public class ComputeClientV2 {
+
+    private final String projectId;
+
+    @Getter
+    private final Compute compute;
+
+    public ComputeClientV2(String projectId, Compute compute) {
+        this.projectId = projectId;
+        this.compute = compute;
+    }
+
+    /**
+     * Updates the labels of a specified {@code instance} by merging or replacing them with {@code newLabels}.
+     * <p>This method adds any new labels found in {@code newLabels} to the instance's existing labels and updates
+     * the values of any existing labels if they are also present in {@code newLabels}. Labels existing on the instance
+     * that are not in {@code newLabels} remain unchanged. This operation can only result in the addition of new
+     * labels or the modification of existing ones.
+     *
+     * @param instance  the instance whose labels are to be updated; must not be {@code null}
+     * @param newLabels the new labels to be merged with or replace the existing labels of the instance; must not be {@code null}
+     * @throws IOException if an I/O error occurs during the label update process.
+     */
+    public void updateInstanceLabels(Instance instance, Map<String, String> newLabels) throws IOException {
+        var allLabels = instance.getLabels();
+        allLabels.putAll(newLabels);
+        var labelsRequest = new InstancesSetLabelsRequest()
+                .setLabels(allLabels)
+                .setLabelFingerprint(instance.getLabelFingerprint());
+        String zoneLink = instance.getZone();
+        String zone = zoneLink.substring(zoneLink.lastIndexOf("/") + 1);
+        compute.instances()
+                .setLabels(projectId, zone, instance.getName(), labelsRequest)
+                .execute();
+    }
+
+    /**
+     * Fetches instances by label key existence and status.
+     * <p>Applies Google Compute Engine aggregated list syntax for filtering:
+     * <a href="https://cloud.google.com/compute/docs/reference/rest/v1/instances/aggregatedList">aggregatedList API</a>.
+     *
+     * @param key the non-empty label key to filter by.
+     * @param status the instance status (RUNNING, STOPPING, etc.) as defined in:
+     * <a href="https://cloud.google.com/compute/docs/instances/instance-lifecycle#instance-states">Instance States</a>.
+     * @return List of {@link Instance} matching criteria, or empty list if none.
+     * @throws IOException for communication issues with Compute Engine API.
+     */
+    public List<Instance> retrieveInstanceByLabelKeyAndStatus(String key, String status) throws IOException {
+        String filter = "labels." + key + ":*" + " AND status=" + status;
+        var response =
+                compute.instances().aggregatedList(projectId).setFilter(filter).execute();
+        var items = response.getItems();
+        if (items == null) {
+            return List.of();
+        }
+        return items.values().stream()
+                .map(InstancesScopedList::getInstances)
+                .filter(Objects::nonNull)
+                .flatMap(List::stream)
+                .collect(Collectors.toList());
+    }
+}