de.fischer.thotti.ec2.clients.EC2ExecutorITHelper.java Source code

Introduction

Here is the source code for de.fischer.thotti.ec2.clients.EC2ExecutorITHelper.java

Source

/*
 * Copyright 2011 Oliver B. Fischer
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package de.fischer.thotti.ec2.clients;

import com.amazonaws.services.ec2.AmazonEC2Client;
import com.amazonaws.services.ec2.model.DescribeInstancesResult;
import com.amazonaws.services.ec2.model.DescribeRegionsResult;
import com.amazonaws.services.ec2.model.Instance;
import com.amazonaws.services.ec2.model.Region;
import com.amazonaws.services.ec2.model.Reservation;
import com.amazonaws.services.ec2.model.TerminateInstancesRequest;
import de.fischer.thotti.awscommon.AWSAccessCredentials;
import de.fischer.thotti.ec2.core.ThottiRuntimeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

public class EC2ExecutorITHelper {
    private Logger logger = LoggerFactory.getLogger(EC2ExecutorITHelper.class);

    private EC2InternalClient.MyAWSCredentials awsCredentials;
    private Map<String, List<String>> existingInstancesMap = new HashMap<String, List<String>>();
    private Map<String, List<String>> currentInstancesMap = new HashMap<String, List<String>>();
    private AmazonEC2Client amazonClient;
    private List<Region> awsRegions = Collections.<Region>emptyList();

    public void retrieveRegions() {
        AmazonEC2Client awsClient = getAmazonClient();

        DescribeRegionsResult regionsResult = awsClient.describeRegions();

        awsRegions = regionsResult.getRegions();
    }

    public void setCredentials(AWSAccessCredentials ec2AccessCredentials) {
        EC2InternalClient.MyAWSCredentials credentials = new EC2InternalClient.MyAWSCredentials();

        credentials.setAccessKeyId(ec2AccessCredentials.getAccessKeyId());
        credentials.setSecretKey(ec2AccessCredentials.getSecretKey());

        awsCredentials = credentials;
    }

    public void retrieveExistingInstances() {
        loadInstanceInformation(existingInstancesMap);
    }

    private void loadInstanceInformation(Map<String, List<String>> instanceMap) {
        AmazonEC2Client awsClient = getAmazonClient();

        // Get all existing instances in all regions
        for (Region region : awsRegions) {
            String regionName = region.getRegionName();

            if (!instanceMap.containsKey(regionName))
                instanceMap.put(regionName, new LinkedList<String>());

            List<String> instanceIDs = instanceMap.get(regionName);

            awsClient.setEndpoint(region.getEndpoint());
            DescribeInstancesResult instancesRes = awsClient.describeInstances();

            List<Reservation> instancesList = instancesRes.getReservations();

            if (instancesList.isEmpty())
                if (logger.isDebugEnabled())
                    logger.debug("No existing instances found in region {}", region.getRegionName());

            for (Reservation reservation : instancesList) {
                for (Instance instance : reservation.getInstances()) {
                    instanceIDs.add(instance.getInstanceId());

                    if (logger.isDebugEnabled()) {
                        logger.debug("Found instance {} in region {}",
                                new Object[] { instance.getInstanceId(), region.getRegionName() });
                    }
                }
            }
        }
    }

    public AmazonEC2Client getAmazonClient() {
        if (null == awsCredentials)
            throw new ThottiRuntimeException("Amazon AWS not set.");

        if (amazonClient == null) {
            amazonClient = new AmazonEC2Client(awsCredentials);
        }

        return amazonClient;
    }

    public void retrieveCurrentInstances() {
        loadInstanceInformation(currentInstancesMap);
    }

    public void terminateNewInstances() {
        for (Region region : awsRegions) {
            String regionName = region.getRegionName();
            List<String> instancesBefore = existingInstancesMap.get(regionName);
            List<String> instancesAfter = currentInstancesMap.get(regionName);

            List<String> addedInstances = new LinkedList(instancesAfter);
            addedInstances.removeAll(instancesBefore);

            terminateInstances(region, addedInstances);
        }
    }

    private void terminateInstances(Region region, List<String> instances) {
        if (instances.isEmpty())
            return;

        AmazonEC2Client client = getAmazonClient();

        client.setEndpoint(region.getEndpoint());

        TerminateInstancesRequest terminateReq = new TerminateInstancesRequest();

        for (String id : instances) {
            if (logger.isDebugEnabled())
                logger.debug("Going to terminate instance {} in region {}.",
                        new Object[] { id, region.getRegionName() });

            terminateReq.withInstanceIds(id);
        }

        client.terminateInstances(terminateReq);
    }
}