public interface AutomationConstants {

    // These NODE_CONFIG_* constants represent key names in the node configs
    String CONFIG_CREATED_DATE = "createdDate";
    String CONFIG_MAX_SESSION = "maxSession";
    String CONFIG_BROWSER = "createdBrowser";
    String CONFIG_OS = "createdOs";
    // This is the value that will be in the desired capabilities that our node registers with
    // Runtime value of the hub instance id that gets passed in as a system property.  Also used in the node config
    String INSTANCE_ID = "instanceId";
    // IP address of the hub that will be passed in as a system property
    String IP_ADDRESS = "ipAddress";
    // Test run UUID
    String UUID = "uuid";
    String WINDOWS_PROPERTY_NAME="node.windows.json";
    String LINUX_PROPERTY_NAME="node.linux.json";
    String EXTRA_CAPABILITIES_PROPERTY_NAME="extraCapabilities";
    String AWS_ACCESS_KEY="awsAccessKey";
    String AWS_PRIVATE_KEY="awsSecretKey";
    String AWS_DEFAULT_RESOURCE_NAME= "aws.properties.default";
    String REAPER_THREAD_CONFIG = "useReaperThread";
} public class AutomationContext {

    private static final Logger log = LoggerFactory.getLogger(AutomationContext.class);
    private static final int DEFAULT_MAX_THREAD_COUNT = 150;
    private static AutomationRunContext context = new AutomationRunContext();

    // Singleton to maintain a context object

    /**
     * Returns the singleton {@link com.rmn.qa.AutomationRunContext context} instance
     * @return
     */
    public static AutomationRunContext getContext() {
        return AutomationContext.context;
    }

    /**
     * Clears out the previous context.  Used for unit testing
     */
    public static void refreshContext() {
        AutomationContext.context = new AutomationRunContext();
    }

    static {
        String totalNodeCount = System.getProperty("totalNodeCount");
        // Default to 150 if node count was not passed in
        if(totalNodeCount == null) {
            getContext().setTotalNodeCount(AutomationContext.DEFAULT_MAX_THREAD_COUNT);
            log.warn("Defaulting node count to " + AutomationContext.DEFAULT_MAX_THREAD_COUNT);
        } else {
            getContext().setTotalNodeCount(Integer.parseInt(totalNodeCount));
        }
    }

} public final class AutomationDynamicNode implements Comparable<AutomationDynamicNode> {

    /** <pre>
     *  RUNNING    - node is running and no further action needs to be taken
     *  EXPIRED    - node has passed its expiration date and needs to be terminated.  A node will not
     *               be marked expired if there is sufficient load on the system to require the node resources
     *  TERMINATED - node has been successfully terminated through the EC2 API
     * </pre>
     **/
    public enum STATUS { RUNNING,EXPIRED,TERMINATED };

    // This is used to set how far past the created date that the node will
    // be marked for termination
    private static final int NODE_INTERVAL_LIFETIME = 55; // 55 minutes

    // TODO: Refactor this to be AutomationRunRequest
    private final String uuid,instanceId,browser,os;
    private Date startDate,endDate;
    private final int nodeCapacity;
    private STATUS status;

    /**
     * Constructor to create a new node representing instance
     * @param uuid UUID of the test run that created this node
     * @param instanceId Instance ID of the instance this node represents
     * @param browser Requested browser of the test run that created this node
     * @param os Requested OS of the test run that created this node
     * @param startDate Date that this node was started
     * @param nodeCapacity Maximum test capacity that this node can run
     */
    public AutomationDynamicNode(String uuid,String instanceId,String browser, String os,Date startDate,int nodeCapacity){
        this.uuid = uuid;
        this.instanceId = instanceId;
        this.browser = browser;
        this.os = os;
        this.startDate = startDate;
        this.endDate = computeEndDate(startDate);
        this.nodeCapacity = nodeCapacity;
        this.status = STATUS.RUNNING;
    }

    /**
     * Updates the status of this node.
     * @param status
     */
    public void updateStatus(STATUS status) {
        this.status = status;
    }

    /**
     * Computes the end date for this node based on the pre configured
     * end time
     * @param dateStarted
     * @return
     */
    private Date computeEndDate(Date dateStarted) {
        Calendar c = Calendar.getInstance();
        c.setTime(dateStarted);
        c.add(Calendar.MINUTE, AutomationDynamicNode.NODE_INTERVAL_LIFETIME);  // number of days to add
        return c.getTime();
    }

    /**
     * Returns the UUID for this node (will be the UUID of the run that
     * resulted in this node being started)
     * @return
     */
    public String getUuid() {
        return uuid;
    }

    /**
     * Returns the instance id of this node
     * @return
     */
    public String getInstanceId() {
        return instanceId;
    }

    /**
     * Returns the browser of this node.  Will be the browser of the
     * run that resulted in this node being started
     * @return
     */
    public String getBrowser() {
        return browser;
    }

    /**
     * Returns the OS of this node
     * @return
     */
    public String getOs() {
        return os;
    }

    /**
     * Returns the date that his node was started.  This will
     * be the time that the node was requested and not necessarily started
     * @return
     */
    public Date getStartDate() {
        return startDate;
    }

    /**
     * Returns the currently scheduled end date of this node.  Note
     * that this can change as this node end date gets pushed back.
     * @return
     */
    public Date getEndDate() {
        return endDate;
    }

    /**
     * Sets the end date for this node.
     * @param endDate Date which will be set for the end date
     */
    public void setEndDate(Date endDate) {
        this.endDate = endDate;
    }

    /**
     * Increments the end date by an hour.  Useful for moving the end date into the next Amazon billing cycle
     */
    public void incrementEndDateByOneHour() {
        // Add 60 minutes so we're as close to the hour as we can be instead of adding 55 again
        setEndDate(AutomationUtils.modifyDate(getEndDate(),60,Calendar.MINUTE));
    }

    /**
     * Returns the total node capacity of this node (total number of browsers
     * that can run simultaneously)
     * @return
     */
    public int getNodeCapacity() {
        return nodeCapacity;
    }

    /**
     * Returns the current status of this node.
     * @return
     */
    public STATUS getStatus(){
        return status;
    }

    @Override
    public int compareTo(AutomationDynamicNode o) {
        return this.startDate.compareTo(getStartDate());
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }

        AutomationDynamicNode that = (AutomationDynamicNode) o;

        if (!instanceId.equals(that.instanceId)) {
            return false;
        }
        if (!uuid.equals(that.uuid)) {
            return false;
        }

        return true;
    }

    @Override
    public int hashCode() {
        int result = uuid.hashCode();
        result = 31 * result + instanceId.hashCode();
        return result;
    }

    @Override
    public String toString() {
        return String.format("Node - UUID: %s Instance ID: %s Created Date: %s",uuid,instanceId, startDate);
    }
} public class AutomationRequestMatcher implements RequestMatcher {

    private static final Logger log = LoggerFactory.getLogger(AutomationRequestMatcher.class);

    /**
     * {@inheritDoc}
     */
    @Override
    public int getNumFreeThreadsForParameters(ProxySet proxySet, AutomationRunRequest runRequest) {
        // This will keep a count of the number of instances that can run our requested test
        int totalFreeSlots = 0;
        // Current runs registered with the hub.  Make a copy of the set so we don't muck with the original set of registered runs
        Map<String,Integer> currentRuns = new HashMap<>();
        // Add them all into our map so we can keep track of both the runs we need to delete from our free count as well as
        // tests that are in progress for the run.
        for(String key : AutomationContext.getContext().getRunUuids()) {
            currentRuns.put(key,0);
        }
        for(RemoteProxy proxy : proxySet) {
            int matchingCapableSlots = 0;
            int runningSessions = 0;
            int matchingRunningSessions = 0;
            int maxNodeThreadsAvailable = proxy.getMaxNumberOfConcurrentTestSessions();
            Map<String,Object> config = proxy.getConfig();
            String instanceId = null;
            if(config.containsKey(AutomationConstants.INSTANCE_ID)) {
                instanceId = (String)config.get(AutomationConstants.INSTANCE_ID);
            }
            boolean nodeMarkedForTermination = false;
            if(instanceId != null) {
                AutomationDynamicNode node = AutomationContext.getContext().getNode(instanceId);
                // If this node has been spun up and it is no longer in the running state, go to the next test slot
                // as we cannot consider this node to be a free resource
                if(node != null) {// There really shouldn't ever be a null node here but adding the check regardless
                    if(node.getStatus() != AutomationDynamicNode.STATUS.RUNNING) {
                        // If this is a dynamic node and its not in the running state, we should not be calculating its resources as available
                        nodeMarkedForTermination = true;
                    }
                }
            }
            if(instanceId != null) {
                log.info(String.format("Analyzing node %s...",instanceId));
            } else {
                log.info("Analyzing node...");
            }
            for (TestSlot testSlot : proxy.getTestSlots()) {
                //TODO Do selenium flavor of browsers to match here from RMN
                //TODO Better property matching
                TestSession session = testSlot.getSession();
                Map<String,Object> testSlotCapabilities = testSlot.getCapabilities();
                if(session != null) {
                    Map<String,Object> sessionCapabilities = session.getRequestedCapabilities();
                    Object uuid = sessionCapabilities.get(AutomationConstants.UUID);
                    // If the session has a UUID, go ahead and remove it from our runs that we're going to subtract from our available
                    // node count as this means the run is under way
                    if(uuid != null && currentRuns.containsKey(uuid)) {
                        int previousCount = currentRuns.get(uuid);
                        currentRuns.put((String)uuid,previousCount + 1);
                    }
                    if(runRequest.matchesCapabilities(testSlotCapabilities)) {
                        matchingRunningSessions++;
                    }
                    runningSessions++;
                }
                if(runRequest.matchesCapabilities(testSlotCapabilities)) {
                    matchingCapableSlots++;
                }
            }
            log.info(String.format("Node had %d matching running sessions and %d matching capable slots",matchingRunningSessions,matchingCapableSlots));
            int nodeFreeSlots;
            // If the node is marked for termination, we need to subtract matching running sessions from our free count, and make sure to not add
            // any capable slots, as they're really not even 'capable' since the node will be shutdown
            if(nodeMarkedForTermination) {
                log.info(String.format("Node marked for termination.  Subtracting %d sessions from %d total free slots",matchingRunningSessions,totalFreeSlots));
                totalFreeSlots -= matchingRunningSessions;
            } else {
                // Decrement the running sessions only if running + free is more than the total threads the node can handle. This will handle
                // load from a capacity standpoint of a node
                if((runningSessions + matchingCapableSlots) > maxNodeThreadsAvailable ) {
                    if(matchingCapableSlots < maxNodeThreadsAvailable) {
                        log.debug(String.format("Subtracting %d running sessions from %d capable sessions",runningSessions,matchingCapableSlots));
                        nodeFreeSlots = matchingCapableSlots - runningSessions;
                    } else {
                        log.debug(String.format("Subtracting %d running sessions from %d maximum node thread limit",matchingCapableSlots,maxNodeThreadsAvailable));
                        nodeFreeSlots = maxNodeThreadsAvailable - runningSessions;
                    }
                } else {
                    log.debug(String.format("%d free node slots derived from matching capable slots",matchingCapableSlots));
                    nodeFreeSlots = matchingCapableSlots;
                    // If there were any running sessions that match this browser, we need to subtract them from the capable sessions
                    if(matchingRunningSessions != 0) {
                        log.debug(String.format("%d matching running sessions will be subtracted from %d node free slots",matchingRunningSessions,nodeFreeSlots));
                        nodeFreeSlots -= matchingRunningSessions;
                    }
                }
                // If nodeFreeSlots is negative, go ahead and subtract the capable sessions instead
                if(nodeFreeSlots < 0) {
                    log.warn("The number of free node slots was less than 0.  Resetting to 0.");
                    nodeFreeSlots = 0;
                }
                totalFreeSlots += nodeFreeSlots;
            }
        }
        // Any runs still in this set means that run has not started yet, so we should consider this in our math
        for(String uuid : currentRuns.keySet()) {
            AutomationRunRequest request = AutomationContext.getContext().getRunRequest(uuid);
            // If we're not dealing with an old run request that just never started, go ahead and decrement
            // the value from available nodes on this hub
            if(request != null && !AutomationUtils.isCurrentTimeAfterDate(request.getCreatedDate(), 90, Calendar.SECOND) ) {
                if(!runRequest.matchesCapabilities(request)) {
                    log.warn(String.format("Requested run %s did not match pending run %s so count will not be included",runRequest,request));
                    continue;
                }
                int currentlyRunningTestsForRun = currentRuns.get(uuid);
                // If some of the tests are underway, subtract the currently running tests from the number of total tests, and then subtract that
                // from the number of free slots.  This way we're including tests that may not have started yet in our free resource check
                if(currentlyRunningTestsForRun < request.getThreadCount()) {
                    int countToSubtract = request.getThreadCount() - currentlyRunningTestsForRun;
                    log.debug(String.format("In progress run has %d threads that will be subtracted from our total free count %d.",countToSubtract,totalFreeSlots));
                    totalFreeSlots -= countToSubtract;
                }
            }
        }
        // Make sure we don't return a negative number to the caller
        if(totalFreeSlots < 0) {
            log.info("The number of total free node slots was less than 0.  Resetting to 0.");
            totalFreeSlots = 0;
        }
        log.info(String.format("Returning %s free slots for request %s",totalFreeSlots,runRequest));
        return totalFreeSlots;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public int getNumInProgressTests(ProxySet proxySet, AutomationRunRequest runRequest) {
        int inProgressTests = 0;
        for(RemoteProxy proxy : proxySet) {
            for(TestSlot testSlot : proxy.getTestSlots() ) {
                TestSession session = testSlot.getSession();
                if(session != null) {
                    if(runRequest.matchesCapabilities(session.getRequestedCapabilities())) {
                        inProgressTests++;
                    }
                }
            }
        }
        return inProgressTests;
    }

} public final class AutomationRunRequest {

    private final String uuid;
    private final Integer threadCount;
    private final String browser;
    private final String browserVersion;
    private final String os;
    private final Date createdDate;

    // Require callers to have required variables through constructor below
    private AutomationRunRequest() {
        this(null);
    }

    /**
     * Constructs a run request instance for the specified browser
     * @param browser Browser for the requesting test run
     */
    public AutomationRunRequest(String browser) {
        this(null,null,browser);
    }

    /**
     * Constructs a run request object
     * @param uuid UUID to represent the requesting test run
     * @param threadCount Number of threads for the requesting test run
     * @param browser Browser for the requesting test run
     */
    public AutomationRunRequest(String uuid, Integer threadCount,String browser) {
        this(uuid,threadCount,browser,null,null);
    }

    /**
     * Constructs a run request object
     * @param runUuid UUID to represent the requesting test run
     * @param threadCount Number of threads for the requesting test run
     * @param browser Browser for the requesting test run
     * @param browserVersion Browser version for the requesting test run
     * @param os OS for the requesting test run
     */
    public AutomationRunRequest(String runUuid, Integer threadCount, String browser, String browserVersion, String os) {
        this(runUuid, threadCount, browser, browserVersion, os, new Date());
    }

    /**
     * Constructs a run request object
     * @param runUuid UUID to represent the requesting test run
     * @param threadCount Number of threads for the requesting test run
     * @param browser Browser for the requesting test run
     * @param browserVersion Browser version for the requesting test run
     * @param os OS for the requesting test run
     * @param createdDate Date that the test run request was received
     */
    @VisibleForTesting
    public AutomationRunRequest(String runUuid, Integer threadCount, String browser,String browserVersion, String os, Date createdDate) {
        this.uuid = runUuid;
        this.threadCount = threadCount;
        this.browser = browser;
        this.browserVersion = browserVersion;
        this.os = os;
        this.createdDate = createdDate;
    }

    /**
     * Generates a AutomationRunRequest object from the capabilities passed in
     * @param capabilities
     * @return
     */
    public static AutomationRunRequest requestFromCapabilities(Map<String,Object> capabilities) {
        String capabilityBrowser = (String)capabilities.get(CapabilityType.BROWSER_NAME);
        String capabilityBrowserVersion = null;
        if(capabilities.containsKey(CapabilityType.VERSION)) {
            capabilityBrowserVersion = (String)capabilities.get(CapabilityType.VERSION);
        }
        String capabilityOs = (String)capabilities.get(CapabilityType.PLATFORM);
        return new AutomationRunRequest(null,null,capabilityBrowser,capabilityBrowserVersion,capabilityOs);
    }

    /**
     * Returns the UUID for this run request
     * @return
     */
    public String getUuid() {
        return uuid;
    }

    /**
     * Returns the thread count requested by this run
     * @return
     */
    public int getThreadCount() {
        return threadCount;
    }

    /**
     * Returns the browser (e.g. 'chrome', 'firefox', etc) for this run request
     * @return
     */
    public String getBrowser() {
        return browser;
    }

    /**
     * Returns the version of the browser (e.g. '27' for Firefox)
     * @return
     */
    public String getBrowserVersion() { return browserVersion; }

    /**
     * Returns the OS (e.g. 'linux')
     * @return
     */
    public String getOs() { return os; }
    /**
     * Returns the created date for this run request
     * @return
     */
    public Date getCreatedDate() {
        return createdDate;
    }

    /**
     * Returns true if this run request is less than 2 minutes old, false otherwise
     * @return
     */
    public boolean isNewRun() {
        Calendar c = Calendar.getInstance();
        c.setTime(createdDate);
        c.add(Calendar.MINUTE,2);
        return new Date().before(c.getTime());
    }

    @Override
    public String toString() {
        StringBuilder builder = new StringBuilder();
        builder.append("Run Request");
        if(!StringUtils.isEmpty(uuid)) {
            builder.append(" - UUID: ").append(uuid);
        }
        if(threadCount != null) {
            builder.append(" - Thread count: ").append(threadCount);
        }
        if(!StringUtils.isEmpty(browser)) {
            builder.append(" - Browser: ").append(browser);
        }
        if(!StringUtils.isEmpty(os)) {
            builder.append(" - OS: ").append(os);
        }
        return builder.toString();
    }

    /**
     * Returns true if this run request matches the capabilities passed in.  Includes browser, browser version, and OS
     * @param capabilities
     * @return
     */
    public boolean matchesCapabilities(Map<String,Object> capabilities) {
        String capabilityBrowser = (String)capabilities.get(CapabilityType.BROWSER_NAME);
        String capabilityBrowserVersion = (String)capabilities.get(CapabilityType.VERSION);
        String capabilityOs = (String)capabilities.get(CapabilityType.PLATFORM);
        if(!AutomationUtils.lowerCaseMatch(browser, capabilityBrowser)) {
            return false;
        }
        if(browserVersion != null && !AutomationUtils.lowerCaseMatch(browserVersion,capabilityBrowserVersion))  {
            return false;
        }
        if(os != null && !AutomationUtils.lowerCaseMatch(os, capabilityOs)) {
            // If either OS has 'ANY' for the platform, that means it should be a match regardless and we don't have to count this as a non-match
            if(!AutomationUtils.lowerCaseMatch(os,"any") && !AutomationUtils.lowerCaseMatch(capabilityOs,"any")) {
                return false;
            }
        }
        return true;
    }

    /**
     * Returns true if this run request matches the run request passed in.  Includes browser, browser version, and OS
     * @param otherRequest
     * @return
     */
    public boolean matchesCapabilities(AutomationRunRequest otherRequest) {
        if(!AutomationUtils.lowerCaseMatch(browser, otherRequest.getBrowser())) {
            return false;
        }
        if(browserVersion != null && browserVersion != otherRequest.getBrowserVersion()) {
            return false;
        }
        if(os != null && !AutomationUtils.lowerCaseMatch(os, otherRequest.getOs())) {
            // If either OS has 'ANY' for the platform, that means it should be a match regardless and we don't have to count this as a non-match
            if(!AutomationUtils.lowerCaseMatch(os,"any") && !AutomationUtils.lowerCaseMatch(otherRequest.getOs(),"any")) {
                return false;
            }
        }
        return true;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;

        AutomationRunRequest that = (AutomationRunRequest) o;

        if (!browser.equals(that.browser)) return false;
        if (browserVersion != null ? !browserVersion.equals(that.browserVersion) : that.browserVersion != null)
            return false;
        if (os != null ? !os.equals(that.os) : that.os != null) return false;

        return true;
    }

    @Override
    public int hashCode() {
        int result = browser.hashCode();
        result = 31 * result + (browserVersion != null ? browserVersion.hashCode() : 0);
        result = 31 * result + (os != null ? os.hashCode() : 0);
        return result;
    }
} public final class AutomationUtils {

    /**
     * Modifies the specified date
     * @param dateToModify Date to modify
     * @param unitsToModify Number of units to modify (e.g. 6 for 6 seconds)
     * @param unitType Measurement type (e.g. Calendar.SECONDS)
     * @return Modified date
     */
    public static Date modifyDate(Date dateToModify,int unitsToModify,int unitType) {
        Calendar c = Calendar.getInstance();
        c.setTime(dateToModify);
        // Add 60 seconds so we're as close to the hour as we can be instead of adding 55 again
        c.add(unitType,unitsToModify);
        return c.getTime();
    }

    /**
     * Returns true if the current time is after the specified date, false otherwise
     * @param dateToCheck Date to check against the current time
     * @param unitsToCheckWith Number of units to add/subtract from dateToCheck
     * @param unitType Unit type (e.g. Calendar.MINUTES)
     * @return
     */
    public static boolean isCurrentTimeAfterDate(Date dateToCheck,int unitsToCheckWith,int unitType) {
        Date targetDate = AutomationUtils.modifyDate(dateToCheck,unitsToCheckWith,unitType);
        return new Date().after(targetDate);
    }

    /**
     * Returns true if the strings are lower case equal
     * @param string1 First string to compare
     * @param string2 Second string to compare
     * @return
     */
    public static boolean lowerCaseMatch(String string1, String string2) {
        string2 = string2.toLowerCase().replace(" ","");
        return string2.equals(string1.toLowerCase().replace(" ", ""));
    }

} public interface RequestMatcher {

    /**
     * Returns the number of free threads which satisfy the conditions of the parameters
     * @param proxySet Set of current registered proxy objects
     * @param runRequest Request used to match against
     * @return
     */
    int getNumFreeThreadsForParameters(ProxySet proxySet, AutomationRunRequest runRequest);

    /**
     * Returns the number of in progress tests which match the request passed in
     * @param proxySet Set of current registered proxy objects
     * @param runRequest Request used to match against
     * @return
     */
    int getNumInProgressTests(ProxySet proxySet, AutomationRunRequest runRequest);

} public interface RegistryRetriever {

    /**
     * Retrieves a Selenium Registry
     * @return
     */
    Registry retrieveRegistry();

} public class NodesCouldNotBeStartedException extends Exception {
    public NodesCouldNotBeStartedException(String msg) {
        super(msg);
    }

    public NodesCouldNotBeStartedException(String msg, Throwable t) {
        super(msg,t);
    }
} public class AwsVmManager implements VmManager {

    private static final Logger log = LoggerFactory.getLogger(AwsVmManager.class);
    public static final DateFormat NODE_DATE_FORMAT = new SimpleDateFormat("MM/dd/yyyy HH:mm:ss");

    private AmazonEC2Client client;
    BasicAWSCredentials credentials;
    private Properties awsProperties;
    public static final int CHROME_THREAD_COUNT = 6;
    public static final int FIREFOX_IE_THREAD_COUNT = 1;

    private String region;

    static {

        // Read and write dates from node config in UTC format
        NODE_DATE_FORMAT.setTimeZone(TimeZone.getTimeZone("UTC"));
    }

    /**
     * Creates a new AwsVMManager instance using
     * the reagion from the properties file
     */
    public AwsVmManager() {
        awsProperties = initAWSProperties();
        this.region = awsProperties.getProperty("region");
        /**
         * By default we use the credentials provided in the configuration files.
         * If there are none we fall back to IAM roles.
         */
        try {
            credentials = getCredentials();
            client = new AmazonEC2Client(credentials);
        } catch (IllegalArgumentException e) {
            log.info("Falling back to IAM roles for authorization since no credentials provided in system properties", e);
            client = new AmazonEC2Client();
        }
        client.setEndpoint(awsProperties.getProperty(region + "_endpoint"));
    }

    /**
     * Creates a new AwsVmManager instance.
     *
     * @param  client {@link com.amazonaws.services.ec2.AmazonEC2Client Client} to use for AWS interaction
     * @param  properties {@link java.util.Properties Properties} to use for EC2 config loading
     * @param  region      Region inside of AWS to use
     */
    public AwsVmManager(final AmazonEC2Client client, final Properties properties, final String region) {
        this.client = client;
        this.awsProperties = properties;
        this.region = region;
    }

    /**
     * Initializes the AWS properties from the default properties file. Allows the user to override the default
     * properties if desired
     *
     * @return
     */
    Properties initAWSProperties() {
        Properties properties = new Properties();
        String propertiesLocation = System.getProperty("propertyFileLocation");

        // If the user passed in an AWS config file, go ahead and use it instead of the default one
        if (propertiesLocation != null) {
            File f = new File(propertiesLocation);
            try {
                InputStream is = new FileInputStream(f);
                properties.load(is);
            } catch (IOException e) {
                throw new RuntimeException("Could not load custom aws.properties", e);
            }
        } else {
            InputStream stream = AwsVmManager.class.getClassLoader().getResourceAsStream(
                    AutomationConstants.AWS_DEFAULT_RESOURCE_NAME);
            try {
                properties.load(stream);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        return properties;
    }

    /**
     * Retrieves AWS {@link com.amazonaws.auth.BasicAWSCredentials credentials} from the configuration file.
     *
     * @return
     */
    @VisibleForTesting
    BasicAWSCredentials getCredentials() {

        // Give the system property credentials precedence over ones found in the config file
        String accessKey = System.getProperty(AutomationConstants.AWS_ACCESS_KEY);
        if (accessKey == null) {
            accessKey = awsProperties.getProperty(AutomationConstants.AWS_ACCESS_KEY);
            if (accessKey == null) {
                throw new IllegalArgumentException(String.format(
                        "AWS Access Key must be passed in by the [%s] system property or be present in the AWS config file",
                        AutomationConstants.AWS_ACCESS_KEY));
            }
        }

        String privateKey = System.getProperty(AutomationConstants.AWS_PRIVATE_KEY);
        if (privateKey == null) {
            privateKey = awsProperties.getProperty(AutomationConstants.AWS_PRIVATE_KEY);
            if (privateKey == null) {
                throw new IllegalArgumentException(String.format(
                        "AWS Private Key must be passed in by the [%s] system property or be present in the AWS config file",
                        AutomationConstants.AWS_PRIVATE_KEY));
            }
        }

        return new BasicAWSCredentials(accessKey, privateKey);
    }

    public List<Instance> launchNodes(final String amiId, final String instanceType, final int numberToStart,
            final String userData, final boolean terminateOnShutdown) throws NodesCouldNotBeStartedException {
        RunInstancesRequest runRequest = new RunInstancesRequest();
        runRequest.withImageId(amiId).withInstanceType(instanceType).withMinCount(numberToStart)
                  .withMaxCount(numberToStart).withUserData(userData);
        if (terminateOnShutdown) {
            runRequest.withInstanceInitiatedShutdownBehavior("terminate");
        }

        log.info("Setting image id: " + runRequest.getImageId());
        log.info("Setting instance type: " + runRequest.getInstanceType());

        String subnetKey = awsProperties.getProperty(region + "_subnet_id");
        if (subnetKey != null) {
            log.info("Setting subnet: " + subnetKey);
            runRequest.withSubnetId(subnetKey);
        }

        String securityGroupKey = awsProperties.getProperty(region + "_security_group");
        if (securityGroupKey != null) {

            String[] splitSecurityGroupdIds = securityGroupKey.split(",");

            List securityGroupIdsAryLst = new ArrayList();
            for (int i = 0; i < splitSecurityGroupdIds.length; i++) {

                log.info("Setting security group(s): " + splitSecurityGroupdIds[i]);
                securityGroupIdsAryLst.add(splitSecurityGroupdIds[i]);
            }
            runRequest.setSecurityGroupIds(securityGroupIdsAryLst);
        }

        String keyName = awsProperties.getProperty(region + "_key_name");
        if (keyName != null) {
            log.info("Setting keyname:" + keyName);
            runRequest.withKeyName(keyName);
        }

        log.info("Sending run request to AWS...");

        RunInstancesResult runInstancesResult = getResults(runRequest, 0);
        log.info("Run request result returned.  Adding tags");

        // Tag the instances with the standard RMN AWS data
        List<Instance> instances = runInstancesResult.getReservation().getInstances();
        if (instances.size() == 0) {
            throw new NodesCouldNotBeStartedException(String.format(
                    "Error starting up nodes -- count was zero and did not match expected count of %d", numberToStart));
        }

        associateTags(new Date().toString(), instances);
        return instances;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public List<Instance> launchNodes(final String uuid, String os, final String browser, final String hubHostName,
            final int nodeCount, final int maxSessions) throws NodesCouldNotBeStartedException {

        // Unspecified OS will default to Linux
        if (null == os) {
            if (AutomationUtils.lowerCaseMatch(browser, "internet explorer")) {
                os = "windows";
            } else {
                os = "linux";
            }
        }

        String userData = getUserData(uuid, hubHostName, browser, os, maxSessions);
        String amiId = awsProperties.getProperty(getAmiIdForOs(os, browser));
        String instanceType = awsProperties.getProperty("node_instance_type_" + browser);
        return this.launchNodes(amiId, instanceType, nodeCount, userData, false);
    }

    /**
     * Attempts to run the {@link com.amazonaws.services.ec2.model.RunInstancesRequest RunInstancesRequest}, falling
     * back on alternative subnets if capacity is full in the current region.
     *
     * @param   request
     * @param   requestNumber
     *
     * @return
     *
     * @throws  NodesCouldNotBeStartedException
     */
    private RunInstancesResult getResults(final RunInstancesRequest request, int requestNumber)
        throws NodesCouldNotBeStartedException {
        RunInstancesResult runInstancesResult;
        try {
            if(client == null){
                throw new RuntimeException("The client is not initialized");
            }
            runInstancesResult = client.runInstances(request);
        } catch (AmazonServiceException e) {

            // If there is insufficient capacity in this subnet / availability zone, then we want to try other
            // configured subnets
            if ("InsufficientInstanceCapacity".equals(e.getErrorCode())
                    || "VolumeTypeNotAvailableInZone".equals(e.getErrorCode())) {
                log.error(String.format("Insufficient capacity in subnet [%s]: %s", request.getSubnetId(), e));
                requestNumber = requestNumber + 1;

                String fallBackSubnetId = awsProperties.getProperty(region + "_subnet_fallback_id_" + requestNumber);

                // Make sure and only try to recursively loop so as long as we have a valid fallback subnet id.  Logic
                // to also
                // prevent an accidental infinite loop
                if (fallBackSubnetId != null && requestNumber < 5) {
                    log.info("Setting fallback subnet: " + fallBackSubnetId);

                    // Modify the original request with the new subnet ID we're trying to fallback on
                    request.withSubnetId(fallBackSubnetId);
                } else {
                    throw new NodesCouldNotBeStartedException(
                        "Sufficient resources were not available in any of the availability zones");
                }

                return getResults(request, requestNumber);
            } else {

                // We got an error other than insufficient capacity, and should just throw it for the caller to handle
                throw e;
            }
        }

        return runInstancesResult;
    }

    /**
     * Assigns the tags asynchronously to AWS.
     *
     * @param  threadName
     * @param  instances
     */
    @VisibleForTesting
    void associateTags(final String threadName, final Collection<Instance> instances) {
        Thread reportThread = new AwsTagReporter(threadName, client, instances, awsProperties);
        reportThread.start();
    }

    /**
     * Gets the instance ID based on the OS that is chosen.
     *
     * @param   os       OS for the requested test run
     * @param   browser  Browser for the requested test run
     *
     * @return
     */
    private String getAmiIdForOs(final String os, final String browser) {
        String requestedProperty;
        if (os.equals("windows") || browser.equals(BrowserType.IE)) {
            requestedProperty = region + "_windows_node_ami";
        } else if (os.equals("linux")) {
            requestedProperty = region + "_linux_node_ami";
        } else {
            throw new RuntimeException("Unsupported OS: " + os);
        }

        return requestedProperty;
    }

    /**
     * Terminates the specified instance.
     *
     * @param  instanceId  Id of the instance to terminate
     */
    public boolean terminateInstance(final String instanceId) {
        TerminateInstancesRequest terminateRequest = new TerminateInstancesRequest();
        terminateRequest.withInstanceIds(instanceId);

        if(client == null){
            throw new RuntimeException("The client is not initialized");
        }
        TerminateInstancesResult result = client.terminateInstances(terminateRequest);
        List<InstanceStateChange> stateChanges = result.getTerminatingInstances();
        boolean terminatedInstance = false;
        for (InstanceStateChange stateChange : stateChanges) {
            if (instanceId.equals(stateChange.getInstanceId())) {
                terminatedInstance = true;

                InstanceState currentState = stateChange.getCurrentState();
                if (currentState.getCode() != 32 && currentState.getCode() != 48) {
                    log.error(String.format(
                            "Machine state for id %s should be terminated (48) or shutting down (32) but was %s instead",
                            instanceId, currentState.getCode()));
                    return false;
                }
            }
        }

        if (!terminatedInstance) {
            log.error("Matching terminated instance was not found for instance " + instanceId);
            return false;
        }

        return true;
    }

    @Override
    public List<Reservation> describeInstances(final DescribeInstancesRequest describeInstancesRequest) {
        return client.describeInstances(describeInstancesRequest).getReservations();
    }

    /**
     * Returns a zip file containing the necessary user data for the images we're going to spin up.
     *
     * @param   uuid         UUID of the test run
     * @param   hubHostName  Resolvable host name of the hub the node will register with
     * @param   browser      Browser for the requested test run
     * @param   os           OS for the requested test run
     * @param   maxSessions  Maximum simultaneous test sessions
     *
     * @return
     */
    @VisibleForTesting
    String getUserData(final String uuid, final String hubHostName, final String browser, final String os,
            final int maxSessions) {
        try(ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
                ZipOutputStream zos = new ZipOutputStream(outputStream);
        ) {

            // Pull the node config out so we can write it to the zip file
            ZipEntry nodeConfigZipEntry = new ZipEntry("nodeConfigTemplate.json");
            zos.putNextEntry(nodeConfigZipEntry);

            String nodeConfigContents = getNodeConfig(uuid, hubHostName, browser, os, maxSessions);
            zos.write(nodeConfigContents.getBytes());
            zos.closeEntry();

            // Now pull out the S3 config file so we can include it in the zip
            // that will be sent in as userdata
            ZipEntry s3ZipEntry = new ZipEntry(".s3cfg");
            zos.putNextEntry(s3ZipEntry);

            String s3Contents = getS3Config();
            zos.write(s3Contents.getBytes());
            zos.closeEntry();

            // Make sure and close the zip before encoding it
            zos.close();
            return new String(Base64.encodeBase64(outputStream.toByteArray()));

        } catch (IOException ex) {
            throw new RuntimeException("Error getting user data", ex);
        }
    }

    /**
     * Reads the hub.json file and returns its contents as a Base64-encoded string.
     *
     * @return
     */
    @VisibleForTesting
    String getNodeConfig(final String uuid, final String hostName, final String browser, final String os,
            final int maxSessions) {
        String resourceName;
        if (os.equals("windows")) {
            resourceName = AutomationConstants.WINDOWS_PROPERTY_NAME;
        } else if (os.equals("linux")) {
            resourceName = AutomationConstants.LINUX_PROPERTY_NAME;
        } else {
            throw new RuntimeException("Unexpected OS for prop config: " + os);
        }

        String nodeConfig = getFileContents(resourceName);
        nodeConfig = nodeConfig.replaceAll("<MAX_SESSION>", String.valueOf(maxSessions));
        nodeConfig = nodeConfig.replaceAll("<MAX_SESSION_FIREFOX>",
                String.valueOf(AwsVmManager.FIREFOX_IE_THREAD_COUNT));
        nodeConfig = nodeConfig.replaceAll("<MAX_SESSION_IE>", String.valueOf(AwsVmManager.FIREFOX_IE_THREAD_COUNT));
        nodeConfig = nodeConfig.replaceAll("<MAX_SESSION_CHROME>", String.valueOf(AwsVmManager.CHROME_THREAD_COUNT));
        nodeConfig = nodeConfig.replaceAll("<UUID>", uuid);
        nodeConfig = nodeConfig.replaceAll("<CREATED_BROWSER>", browser);
        nodeConfig = nodeConfig.replaceAll("<CREATED_OS>", os);

        Date createdDate = Calendar.getInstance().getTime();

        // Pass in the created date so we can know when this node was spun up
        nodeConfig = nodeConfig.replaceAll("<CREATED_DATE>", AwsVmManager.NODE_DATE_FORMAT.format(createdDate));
        nodeConfig = nodeConfig.replaceFirst("<HOST_NAME>", hostName);
        return nodeConfig;
    }

    /**
     * Returns the S3 config file replaced with the appropriate AWS key/secret.
     *
     * @return
     */
    @VisibleForTesting
    String getS3Config() {
        String s3Config = getFileContents(".s3cfg");
        if (credentials != null) {
            s3Config = s3Config.replaceAll("<ACCESS_KEY>", credentials.getAWSAccessKeyId());
            s3Config = s3Config.replaceAll("<SECRET_KEY>", credentials.getAWSSecretKey());
        }
        return s3Config;
    }

    /**
     * Returns the contents of the specified resource as a string.
     *
     * @param   resourceName
     *
     * @return
     */
    private static String getFileContents(final String resourceName) {
        String fileContents = "";
        try {
            InputStream stream = AwsVmManager.class.getClassLoader().getResourceAsStream(resourceName);
            BufferedReader br = new BufferedReader(new InputStreamReader(stream));
            String strLine;
            while ((strLine = br.readLine()) != null) {
                fileContents += strLine + "\n";
            }

            stream.close();
        } catch (IOException e) {
            throw new RuntimeException("Error loading resource: -" + resourceName, e);
        }

        return fileContents;
    }
} public interface VmManager {

    /**
     * Launches the specified instances
     * @param uuid UUID of the requesting test run
     * @param os OS of the requesting test run
     * @param browser Browser of the requesting test run
     * @param hubHostName Hub host name for the nodes to register with
     * @param nodeCount Number of nodes to be started
     * @param maxSessions Number of max sessions per node
     * @return
     */
    // TODO Refactor into AutomationRunRequest
    List<Instance> launchNodes(String uuid, String os, String browser, String hubHostName, int nodeCount, int maxSessions) throws NodesCouldNotBeStartedException;

    /**
     * Terminates the specified instance
     * @param instanceId
     */
    // TODO Rename to be node or instance in the name
    boolean terminateInstance(String instanceId);

    /**
     * Returns a list of reservations as defined in the {@link com.amazonaws.services.ec2.model.DescribeInstancesRequest DescribeInstancesRequest}
     * @param describeInstancesRequest
     * @return
     */
    List<Reservation> describeInstances(DescribeInstancesRequest describeInstancesRequest);
} public class AutomationHubCleanupTask extends AbstractAutomationCleanupTask {

    private static final Logger log = LoggerFactory.getLogger(AutomationHubCleanupTask.class);

    private VmManager ec2;
    private final String instanceId;
    @VisibleForTesting
    static final String NAME = "Hub Cleanup Task";
    protected static Date createdDate = null;
    protected static Date endDate = null;
    static boolean errorEncountered = false;

    /**
     * Constructs a hub cleanup task with the specified parameters
     * @param registryRetriever Context retrieval mechanism to use
     * @param ec2 EC2 implementation to use for interaction with the hub
     * @param instanceId Instance ID of the hub to cleanup
     */
    public AutomationHubCleanupTask(RegistryRetriever registryRetriever, VmManager ec2,String instanceId) {
        super(registryRetriever);
        this.ec2 = ec2;
        this.instanceId = instanceId;
    }

    /**
     * Returns the ProxySet to be used for cleanup purposes.
     * @return
     */
    @VisibleForTesting
    protected ProxySet getProxySet() {
        return registryRetriever.retrieveRegistry().getAllProxies();
    }

    @Override
    public String getDescription() {
        return AutomationHubCleanupTask.NAME;
    }

    /**
     * Returns the created date which is pulled from the grid configuration
     * @return
     */
    protected Object getCreatedDate() {
        GridHubConfiguration config = registryRetriever.retrieveRegistry().getConfiguration();
        Object createdDate = config.getAllParams().get(AutomationConstants.CONFIG_CREATED_DATE);
        return createdDate;
    }

    // We're going to continuously monitor this hub to see if we can shut it down.  If were approaching the next
    // billing cycle and the hub has no nodes, we will terminate the hub.  Also, if there is an error parsing the created date,
    // we will terminate it at our earliest convenience (no nodes) without regard to the creation time
    @Override
    public void doWork() {
        log.info("Performing cleanup on hub.");
        synchronized (AutomationHubCleanupTask.class) {
            if(createdDate == null) {
                Object createdDate = getCreatedDate();
                try{
                    AutomationHubCleanupTask.createdDate = AwsVmManager.NODE_DATE_FORMAT.parse((String)createdDate);
                } catch(ParseException pe) {
                    String message = "Error parsing created date for hub: " + pe;
                    log.error(message);
                    errorEncountered = true;
                }
                if(!errorEncountered) {
                    // Set the end date to be 55 minutes + creation date
                    AutomationHubCleanupTask.endDate = AutomationUtils.modifyDate(AutomationHubCleanupTask.createdDate, 55, Calendar.MINUTE);
                }
            }
        }
        Date currentTime = new Date();
        if(errorEncountered || currentTime.after(AutomationHubCleanupTask.endDate)) {
            // If we're into the next billing cycle, don't shut the hub down
            if(errorEncountered && getProxySet().isEmpty()) {
                log.info("Error was encountered parsing the created date, so the hub will be shutdown as it is empty.");
                ec2.terminateInstance(instanceId);
            } else if(AutomationUtils.isCurrentTimeAfterDate(AutomationHubCleanupTask.endDate, 6, Calendar.MINUTE)) {
                log.info("Current date: " + new Date());
                log.info("Current end date: " + AutomationHubCleanupTask.endDate);
                log.info(String.format("Hub [%s] ran into the next billing cycle.  Increasing end date.", instanceId));
                AutomationHubCleanupTask.endDate = AutomationUtils.modifyDate(AutomationHubCleanupTask.endDate,60,Calendar.MINUTE);
                return;
            }else if(getProxySet().isEmpty()){
                log.warn("No running nodes found after hub expiration time -- terminating hub: " + instanceId);
                ec2.terminateInstance(instanceId);
            } else {
                // This means tests are still running and we cannot terminate yet
                log.info("Hub could not be shutdown yet.");
                return;
            }
        }

    }
} public class AutomationNodeCleanupTask extends AbstractAutomationCleanupTask {

    private static final Logger log = LoggerFactory.getLogger(AutomationNodeCleanupTask.class);

    private VmManager ec2;
    private RequestMatcher requestMatcher;
    @VisibleForTesting
    static final String NAME = "Node Cleanup Task";

    /**
     * Constructs a cleanup task with the specified options
     * @param registryRetriever Context retrieval mechanism
     * @param ec2 EC2 implementation
     * @param requestMatcher Request matcher implementation
     */
    public AutomationNodeCleanupTask(RegistryRetriever registryRetriever,VmManager ec2,RequestMatcher requestMatcher) {
        super(registryRetriever);
        this.ec2 = ec2;
        this.requestMatcher = requestMatcher;
    }

    @Override
    public String getDescription() {
        return AutomationNodeCleanupTask.NAME;
    }

    /**
     * Returns the ProxySet to be used for cleanup purposes.
     * @return
     */
    @VisibleForTesting
    protected ProxySet getProxySet() {
        return registryRetriever.retrieveRegistry().getAllProxies();
    }

    // We're going to continuously iterate over registered nodes with the hub.  If they're expired, we're going to mark them for removal.
    // If nodes marked for removal are used into the next billing cycle, then we're going to move their end date back again and put them back
    // into the running queue
    @Override
    public void doWork() {
        log.info("Performing cleanup on nodes.");
        AutomationRunContext context = AutomationContext.getContext();
        Map<String,AutomationDynamicNode> nodes = context.getNodes();
        synchronized (nodes) {
            Iterator<String> iterator = nodes.keySet().iterator();
            Date nowDate = new Date();
            while(iterator.hasNext()) {
                String instanceId = iterator.next();
                AutomationDynamicNode node = nodes.get(instanceId);
                AutomationDynamicNode.STATUS nodeStatus = node.getStatus();
                // If the current time is after the scheduled end time for this node and the node is still running, go ahead and queue it to be removed
                if(nodeStatus == AutomationDynamicNode.STATUS.RUNNING && nowDate.after(node.getEndDate())) {
                    if(canNodeShutDown(node)) {
                        log.info(String.format("Updating node %s to 'EXPIRED' status.  Start date [%s] End date [%s]",instanceId,node.getStartDate(),node.getEndDate()));
                        node.updateStatus(AutomationDynamicNode.STATUS.EXPIRED);
                    }
                } else if(nodeStatus == AutomationDynamicNode.STATUS.EXPIRED) {
                    // See if we're in the next billing cycle (create + 55 + 6, which should equal 61 minutes and would safely be in the next billing cycle)
                    if(AutomationUtils.isCurrentTimeAfterDate(node.getEndDate(), 6,Calendar.MINUTE)) {
                        node.incrementEndDateByOneHour();
                        log.info(String.format("Node [%s] was still running after initial allotted time.  Resetting status and increasing end date to %s.",instanceId, node.getEndDate()));
                        node.updateStatus(AutomationDynamicNode.STATUS.RUNNING);
                    } else if(isNodeCurrentlyEmpty(instanceId)) {
                        log.info(String.format("Terminating node %s and updating status to 'TERMINATED'",instanceId));
                        // Delete node
                        ec2.terminateInstance(instanceId);
                        node.updateStatus(AutomationDynamicNode.STATUS.TERMINATED);
                    }
                } else if(nodeStatus == AutomationDynamicNode.STATUS.TERMINATED) {
                    // If the current time is more than 30 minutes after the node end date, we should remove it from being tracked
                    if(System.currentTimeMillis()  > node.getEndDate().getTime() + (30 * 60 * 1000) ) {
                        // Remove it, and this will remove from tracking since we're referencing the collection
                        log.info(String.format("Removing node [%s] from internal tracking set",instanceId));
                        iterator.remove();
                    }
                }
            }
        }
    }

    /**
     * Returns true if the specified node can be safely shut down, false otherwise
     * @param node
     * @return
     */
    // If free slots are greater than OR equal to our node capacity, that means we have enough wiggle room to go ahead and delete this node.
    // We also need to check to make sure there are no registered runs that have not yet started up
    private boolean canNodeShutDown(AutomationDynamicNode node) {
        // If a new run is queued up, we cannot shut down nodes until it starts, so short circuit out of here
        if(AutomationContext.getContext().isNewRunQueuedUp()) {
            log.warn(String.format("Node %s cannot be shutdown yet as a new run is queued up.", node.getInstanceId()));
            return false;
        }
        ProxySet proxySet = getProxySet();
        // Iterate over our set of proxies until we find the proxy which represents the node we're trying to shut down
        Set<AutomationRunRequest> existingSlots = new HashSet<>();
        Map<String,RemoteProxy> proxies = new HashMap<>();
        for(RemoteProxy proxy : proxySet) {
            Object instanceId = proxy.getConfig().get(AutomationConstants.INSTANCE_ID);
            if(node.getInstanceId().equals(instanceId)) {
                // Associate this proxy with the instance id for use later
                proxies.put((String)instanceId,proxy);
                // Once we've found our target proxy, go ahead and compile a list of browsers this node uses
                for(TestSlot testSlot : proxy.getTestSlots()) {
                    AutomationRunRequest requestRepresentation = AutomationRunRequest.requestFromCapabilities(testSlot.getCapabilities());
                    existingSlots.add(requestRepresentation);
                }
            }
        }
        if(existingSlots.size() == 0) {
            log.error("No browsers found for node: " + node.getInstanceId());
        }
        for(AutomationRunRequest request : existingSlots) {
            int freeSlotsForBrowser = requestMatcher.getNumFreeThreadsForParameters(getProxySet(), request);
            if(freeSlotsForBrowser == 0) {
                log.info(String.format("No free slots exist so node will not be shutdown. Node: %s Request: %s Browser: %s",node.getInstanceId(),request,request.getBrowser()));
                return false;
            }
            // Once we find at least one browser that does not have enough capacity to shut this node down
            // we can return false to say we cannot shut this node down.
            RemoteProxy proxy = proxies.get(node.getInstanceId());
            if(proxy == null) {
                log.error(String.format("Proxy was not found for node %s", node.getInstanceId()));
                return false;
            }
            int matchingSlots = 0;
            // Go over the test slots and get the number of browsers this node can run
            for(TestSlot testSlot : proxy.getTestSlots()) {
                // If this test slot matches the browser, increment our count
                if(request.matchesCapabilities(testSlot.getCapabilities())) {
                    matchingSlots ++;
                }
            }
            log.info(String.format("%d matching slots were found for node %s",matchingSlots,node.getInstanceId()));
            // Get the lesser number between the total node capacity and browser specific capacity
            int finalNum = (matchingSlots < node.getNodeCapacity()) ? matchingSlots : node.getNodeCapacity();
            if(freeSlotsForBrowser < finalNum) {
                // If there are no running tests which match the browser, we don't need to honor this browser for shutdown logic
                // as it is not currently needed
                int inProgressTests = requestMatcher.getNumInProgressTests(getProxySet(),request);
                if(inProgressTests != 0) {
                    log.info(String.format("Current load will not allow for node to shutdown right now. Node: %s Request: %s Free Slots: %s Node Slots: %s",node.getInstanceId(),request,freeSlotsForBrowser,finalNum));
                    return false;
                } else {
                    log.info(String.format("Tests are not in progress. Node: %s Request: %s Free Slots: %s Node Slots: %s",node.getInstanceId(),request,freeSlotsForBrowser,finalNum));
                }
            } else {
                log.info(String.format("Load suitable for Request [%s].  Free slots [%s] Node slots [%s]",request,freeSlotsForBrowser,finalNum));
            }
        }
        // If we iterated over every browser for the node and load was not heavy enough, we can safely shut this node down
        return true;
    }

    /**
     * Returns true if the specified node is empty and has no runs on it, and false otherwise
     * @param instanceToFind
     * @return
     */
    public boolean isNodeCurrentlyEmpty(String instanceToFind) {
        ProxySet proxySet = getProxySet();
        for(RemoteProxy proxy : proxySet){
            List<TestSlot> slots = proxy.getTestSlots();
            Object instanceId = proxy.getConfig().get(AutomationConstants.INSTANCE_ID);
            // If the instance id's do not match, this means this is not the node we are looking for
            // and we should continue on to the next one
            if(!instanceToFind.equals(instanceId)) {
                continue;
            }
            // Now that we found the matching node, iterate over all its test slots to see if any sessions are running
            for (TestSlot testSlot : slots) {
                // If we find a running session, this means the node is occupied, so we should return false
                if(testSlot.getSession() != null) {
                    return false;
                }
            }
            // If we reached this point, this means we found our target node AND it had no sessions, meaning the node was empty
            return true;
        }
        // If we didn't find a matching node, we're going to say the nodes is empty so we can terminate it
        log.warn("No matching node was found in the proxy set.  Instance id: " + instanceToFind);
        return true;
    }
} public class AutomationNodeRegistryTask extends AbstractAutomationCleanupTask {

    private static final Logger log = LoggerFactory.getLogger(AutomationNodeRegistryTask.class);
    @VisibleForTesting
    static final String NAME = "Node Registry Task";

    /**
     * Constructs a registry task with the specified context retrieval mechanism
     * @param registryRetriever Represents the retrieval mechanism you wish to use
     */
    public AutomationNodeRegistryTask(RegistryRetriever registryRetriever) {
        super(registryRetriever);
    }

    /**
     * Returns the ProxySet to be used for cleanup purposes.
     * @return
     */
    protected ProxySet getProxySet() {
        return registryRetriever.retrieveRegistry().getAllProxies();
    }

    @Override
    public String getDescription() {
        return AutomationNodeRegistryTask.NAME;
    }

    // We're going to continuously iterate over registered nodes with the hub.  If they're expired, we're going to mark them for removal.
    // If nodes marked for removal are used into the next billing cycle, then we're going to move their end date back again and put them back
    // into the running queue
    @Override
    public void doWork() {
        log.info("Looking for unregistered nodes");
        ProxySet proxySet = getProxySet();
        if(proxySet != null && proxySet.size() > 0) {
            for(RemoteProxy proxy : proxySet) {
                Map<String,Object> config = proxy.getConfig();
                // If the config has an instanceId in it, this means this node was dynamically started and we should
                // track it if we are not already
                if(config.containsKey(AutomationConstants.INSTANCE_ID)) {
                    String instanceId = (String)config.get(AutomationConstants.INSTANCE_ID);
                    AutomationRunContext context = AutomationContext.getContext();
                    // If this node is already in our context, that means we are already tracking this node to terminate
                    if(!context.nodeExists(instanceId)) {
                        Date createdDate = getDate(config);
                        // If we couldn't parse the date out, we are sort of out of luck
                        if(createdDate == null) {
                            break;
                        }
                        proxy.getConfig();
                        String uuid = (String)config.get(AutomationConstants.UUID);
                        int threadCount = (Integer)config.get(AutomationConstants.CONFIG_MAX_SESSION);
                        String browser = (String)config.get(AutomationConstants.CONFIG_BROWSER);
                        String os = (String)config.get(AutomationConstants.CONFIG_OS);
                        AutomationDynamicNode node = new AutomationDynamicNode(uuid,instanceId,browser,os,createdDate,threadCount);
                        log.info("Unregistered dynamic node found: " + node);
                        context.addNode(node);
                    }
                }
            }
        }
    }

    /**
     * Attempts to parse the created date of the node from the capabilities object
     * @param capabilities
     * @return
     */
    private Date getDate(Map<String,Object> capabilities) {
        String stringDate = (String)capabilities.get(AutomationConstants.CONFIG_CREATED_DATE);
        Date returnDate = null;
        try{
            returnDate = AwsVmManager.NODE_DATE_FORMAT.parse(stringDate);
        } catch (ParseException pe) {
            log.error(String.format("Error trying to parse created date [%s]: %s", stringDate, pe));
        }
        return returnDate;
    }
} public class AutomationReaperTask extends AbstractAutomationCleanupTask {

    private static final Logger log = LoggerFactory.getLogger(AutomationReaperTask.class);
    @VisibleForTesting static final String NAME = "VM Reaper Task";

    private VmManager ec2;

    /**
     * Constructs a registry task with the specified context retrieval mechanism
     * @param registryRetriever Represents the retrieval mechanism you wish to use
     */
    public AutomationReaperTask(RegistryRetriever registryRetriever,VmManager ec2) {
        super(registryRetriever);
        this.ec2 = ec2;
    }
    @Override
    public void doWork() {
        log.info("Running " + AutomationReaperTask.NAME);
        DescribeInstancesRequest describeInstancesRequest = new DescribeInstancesRequest();
        Filter filter = new Filter("tag:LaunchSource");
        filter.withValues("SeleniumGridScalerPlugin");
        describeInstancesRequest.withFilters(filter);
        List<Reservation> reservations = ec2.describeInstances(describeInstancesRequest);
        for(Reservation reservation : reservations) {
            for(Instance instance : reservation.getInstances()) {
                // Look for orphaned nodes
                Date threshold = AutomationUtils.modifyDate(new Date(),-30, Calendar.MINUTE);
                String instanceId = instance.getInstanceId();
                // If we found a node old enough AND we're not internally tracking it, this means this is an orphaned node and we should terminate it
                if(threshold.after(instance.getLaunchTime()) && !AutomationContext.getContext().nodeExists(instanceId)) {
                    log.info("Terminating orphaned node: " + instanceId);
                    ec2.terminateInstance(instanceId);
                }
            }
        }
    }

    @Override
    public String getDescription() {
        return AutomationReaperTask.NAME;
    }
} public class AutomationRunCleanupTask extends AbstractAutomationCleanupTask {

    private static final Logger log = LoggerFactory.getLogger(AutomationRunCleanupTask.class);
    @VisibleForTesting
    static final String NAME = "Run Cleanup Task";

    /**
     * Constructs a run cleanup task with the specified context retrieval mechanism
     * @param registryRetriever
     */
    public AutomationRunCleanupTask(RegistryRetriever registryRetriever) {
        super(registryRetriever);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public String getDescription() {
        return AutomationRunCleanupTask.NAME;
    }

    // Basically we want to continuously monitor all runs that are registered and ensure that at least
    // one test is being run that belongs to that test run.  If there are any orphaned runs, we will
    // go ahead and remove it from AutomationRunContext providing they are old enough
    /**
     * {@inheritDoc}
     */
    @Override
    public void doWork() {
        log.info("Performing cleanup on runs.");
        AutomationRunContext context = AutomationContext.getContext();
        context.cleanUpRunRequests(getProxySet());
    }

    /**
     * Returns the {@link org.openqa.grid.internal.ProxySet ProxySet} to be used for cleanup purposes.
     * @return
     */
    @VisibleForTesting
    protected ProxySet getProxySet() {
        return registryRetriever.retrieveRegistry().getAllProxies();
    }
}
import com.amazonaws.services.ec2.model.Instance;
import com.rmn.qa.AutomationConstants;
import com.rmn.qa.AutomationContext;
import com.rmn.qa.AutomationDynamicNode;
import com.rmn.qa.AutomationRequestMatcher;
import com.rmn.qa.AutomationRunRequest;
import com.rmn.qa.AutomationUtils;
import com.rmn.qa.RequestMatcher;
import com.rmn.qa.RegistryRetriever;
import com.rmn.qa.NodesCouldNotBeStartedException;
import com.rmn.qa.aws.AwsVmManager;
import com.rmn.qa.aws.VmManager;
import com.rmn.qa.task.AutomationHubCleanupTask;
import com.rmn.qa.task.AutomationNodeCleanupTask;
import com.rmn.qa.task.AutomationNodeRegistryTask;
import com.rmn.qa.task.AutomationReaperTask;
import com.rmn.qa.task.AutomationRunCleanupTask;
import org.openqa.grid.internal.ProxySet;
import org.openqa.grid.internal.Registry;
import org.openqa.grid.selenium.GridLauncher;
import org.openqa.grid.web.servlet.RegistryBasedServlet;
import org.openqa.selenium.remote.BrowserType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Date;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
/*
 * Copyright (C) 2014 RetailMeNot, Inc.
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 */
package com.rmn.qa.servlet;



/**
 * Servlet used to register new {@link com.rmn.qa.AutomationRunRequest runs} as well as delete existing {@link com.rmn.qa.AutomationRunRequest runs}.
 * New {@link com.rmn.qa.AutomationRunRequest runs} will automatically spawn up new {@link com.rmn.qa.AutomationDynamicNode nodes} as needed
 * @author mhardin
 */
public class AutomationTestRunServlet extends RegistryBasedServlet implements RegistryRetriever {

    private static final long serialVersionUID = 8484071790930378855L;
    private static final Logger log = LoggerFactory.getLogger(AutomationTestRunServlet.class);
    // Start times
    private static final long START_DELAY_IN_SECONDS = 60L; // 60 start delay for run and node cleanup tasks
    private static final long EXPIRED_POLLING_TIME_IN_SECONDS = 15L; // Look for nodes to clean up every 15 seconds
    private static final long HUB_TERMINATE_START_DELAY_IN_MINUTES= 5L; // Delay 5 minutes to start trying to shutdown the hub
    // Polling times
    private static final long HUB_TERMINATION_POLLING_TIME_IN_MINUTES = 1L; // Look every minute to shutdown the hub
    private static final long NODE_REGISTRATION_POLLING_TIME_IN_MINUTES = 15L; // Look every 15 minutes for new unregistered nodes
    private static final long TEST_RUN_CLEANUP_POLLING_TIME_IN_SECONDS = 60L; // Look for runs to clean up every 60 seconds

    // We override these for unit testing
    private VmManager ec2;
    private RequestMatcher requestMatcher;

    /**
     * Constructs a test run servlet with default values
     */
    public AutomationTestRunServlet() {
        this(null, true, new AwsVmManager(), new AutomationRequestMatcher());
    }

    /**
     * Constructs a test run servlet with customized values
     * @param registry Selenium registry object to use from Grid
     * @param initThreads Set to true if you want the cleanup threads initialized
     * @param ec2 EC2 implementation that you wish to use
     * @param requestMatcher RequestMatcher implementation you wish you use
     */
    public AutomationTestRunServlet(Registry registry, boolean initThreads, VmManager ec2,RequestMatcher requestMatcher) {
        super(registry);
        setManageEc2(ec2);
        setRequestMatcher(requestMatcher);
        // Start up our cleanup thread that will cleanup unused runs
        if(initThreads) {
            this.initCleanupThreads();
        }
    }

    private void initCleanupThreads() {
        // Wrapper to lazily fetch the Registry object as this is not populated at instantiation time
        // Spin up a scheduled thread to poll for unused test runs and clean up them
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(new AutomationRunCleanupTask(this),
                AutomationTestRunServlet.START_DELAY_IN_SECONDS, AutomationTestRunServlet.TEST_RUN_CLEANUP_POLLING_TIME_IN_SECONDS, TimeUnit.SECONDS);
        // Spin up a scheduled thread to clean up and terminate nodes that were spun up
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(new AutomationNodeCleanupTask(this,ec2,requestMatcher),
                AutomationTestRunServlet.START_DELAY_IN_SECONDS,AutomationTestRunServlet.EXPIRED_POLLING_TIME_IN_SECONDS, TimeUnit.SECONDS);
        // Spin up a scheduled thread to register unregistered dynamic nodes (will happen if hub gets shut down)
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(new AutomationNodeRegistryTask(this),
                AutomationTestRunServlet.HUB_TERMINATE_START_DELAY_IN_MINUTES,AutomationTestRunServlet.NODE_REGISTRATION_POLLING_TIME_IN_MINUTES, TimeUnit.MINUTES);
        String instanceId = System.getProperty(AutomationConstants.INSTANCE_ID);
        if(instanceId != null && instanceId.length() > 0) {
            log.info("Instance ID detected.  Hub termination thread will be started.");
            Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(new AutomationHubCleanupTask(this,ec2,instanceId),
                 AutomationTestRunServlet.HUB_TERMINATE_START_DELAY_IN_MINUTES,AutomationTestRunServlet.HUB_TERMINATION_POLLING_TIME_IN_MINUTES, TimeUnit.MINUTES);
        } else {
            log.info("Hub is not a dynamic hub -- termination logic will not be started");
        }
        String runReaperThread = System.getProperty(AutomationConstants.REAPER_THREAD_CONFIG);
        // Reaper thread defaults to on unless specified not to run
        if(!"false".equalsIgnoreCase(runReaperThread)) {
            // Spin up a scheduled thread to terminate orphaned instances
            Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(new AutomationReaperTask(this,ec2),
                    AutomationTestRunServlet.HUB_TERMINATE_START_DELAY_IN_MINUTES,AutomationTestRunServlet.NODE_REGISTRATION_POLLING_TIME_IN_MINUTES, TimeUnit.MINUTES);
        } else {
            log.info("Reaper thread not running due to config flag.");
        }
    }

    void setManageEc2(VmManager ec2) {
        this.ec2 = ec2;
    }

    void setRequestMatcher(RequestMatcher requestMatcher) {
        this.requestMatcher = requestMatcher;
    }

    protected ProxySet getProxySet() {
        return super.getRegistry().getAllProxies();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    protected void doGet(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
        this.doPost(request,response);
    }


    /**
     * Attempts to register a new run request with the server.
     * Returns a 201 if the request can be fulfilled but AMIs must be started
     * Returns a 202 if the request can be fulfilled
     * Returns a 400 if the required parameters are not passed in.
     * Returns a 409 if the server is at full node capacity
     */
    @Override
    protected void doPost(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
        response.setContentType("text/html");
        response.setCharacterEncoding("UTF-8");

        String browserRequested = request.getParameter("browser");
        String browserVersion = request.getParameter("browserVersion");
        String osRequested = request.getParameter("os");
        String threadCount = request.getParameter("threadCount");
        String uuid = request.getParameter(AutomationConstants.UUID);
        // Return a 400 if any of the required parameters are not passed in
        // Check for uuid first as this is the most important variable
        if (uuid == null) {
            String msg = "Parameter 'uuid' must be passed in as a query string parameter";
            log.error(msg);
            response.sendError(HttpServletResponse.SC_BAD_REQUEST, msg);
            return;
        }
        if (browserRequested == null) {
            String msg = "Parameter 'browser' must be passed in as a query string parameter";
            log.error(msg);
            response.sendError(HttpServletResponse.SC_BAD_REQUEST, msg);
            return;
        }
        if (threadCount == null) {
            String msg = "Parameter 'threadCount' must be passed in as a query string parameter";
            log.error(msg);
            response.sendError(HttpServletResponse.SC_BAD_REQUEST, msg);
            return;
        }
        Integer threadCountRequested = Integer.valueOf(threadCount);
        AutomationRunRequest runRequest = new AutomationRunRequest(uuid, threadCountRequested, browserRequested, browserVersion, osRequested);
        log.info(String.format("Server request received.  Browser [%s] - Requested Node Count [%s] - Request UUID [%s]", browserRequested, threadCountRequested, uuid));
        boolean amisNeeded;
        int amiThreadsToStart=0;
        int currentlyAvailableNodes;
        // Synchronize this block until we've added the run to our context for other potential threads to see
        synchronized (AutomationTestRunServlet.class) {
            int remainingNodesAvailable = AutomationContext.getContext().getTotalThreadsAvailable(getProxySet());
            // If the number of nodes this grid hub can actually run is less than the number requested, this hub can not fulfill this run at this time
            if(remainingNodesAvailable < runRequest.getThreadCount()) {
                log.error(String.format("Requested node count of [%d] could not be fulfilled due to hub limit. [%d] nodes available - Request UUID [%s]", threadCountRequested, remainingNodesAvailable, uuid));
                response.sendError(HttpServletResponse.SC_CONFLICT, "Server cannot fulfill request due to configured node limit being reached.");
                return;
            }
            // Get the number of matching, free nodes to determine if we need to start up AMIs or not
            currentlyAvailableNodes = requestMatcher.getNumFreeThreadsForParameters(getProxySet(),runRequest);
            // If the number of available nodes is less than the total number requested, we will have to spin up AMIs in order to fulfill the request
            amisNeeded = currentlyAvailableNodes < threadCountRequested;
            if(amisNeeded) {
                // Get the difference which will be the number of additional nodes we need to spin up to supplement existing nodes
                amiThreadsToStart = threadCountRequested - currentlyAvailableNodes;
            }
            // If the browser requested is not supported by AMIs, we need to not unnecessarily spin up AMIs
            if(amisNeeded && !browserSupportedByAmis(browserRequested)) {
                response.sendError(HttpServletResponse.SC_GONE,"Request cannot be fulfilled and browser is not supported by AMIs");
                return;
            }
            // Add the run to our context so we can track it
            AutomationRunRequest newRunRequest = new AutomationRunRequest(uuid, threadCountRequested, browserRequested, browserVersion, osRequested);
            boolean addSuccessful = AutomationContext.getContext().addRun(newRunRequest);
            if(!addSuccessful) {
                log.warn(String.format("Test run already exists for the same UUID [%s]", uuid));
                response.sendError(HttpServletResponse.SC_BAD_REQUEST, "Test run already exists with the same UUID.");
                return;
            }
        }
        if (amisNeeded) {
            // Start up AMIs as that will be required
            log.warn(String.format("Insufficient nodes to fulfill request. New AMIs will be queued up. Requested [%s] - Available [%s] - Request UUID [%s]", threadCountRequested, currentlyAvailableNodes, uuid));
            try{
                startNodes(uuid,amiThreadsToStart,browserRequested,osRequested);
            } catch(NodesCouldNotBeStartedException e) {