ClusterRemovableAgentRemover { @VisibleForTesting Completable doRemoveAgents() { return Completable.defer(() -> { if (!configuration.isRemovingAgentsEnabled()) { logger.debug("Removing agents is not enabled"); return Completable.complete(); } long now = titusRuntime.getClock().wallTime(); List<AgentInstanceGroup> eligibleInstanceGroups = agentManagementService.getInstanceGroups().stream() .filter(agentManagementService::isOwnedByFenzo) .filter(ig -> ig.getLifecycleStatus().getState() == InstanceGroupLifecycleState.Active || ig.getLifecycleStatus().getState() == InstanceGroupLifecycleState.PhasedOut) .collect(Collectors.toList()); logger.debug("Eligible instance groups: {}", eligibleInstanceGroups); long totalAgentsToRemove = 0; Map<AgentInstanceGroup, List<AgentInstance>> agentInstancesPerInstanceGroup = new HashMap<>(); for (AgentInstanceGroup instanceGroup : eligibleInstanceGroups) { List<AgentInstance> agentInstances = agentManagementService.getAgentInstances(instanceGroup.getId()) .stream() .filter(i -> { Map<String, String> attributes = i.getAttributes(); if (attributes.containsKey(ClusterOperationsAttributes.NOT_REMOVABLE)) { return false; } String removableTimestampValue = attributes.get(ClusterOperationsAttributes.REMOVABLE); if (!Strings.isNullOrEmpty(removableTimestampValue)) { Long parsedRemovableTimestamp = Longs.tryParse(removableTimestampValue); long removableTimestamp = parsedRemovableTimestamp == null ? 0L : parsedRemovableTimestamp; return hasTimeElapsed(removableTimestamp, now, configuration.getAgentInstanceRemovableGracePeriodMs()); } return false; }) .collect(Collectors.toList()); totalAgentsToRemove += agentInstances.size(); agentInstancesPerInstanceGroup.put(instanceGroup, agentInstances); } totalAgentsToRemoveGauge.set(totalAgentsToRemove); logger.debug("Agent instances per instance group: {}", agentInstancesPerInstanceGroup); if (totalAgentsToRemove <= 0) { totalEligibleAgentsToRemoveGauge.set(0); return Completable.complete(); } long totalEligibleAgentsToRemove = 0; Map<AgentInstanceGroup, List<AgentInstance>> eligibleAgentInstancesPerInstanceGroup = new HashMap<>(); Map<String, Long> numberOfTasksOnAgents = getNumberOfTasksOnAgents(); logger.debug("numberOfTasksOnAgents: {}", numberOfTasksOnAgents); for (Map.Entry<AgentInstanceGroup, List<AgentInstance>> entry : agentInstancesPerInstanceGroup.entrySet()) { List<AgentInstance> eligibleAgentInstances = entry.getValue().stream() .filter(i -> agentInstancesBeingTerminated.getIfPresent(i.getId()) == null) .filter(i -> numberOfTasksOnAgents.getOrDefault(i.getId(), 0L) <= 0) .collect(Collectors.toList()); totalEligibleAgentsToRemove += eligibleAgentInstances.size(); eligibleAgentInstancesPerInstanceGroup.put(entry.getKey(), eligibleAgentInstances); } totalEligibleAgentsToRemoveGauge.set(totalEligibleAgentsToRemove); logger.debug("Eligible agent instances per instance group: {}", eligibleAgentInstancesPerInstanceGroup); if (totalEligibleAgentsToRemove <= 0) { return Completable.complete(); } long maxTokensToTake = Math.min(TOKEN_BUCKET_CAPACITY, totalEligibleAgentsToRemove); Optional<Pair<Long, ImmutableTokenBucket>> takeOpt = this.lastTokenBucket.tryTake(1, maxTokensToTake); if (takeOpt.isPresent()) { Pair<Long, ImmutableTokenBucket> takePair = takeOpt.get(); this.lastTokenBucket = takePair.getRight(); long tokensAvailable = takePair.getLeft(); long tokensUsed = 0; logger.debug("Attempting to terminate {} agent instances", tokensAvailable); List<Completable> actions = new ArrayList<>(); for (Map.Entry<AgentInstanceGroup, List<AgentInstance>> entry : eligibleAgentInstancesPerInstanceGroup.entrySet()) { long tokensRemaining = tokensAvailable - tokensUsed; if (tokensRemaining <= 0) { break; } List<AgentInstance> agentInstances = entry.getValue(); if (agentInstances.isEmpty()) { continue; } AgentInstanceGroup instanceGroup = entry.getKey(); String instanceGroupId = instanceGroup.getId(); int agentCountEligibleToTerminateInInstanceGroup = instanceGroup.getCurrent() - instanceGroup.getMin(); int agentCountToTerminate = Math.min((int) tokensRemaining, agentCountEligibleToTerminateInInstanceGroup); if (agentCountToTerminate > 0) { List<String> agentInstanceIdsToTerminate = agentInstances.stream() .limit(agentCountToTerminate) .map(AgentInstance::getId) .collect(Collectors.toList()); logger.info("Terminating in instance group: {} agent instances({}): {}", instanceGroupId, agentInstanceIdsToTerminate.size(), agentInstanceIdsToTerminate); actions.add(createTerminateAgentsCompletable(instanceGroupId, agentInstanceIdsToTerminate)); tokensUsed += agentInstanceIdsToTerminate.size(); } } totalAgentsBeingRemovedGauge.set(tokensUsed); return Completable.concat(actions); } return Completable.complete(); }).doOnCompleted(() -> logger.debug("Completed cluster agent removal")) .timeout(REMOVE_AGENTS_COMPLETABLE_TIMEOUT_MS, TimeUnit.MILLISECONDS); } @Inject  ClusterRemovableAgentRemover(TitusRuntime titusRuntime,
                                        ClusterOperationsConfiguration configuration,
                                        AgentManagementService agentManagementService,
                                        V3JobOperations v3JobOperations);  ClusterRemovableAgentRemover(TitusRuntime titusRuntime,
                                        ClusterOperationsConfiguration configuration,
                                        AgentManagementService agentManagementService,
                                        V3JobOperations v3JobOperations,
                                        Scheduler scheduler); @Activator void enterActiveMode(); @PreDestroy void shutdown();  }
