fix: clean up deployment artifacts

muboheng · muboheng · commit 78fdfac7c43c · 2025-09-22T18:11:08.000-07:00
diff --git a/src/main/java/com/aws/greengrass/componentmanager/ComponentManager.java b/src/main/java/com/aws/greengrass/componentmanager/ComponentManager.java
diff --git a/src/main/java/com/aws/greengrass/deployment/DefaultDeploymentTask.java b/src/main/java/com/aws/greengrass/deployment/DefaultDeploymentTask.java
@@ -112,6 +112,8 @@ public DeploymentResult call() throws InterruptedException {
         Future<List<ComponentIdentifier>> resolveDependenciesFuture = null;
         Future<Void> preparePackagesFuture = null;
         Future<DeploymentResult> deploymentMergeFuture = null;
+        DeploymentResult deploymentResult = null;
+        List<ComponentIdentifier> desiredPackages = null;
         DeploymentDocument deploymentDocument = deployment.getDeploymentDocumentObj();
         try {
             logger.atInfo().setEventType(DEPLOYMENT_TASK_EVENT_TYPE)
@@ -131,7 +133,7 @@ public DeploymentResult call() throws InterruptedException {
             resolveDependenciesFuture = executorService.submit(() ->
                     dependencyResolver.resolveDependencies(deploymentDocument, nonTargetGroupsToRootPackagesMap));
 
-            List<ComponentIdentifier> desiredPackages = resolveDependenciesFuture.get();
+            desiredPackages = resolveDependenciesFuture.get();
 
             // download configuration if large
             List<String> requiredCapabilities = deploymentDocument.getRequiredCapabilities();
@@ -163,23 +165,18 @@ public DeploymentResult call() throws InterruptedException {
             //   on the device before the deployment started, and finding one of type nucleus.
             Optional<DeploymentPackageConfiguration> incomingNucleusComponentConfiguration =
                     deploymentDocument.getDeploymentPackageConfigurationList() == null ? Optional.empty() :
-                    deploymentDocument.getDeploymentPackageConfigurationList().stream()
-                        .filter(c -> c.getPackageName().equals(deviceConfiguration.getNucleusComponentName()))
-                        .findAny();
+                            deploymentDocument.getDeploymentPackageConfigurationList().stream()
+                                    .filter(c -> c.getPackageName()
+                                            .equals(deviceConfiguration.getNucleusComponentName()))
+                                    .findAny();
 
             if (incomingNucleusComponentConfiguration.isPresent()
-                    && incomingNucleusComponentConfiguration
-                        .get()
-                        .getConfigurationUpdateOperation() != null
-                    && incomingNucleusComponentConfiguration
-                        .get()
-                        .getConfigurationUpdateOperation()
-                        .getValueToMerge() != null
-                    && incomingNucleusComponentConfiguration
-                        .get()
-                        .getConfigurationUpdateOperation()
-                        .getValueToMerge()
-                        .containsKey(DEVICE_PARAM_DEPLOYMENT_CONFIGURATION_TIME_SOURCE)) {
+                    && incomingNucleusComponentConfiguration.get().getConfigurationUpdateOperation() != null
+                    && incomingNucleusComponentConfiguration.get().getConfigurationUpdateOperation()
+                    .getValueToMerge() != null
+                    && incomingNucleusComponentConfiguration.get().getConfigurationUpdateOperation()
+                    .getValueToMerge()
+                    .containsKey(DEVICE_PARAM_DEPLOYMENT_CONFIGURATION_TIME_SOURCE)) {
                 logger.atDebug(DEPLOYMENT_TASK_EVENT_TYPE).log(
                         "Incoming nucleus component configuration contains deployment configuration time source");
                 String incomingDeploymentConfigurationTimeSource = Coerce.toString(
@@ -223,28 +220,41 @@ public DeploymentResult call() throws InterruptedException {
             // Block this without timeout because it can take a long time for the device to update the config
             // (if it's not in a safe window).
             DeploymentResult result = deploymentMergeFuture.get();
+            deploymentResult = result;
 
             logger.atInfo(DEPLOYMENT_TASK_EVENT_TYPE).setEventType(DEPLOYMENT_TASK_EVENT_TYPE)
                     .log("Finished deployment task");
 
-            componentManager.cleanupStaleVersions();
             return result;
         } catch (PackageLoadingException | DeploymentTaskFailureException | IOException e) {
             logger.atError().setCause(e).log("Error occurred while processing deployment");
-            return new DeploymentResult(DeploymentResult.DeploymentStatus.FAILED_NO_STATE_CHANGE, e);
+            deploymentResult = new DeploymentResult(DeploymentResult.DeploymentStatus.FAILED_NO_STATE_CHANGE, e);
+            return deploymentResult;
         } catch (ExecutionException e) {
             logger.atError().setCause(e).log("Error occurred while processing deployment");
             Throwable t = e.getCause();
             if (t instanceof InterruptedException) {
                 throw (InterruptedException) t;
             } else {
-                return new DeploymentResult(DeploymentResult.DeploymentStatus.FAILED_NO_STATE_CHANGE, t);
+                deploymentResult = new DeploymentResult(DeploymentResult.DeploymentStatus.FAILED_NO_STATE_CHANGE, t);
+                return deploymentResult;
             }
         } catch (InterruptedException e) {
             // DeploymentTask got interrupted while performing a blocking step.
             cancelDeploymentTask(resolveDependenciesFuture, preparePackagesFuture, deploymentMergeFuture);
             // Populate the exception up to the stack
             throw e;
+        } finally {
+            // Clean up stale artifacts from previous deployments after downloading deployment document
+            // but before preparing packages, so we don't interfere with current deployment
+            Map<String, String> failedDeploymentVersions = null;
+            if (desiredPackages != null) {
+                failedDeploymentVersions = desiredPackages.stream()
+                        .collect(Collectors.toMap(
+                                ComponentIdentifier::getName,
+                                ci -> ci.getVersion().getValue()));
+            }
+            componentManager.cleanupStaleVersions(deploymentResult, failedDeploymentVersions);
         }
     }
 
diff --git a/src/main/java/com/aws/greengrass/deployment/KernelUpdateDeploymentTask.java b/src/main/java/com/aws/greengrass/deployment/KernelUpdateDeploymentTask.java
@@ -79,6 +79,7 @@ public DeploymentResult call() {
             result = deploymentResultCompletableFuture.get();
         } catch (InterruptedException | ExecutionException | CancellationException e) {
             // nothing to report when deployment is cancelled
+            componentManager.cleanupStaleVersions();
             return null;
         }
         componentManager.cleanupStaleVersions();
diff --git a/src/main/java/com/aws/greengrass/lifecyclemanager/UpdateSystemPolicyService.java b/src/main/java/com/aws/greengrass/lifecyclemanager/UpdateSystemPolicyService.java
@@ -209,12 +209,14 @@ private long getTimeToReCheck(long timeout, String deploymentId,
                     try {
                         DeferComponentUpdateRequest deferRequest = fut.get();
                         if (deploymentId.equals(deferRequest.getDeploymentId())) {
-                            long timeToRecheck = currentTimeMillis + deferRequest.getRecheckAfterMs();
+                            Long recheckAfterMs = deferRequest.getRecheckAfterMs();
+                            long recheckValue = recheckAfterMs == null ? 0 : recheckAfterMs;
+                            long timeToRecheck = currentTimeMillis + recheckValue;
                             if (timeToRecheck > maxTimeToReCheck) {
                                 maxTimeToReCheck = timeToRecheck;
                                 logger.atInfo().setEventType("service-update-deferred")
                                         .log("deferred for {} millis with message {}",
-                                                deferRequest.getRecheckAfterMs(),
+                                                recheckValue,
                                                 deferRequest.getMessage());
                             }
                         } else {
diff --git a/src/test/java/com/aws/greengrass/componentmanager/ComponentManagerCleanupTest.java b/src/test/java/com/aws/greengrass/componentmanager/ComponentManagerCleanupTest.java
@@ -0,0 +1,153 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.aws.greengrass.componentmanager;
+
+import com.aws.greengrass.config.Topic;
+import com.aws.greengrass.config.Topics;
+import com.aws.greengrass.deployment.model.DeploymentResult;
+import com.aws.greengrass.lifecyclemanager.GreengrassService;
+import com.aws.greengrass.lifecyclemanager.Kernel;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+import static com.aws.greengrass.componentmanager.KernelConfigResolver.PREV_VERSION_CONFIG_KEY;
+import static com.aws.greengrass.componentmanager.KernelConfigResolver.VERSION_CONFIG_KEY;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+class ComponentManagerCleanupTest {
+
+    @Mock
+    private Kernel kernel;
+    @Mock
+    private ComponentStore componentStore;
+
+    private ComponentManager componentManager;
+
+    @BeforeEach
+    void setup() {
+        componentManager = new ComponentManager(null, null, null, componentStore, kernel, null, null, null);
+    }
+
+    @Test
+    void GIVEN_successful_deployment_WHEN_getVersionsToKeep_THEN_keeps_current_and_previous() {
+        GreengrassService service = createServiceWithBothVersions("TestComponent", "2.0.0", "1.0.0");
+        when(kernel.orderedDependencies()).thenReturn(Collections.singletonList(service));
+        
+        DeploymentResult result = new DeploymentResult(DeploymentResult.DeploymentStatus.SUCCESSFUL, null);
+        Map<String, Set<String>> versions = componentManager.getVersionsToKeep(result, null);
+        
+        assertEquals(new HashSet<>(Arrays.asList("2.0.0", "1.0.0")), versions.get("TestComponent"));
+    }
+
+    @Test
+    void GIVEN_failed_deployment_WHEN_getVersionsToKeep_THEN_keeps_three_versions() {
+        GreengrassService service = createServiceWithBothVersions("TestComponent", "3.0.0", "2.0.0");
+        when(kernel.orderedDependencies()).thenReturn(Collections.singletonList(service));
+        
+        // Mock available versions in component store
+        Map<String, Set<String>> availableVersions = new HashMap<>();
+        availableVersions.put("TestComponent", new HashSet<>(Arrays.asList("1.0.0", "2.0.0", "3.0.0")));
+        when(componentStore.listAvailableComponentVersions()).thenReturn(availableVersions);
+        
+        DeploymentResult result = new DeploymentResult(DeploymentResult.DeploymentStatus.FAILED_NO_STATE_CHANGE, 
+                new RuntimeException());
+        Map<String, Set<String>> versions = componentManager.getVersionsToKeep(result, null);
+        
+        // Should keep current (3.0.0 failed) + previous (2.0.0) + one more (1.0.0) = 3 total
+        assertEquals(new HashSet<>(Arrays.asList("3.0.0", "2.0.0", "1.0.0")), versions.get("TestComponent"));
+    }
+
+    @Test
+    void GIVEN_null_deployment_WHEN_getVersionsToKeep_THEN_keeps_current_and_previous() {
+        GreengrassService service = createServiceWithBothVersions("TestComponent", "2.0.0", "1.0.0");
+        when(kernel.orderedDependencies()).thenReturn(Collections.singletonList(service));
+        
+        Map<String, Set<String>> versions = componentManager.getVersionsToKeep(null, null);
+        
+        assertEquals(new HashSet<>(Arrays.asList("2.0.0", "1.0.0")), versions.get("TestComponent"));
+    }
+
+    @Test
+    void GIVEN_failed_deployment_with_version_info_WHEN_getVersionsToKeep_THEN_includes_failed_version() {
+        GreengrassService service = createServiceWithBothVersions("TestComponent", "1.0.0", "0.9.0");
+        when(kernel.orderedDependencies()).thenReturn(Arrays.asList(service));
+        
+        DeploymentResult failedResult = new DeploymentResult(DeploymentResult.DeploymentStatus.FAILED_NO_STATE_CHANGE, 
+                new RuntimeException("Test failure"));
+        Map<String, String> failedVersions = Collections.singletonMap("TestComponent", "2.0.0");
+        
+        Map<String, Set<String>> versions = componentManager.getVersionsToKeep(failedResult, failedVersions);
+        
+        assertEquals(new HashSet<>(Arrays.asList("1.0.0", "0.9.0", "2.0.0")), versions.get("TestComponent"));
+    }
+
+    @Test
+    void GIVEN_service_with_null_version_topics_WHEN_getVersionsToKeep_THEN_handles_gracefully() {
+        GreengrassService service = createServiceWithNullVersions("TestComponent");
+        when(kernel.orderedDependencies()).thenReturn(Collections.singletonList(service));
+        
+        Map<String, Set<String>> versions = componentManager.getVersionsToKeep(null, null);
+        
+        assertEquals(new HashSet<>(), versions.get("TestComponent"));
+    }
+
+    @Test
+    void GIVEN_failed_deployment_no_available_versions_WHEN_getVersionsToKeep_THEN_keeps_current_and_previous() {
+        GreengrassService service = createServiceWithBothVersions("TestComponent", "2.0.0", "1.0.0");
+        when(kernel.orderedDependencies()).thenReturn(Collections.singletonList(service));
+        
+        // Mock empty available versions in component store
+        when(componentStore.listAvailableComponentVersions()).thenReturn(new HashMap<>());
+        
+        DeploymentResult result = new DeploymentResult(DeploymentResult.DeploymentStatus.FAILED_NO_STATE_CHANGE, 
+                new RuntimeException());
+        Map<String, Set<String>> versions = componentManager.getVersionsToKeep(result, null);
+        
+        // Should only keep current + previous when no additional versions available
+        assertEquals(new HashSet<>(Arrays.asList("2.0.0", "1.0.0")), versions.get("TestComponent"));
+    }
+
+    private GreengrassService createServiceWithBothVersions(String name, String currentVersion, String previousVersion) {
+        GreengrassService service = mock(GreengrassService.class);
+        Topics config = mock(Topics.class);
+        Topic versionTopic = mock(Topic.class);
+        Topic prevVersionTopic = mock(Topic.class);
+        
+        when(service.getName()).thenReturn(name);
+        when(service.getServiceConfig()).thenReturn(config);
+        when(versionTopic.getOnce()).thenReturn(currentVersion);
+        when(prevVersionTopic.getOnce()).thenReturn(previousVersion);
+        when(config.find(VERSION_CONFIG_KEY)).thenReturn(versionTopic);
+        when(config.find(PREV_VERSION_CONFIG_KEY)).thenReturn(prevVersionTopic);
+        
+        return service;
+    }
+
+    private GreengrassService createServiceWithNullVersions(String name) {
+        GreengrassService service = mock(GreengrassService.class);
+        Topics config = mock(Topics.class);
+        
+        when(service.getName()).thenReturn(name);
+        when(service.getServiceConfig()).thenReturn(config);
+        when(config.find(VERSION_CONFIG_KEY)).thenReturn(null);
+        when(config.find(PREV_VERSION_CONFIG_KEY)).thenReturn(null);
+        
+        return service;
+    }
+}
diff --git a/src/test/java/com/aws/greengrass/deployment/DefaultDeploymentTaskCleanupTest.java b/src/test/java/com/aws/greengrass/deployment/DefaultDeploymentTaskCleanupTest.java
@@ -0,0 +1,73 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.aws.greengrass.deployment;
+
+import com.aws.greengrass.componentmanager.ComponentManager;
+import com.aws.greengrass.componentmanager.models.ComponentIdentifier;
+import com.aws.greengrass.deployment.model.DeploymentResult;
+import com.vdurmont.semver4j.Semver;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
+
+@ExtendWith(MockitoExtension.class)
+class DefaultDeploymentTaskCleanupTest {
+
+    @Mock
+    private ComponentManager componentManager;
+
+    @Test
+    void GIVEN_deployment_completes_WHEN_cleanup_called_THEN_passes_deployment_result() {
+        DeploymentResult result = new DeploymentResult(DeploymentResult.DeploymentStatus.SUCCESSFUL, null);
+        
+        // Simulate the cleanup call that happens in DefaultDeploymentTask
+        componentManager.cleanupStaleVersions(result, null);
+        
+        verify(componentManager).cleanupStaleVersions(eq(result), eq(null));
+    }
+
+    @Test
+    void GIVEN_failed_deployment_WHEN_cleanup_called_THEN_passes_failed_result_with_versions() {
+        DeploymentResult result = new DeploymentResult(DeploymentResult.DeploymentStatus.FAILED_NO_STATE_CHANGE, 
+                new RuntimeException("Test failure"));
+        
+        // Simulate the cleanup call that happens in DefaultDeploymentTask for failed deployments
+        Map<String, String> failedVersions = Collections.singletonMap("TestComponent", "1.0.0");
+        componentManager.cleanupStaleVersions(result, failedVersions);
+        
+        verify(componentManager).cleanupStaleVersions(eq(result), any(Map.class));
+    }
+
+    @Test
+    void GIVEN_desiredPackages_WHEN_extracting_versions_THEN_creates_correct_map() {
+        // Simulate the version extraction logic from DefaultDeploymentTask
+        List<ComponentIdentifier> desiredPackages = Arrays.asList(
+                new ComponentIdentifier("ComponentA", new Semver("1.0.0")),
+                new ComponentIdentifier("ComponentB", new Semver("2.1.0"))
+        );
+        
+        Map<String, String> failedDeploymentVersions = desiredPackages.stream()
+                .collect(Collectors.toMap(
+                        ComponentIdentifier::getName,
+                        ci -> ci.getVersion().getValue()));
+        
+        assertEquals("1.0.0", failedDeploymentVersions.get("ComponentA"));
+        assertEquals("2.1.0", failedDeploymentVersions.get("ComponentB"));
+        assertEquals(2, failedDeploymentVersions.size());
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -79,6 +79,7 @@ public DeploymentResult call() {`
`79`	`79`	`result = deploymentResultCompletableFuture.get();`
`80`	`80`	`} catch (InterruptedException \| ExecutionException \| CancellationException e) {`
`81`	`81`	`// nothing to report when deployment is cancelled`
	`82`	`+ componentManager.cleanupStaleVersions();`
`82`	`83`	`return null;`
`83`	`84`	`}`
`84`	`85`	`componentManager.cleanupStaleVersions();`