@Test public void testOnlyFenzoPartitionIsIncluded() { AgentInstanceGroup fenzoInstanceGroup = createInstanceGroup(0).toBuilder().withId("fenzo").build(); AgentInstanceGroup kubeSchedulerInstanceGroup = createInstanceGroup(0).toBuilder().withId("kubeScheduler").build(); when(agentManagementService.getInstanceGroups()).thenReturn(asList(fenzoInstanceGroup, kubeSchedulerInstanceGroup)); when(agentManagementService.isOwnedByFenzo(any(AgentInstanceGroup.class))).thenAnswer(invocation -> { AgentInstanceGroup instanceGroup = invocation.getArgument(0); return instanceGroup.getId().equals("fenzo"); }); AgentInstance fenzoInstance = createRemovableInstance("fenzo", "agentInstance1"); AgentInstance kubeSchedulerInstance = createRemovableInstance("kubeScheduler", "agentInstance2"); when(agentManagementService.getAgentInstances("fenzo")).thenReturn(Collections.singletonList(fenzoInstance)); when(agentManagementService.getAgentInstances("kubeScheduler")).thenReturn(Collections.singletonList(kubeSchedulerInstance)); when(agentManagementService.terminateAgents("fenzo", singletonList("agentInstance1"), true)) .thenReturn(Observable.just(singletonList(Either.ofValue(true)))); ClusterRemovableAgentRemover clusterRemovableAgentRemover = new ClusterRemovableAgentRemover(titusRuntime, configuration, agentManagementService, v3JobOperations, testScheduler); clusterRemovableAgentRemover.doRemoveAgents().await(); verify(agentManagementService).terminateAgents("fenzo", singletonList("agentInstance1"), true); verify(agentManagementService, times(0)).terminateAgents("kubeScheduler", singletonList("agentInstance2"), true); }
@Test public void testClusterAgentRemoval() { AgentInstanceGroup instanceGroup = createInstanceGroup(0); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); AgentInstance agentInstance1 = createRemovableInstance("instanceGroup1", "agentInstance1"); AgentInstance agentInstance2 = createActiveInstance("instanceGroup1", "agentInstance1"); List<AgentInstance> agentInstances = asList(agentInstance1, agentInstance2); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(agentInstances); when(agentManagementService.terminateAgents("instanceGroup1", singletonList("agentInstance1"), true)) .thenReturn(Observable.just(singletonList(Either.ofValue(true)))); Task task = createTask("agentInstance2"); when(v3JobOperations.getTasks()).thenReturn(singletonList(task)); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterRemovableAgentRemover clusterRemovableAgentRemover = new ClusterRemovableAgentRemover(titusRuntime, configuration, agentManagementService, v3JobOperations, testScheduler); clusterRemovableAgentRemover.doRemoveAgents().await(); verify(agentManagementService).terminateAgents("instanceGroup1", singletonList("agentInstance1"), true); }
@Test public void testDoNotRemoveMoreAgentsThanInstanceGroupMin() { AgentInstanceGroup instanceGroup = createInstanceGroup(2); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); AgentInstance agentInstance1 = createActiveInstance("instanceGroup1", "agentInstance1"); AgentInstance agentInstance2 = createActiveInstance("instanceGroup1", "agentInstance2"); List<AgentInstance> agentInstances = asList(agentInstance1, agentInstance2); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(agentInstances); when(agentManagementService.terminateAgents("instanceGroup1", singletonList("agentInstance1"), true)) .thenReturn(Observable.just(singletonList(Either.ofValue(true)))); when(v3JobOperations.getTasks()).thenReturn(emptyList()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterRemovableInstanceGroupAgentRemover clusterRemovableInstanceGroupAgentRemover = new ClusterRemovableInstanceGroupAgentRemover(titusRuntime, configuration, agentManagementService, v3JobOperations, testScheduler); clusterRemovableInstanceGroupAgentRemover.doRemoveAgents().await(); verify(agentManagementService, times(0)).terminateAgents(any(), any(), anyBoolean()); }