package com.intuit.tank.persistence.databases; /* * #%L * Reporting database support * %% * Copyright (C) 2011 - 2015 Intuit Inc. * %% * All rights reserved. This program and the accompanying materials * are made available under the terms of the Eclipse Public License v1.0 * which accompanies this distribution, and is available at * http://www.eclipse.org/legal/epl-v10.html * #L% */ import java.util.ArrayList; import java.util.Date; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; import java.util.UUID; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.Executors; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import javax.annotation.Nonnull; import org.apache.commons.configuration.HierarchicalConfiguration; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.math.NumberUtils; import org.apache.http.HttpStatus; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import com.amazonaws.AmazonServiceException; import com.amazonaws.ClientConfiguration; import com.amazonaws.auth.AWSCredentials; import com.amazonaws.auth.BasicAWSCredentials; import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient; import com.amazonaws.services.dynamodbv2.document.DynamoDB; import com.amazonaws.services.dynamodbv2.document.Table; import com.amazonaws.services.dynamodbv2.model.AttributeDefinition; import com.amazonaws.services.dynamodbv2.model.AttributeValue; import com.amazonaws.services.dynamodbv2.model.BatchWriteItemRequest; import com.amazonaws.services.dynamodbv2.model.BatchWriteItemResult; import com.amazonaws.services.dynamodbv2.model.ComparisonOperator; import com.amazonaws.services.dynamodbv2.model.Condition; import com.amazonaws.services.dynamodbv2.model.ConsumedCapacity; import com.amazonaws.services.dynamodbv2.model.CreateTableRequest; import com.amazonaws.services.dynamodbv2.model.CreateTableResult; import com.amazonaws.services.dynamodbv2.model.DeleteRequest; import com.amazonaws.services.dynamodbv2.model.DeleteTableRequest; import com.amazonaws.services.dynamodbv2.model.DeleteTableResult; import com.amazonaws.services.dynamodbv2.model.DescribeTableRequest; import com.amazonaws.services.dynamodbv2.model.KeySchemaElement; import com.amazonaws.services.dynamodbv2.model.KeyType; import com.amazonaws.services.dynamodbv2.model.ListTablesRequest; import com.amazonaws.services.dynamodbv2.model.ListTablesResult; import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput; import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughputDescription; import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughputExceededException; import com.amazonaws.services.dynamodbv2.model.PutRequest; import com.amazonaws.services.dynamodbv2.model.QueryRequest; import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType; import com.amazonaws.services.dynamodbv2.model.ScanRequest; import com.amazonaws.services.dynamodbv2.model.ScanResult; import com.amazonaws.services.dynamodbv2.model.TableDescription; import com.amazonaws.services.dynamodbv2.model.TableStatus; import com.amazonaws.services.dynamodbv2.model.WriteRequest; import com.intuit.tank.reporting.databases.Attribute; import com.intuit.tank.reporting.databases.IDatabase; import com.intuit.tank.reporting.databases.Item; import com.intuit.tank.reporting.databases.PagedDatabaseResult; import com.intuit.tank.reporting.databases.TankDatabaseType; import com.intuit.tank.results.TankResult; import com.intuit.tank.vm.common.util.MethodTimer; import com.intuit.tank.vm.common.util.ReportUtil; import com.intuit.tank.vm.settings.CloudCredentials; import com.intuit.tank.vm.settings.CloudProvider; import com.intuit.tank.vm.settings.TankConfig; public class AmazonDynamoDatabaseDocApi implements IDatabase { private static final int MAX_NUMBER_OF_RETRIES = 5; private AmazonDynamoDB dynamoDb; private TankConfig config = new TankConfig(); private static final ThreadPoolExecutor EXECUTOR = new ThreadPoolExecutor(10, 50, 60, TimeUnit.SECONDS, new ArrayBlockingQueue<Runnable>(50), Executors.defaultThreadFactory(), new ThreadPoolExecutor.DiscardOldestPolicy()); protected static final int BATCH_SIZE = 25; private static final long MAX_WRITE_UNITS = 1500L; private static Logger logger = LogManager.getLogger(AmazonDynamoDatabaseDocApi.class); /** * * @param dynamoDb */ public AmazonDynamoDatabaseDocApi() { CloudCredentials creds = new TankConfig().getVmManagerConfig().getCloudCredentials(CloudProvider.amazon); if (creds != null && StringUtils.isNotBlank(creds.getKeyId())) { AWSCredentials credentials = new BasicAWSCredentials(creds.getKeyId(), creds.getKey()); this.dynamoDb = new AmazonDynamoDBClient(credentials, new ClientConfiguration()); } else { this.dynamoDb = new AmazonDynamoDBClient(new ClientConfiguration()); } } /** * * @param dynamoDb */ public AmazonDynamoDatabaseDocApi(AmazonDynamoDB dynamoDb) { this.dynamoDb = dynamoDb; } /** * * @{inheritDoc */ @Override public void createTable(String tableName) { try { if (!hasTable(tableName)) { logger.info("Creating table: " + tableName); HierarchicalConfiguration resultsProviderConfig = config.getVmManagerConfig() .getResultsProviderConfig(); long readCapacity = getCapacity(resultsProviderConfig, "read-capacity", 10L); long writeCapacity = getCapacity(resultsProviderConfig, "write-capacity", 50L); ArrayList<AttributeDefinition> attributeDefinitions = new ArrayList<AttributeDefinition>(); attributeDefinitions.add(new AttributeDefinition().withAttributeName( DatabaseKeys.JOB_ID_KEY.getShortKey()).withAttributeType(ScalarAttributeType.S)); attributeDefinitions.add(new AttributeDefinition().withAttributeName( DatabaseKeys.REQUEST_NAME_KEY.getShortKey()).withAttributeType(ScalarAttributeType.S)); ProvisionedThroughput provisionedThroughput = new ProvisionedThroughput().withReadCapacityUnits( readCapacity).withWriteCapacityUnits(writeCapacity); KeySchemaElement hashKeyElement = new KeySchemaElement().withAttributeName( DatabaseKeys.JOB_ID_KEY.getShortKey()).withKeyType(KeyType.HASH); KeySchemaElement rangeKeyElement = new KeySchemaElement().withAttributeName( DatabaseKeys.REQUEST_NAME_KEY.getShortKey()).withKeyType(KeyType.RANGE); CreateTableRequest request = new CreateTableRequest() .withTableName(tableName) .withKeySchema(hashKeyElement, rangeKeyElement) .withAttributeDefinitions(attributeDefinitions) .withProvisionedThroughput(provisionedThroughput); CreateTableResult result = dynamoDb.createTable(request); waitForStatus(tableName, TableStatus.ACTIVE); logger.info("Created table: " + result.getTableDescription().getTableName()); } } catch (Exception t) { logger.error(t, t); throw new RuntimeException(t); } } private long getCapacity(HierarchicalConfiguration resultsProviderConfig, String key, long defaultValue) { long ret = defaultValue; if (resultsProviderConfig != null) { try { String string = resultsProviderConfig.getString(key); if (NumberUtils.isDigits(string)) { return Long.parseLong(string); } } catch (Exception e) { logger.error(e.toString()); } } return ret; } /** * * @{inheritDoc */ @Override public void deleteTable(String tableName) { try { if (hasTable(tableName)) { logger.info("Deleting table: " + tableName); DeleteTableRequest deleteTableRequest = new DeleteTableRequest(tableName); DeleteTableResult result = dynamoDb.deleteTable(deleteTableRequest); logger.info("Deleted table: " + result.getTableDescription().getTableName()); waitForDelete(tableName); } } catch (Exception t) { logger.error(t, t); throw new RuntimeException(t); } } /** * * @{inheritDoc */ @Override public boolean hasTable(String tableName) { String nextTableName = null; do { ListTablesResult listTables = dynamoDb.listTables(new ListTablesRequest() .withExclusiveStartTableName(nextTableName)); for (String name : listTables.getTableNames()) { if (tableName.equalsIgnoreCase(name)) { return true; } } nextTableName = listTables.getLastEvaluatedTableName(); } while (nextTableName != null); return false; } /** * * @{inheritDoc */ @Override public void addTimingResults(final @Nonnull String tableName, final @Nonnull List<TankResult> results, boolean async) { if (!results.isEmpty()) { Runnable task = new Runnable() { public void run() { MethodTimer mt = new MethodTimer(logger, this.getClass(), "addTimingResults (" + results + ")"); List<WriteRequest> requests = new ArrayList<WriteRequest>(); try { for (TankResult result : results) { Map<String, AttributeValue> item = getTimingAttributes(result); PutRequest putRequest = new PutRequest().withItem(item); WriteRequest writeRequest = new WriteRequest().withPutRequest(putRequest); requests.add(writeRequest); } sendBatch(tableName, requests); } catch (Exception t) { logger.error("Error adding results: " + t.getMessage(), t); throw new RuntimeException(t); } mt.endAndLog(); } }; if (async) { EXECUTOR.execute(task); } else { task.run(); } } } /** * * @{inheritDoc */ @Override public Set<String> getTables(String regex) { Set<String> result = new HashSet<String>(); String nextTableName = null; do { ListTablesResult listTables = dynamoDb.listTables(new ListTablesRequest() .withExclusiveStartTableName(nextTableName)); for (String s : listTables.getTableNames()) { if (s.matches(regex)) { result.add(s); } } nextTableName = listTables.getLastEvaluatedTableName(); } while (nextTableName != null); return result; } /** * @{inheritDoc */ @SuppressWarnings("unchecked") @Override public PagedDatabaseResult getPagedItems(String tableName, Object nextToken, String minRange, String maxRange, String instanceId, String jobId) { List<Item> ret = new ArrayList<Item>(); Map<String, AttributeValue> lastKeyEvaluated = (Map<String, AttributeValue>) nextToken; ScanRequest scanRequest = new ScanRequest().withTableName(tableName); Map<String, Condition> conditions = new HashMap<String, Condition>(); if (jobId != null) { Condition jobIdCondition = new Condition(); jobIdCondition.withComparisonOperator(ComparisonOperator.EQ) .withAttributeValueList(new AttributeValue().withS(jobId)); conditions.put(DatabaseKeys.JOB_ID_KEY.getShortKey(), jobIdCondition); } if (StringUtils.isNotBlank(instanceId)) { // add a filter Condition filter = new Condition(); filter.withComparisonOperator(ComparisonOperator.EQ).withAttributeValueList( new AttributeValue().withS(instanceId)); scanRequest.addScanFilterEntry(DatabaseKeys.INSTANCE_ID_KEY.getShortKey(), filter); } Condition rangeKeyCondition = new Condition(); if (minRange != null && maxRange != null) { rangeKeyCondition.withComparisonOperator(ComparisonOperator.BETWEEN.toString()) .withAttributeValueList(new AttributeValue().withS(minRange)) .withAttributeValueList(new AttributeValue().withS(maxRange)); } else if (minRange != null) { rangeKeyCondition.withComparisonOperator(ComparisonOperator.GE.toString()) .withAttributeValueList(new AttributeValue().withS(minRange)); } else if (maxRange != null) { rangeKeyCondition.withComparisonOperator(ComparisonOperator.LT.toString()) .withAttributeValueList(new AttributeValue().withS(maxRange)); } else { rangeKeyCondition = null; } if (rangeKeyCondition != null) { conditions.put(DatabaseKeys.REQUEST_NAME_KEY.getShortKey(), rangeKeyCondition); } scanRequest.withScanFilter(conditions); scanRequest.withExclusiveStartKey(lastKeyEvaluated); ScanResult result = dynamoDb.scan(scanRequest); for (Map<String, AttributeValue> item : result.getItems()) { ret.add(getItemFromResult(item)); } return new PagedDatabaseResult(ret, result.getLastEvaluatedKey()); } /** * * @{inheritDoc */ @Override public List<Item> getItems(String tableName, String minRange, String maxRange, String instanceId, String... jobIds) { List<Item> ret = new ArrayList<Item>(); for (String jobId : jobIds) { Object lastKeyEvaluated = null; do { PagedDatabaseResult pagedItems = getPagedItems(tableName, lastKeyEvaluated, minRange, maxRange, instanceId, jobId); ret.addAll(pagedItems.getItems()); lastKeyEvaluated = pagedItems.getNextToken(); } while (lastKeyEvaluated != null); } return ret; } /** * @{inheritDoc */ @Override public void addItems(final String tableName, List<Item> itemList, final boolean asynch) { if (!itemList.isEmpty()) { final List<Item> items = new ArrayList<Item>(itemList); Runnable task = new Runnable() { public void run() { MethodTimer mt = new MethodTimer(logger, this.getClass(), "addItems (" + items + ")"); List<WriteRequest> requests = new ArrayList<WriteRequest>(); try { for (Item item : items) { Map<String, AttributeValue> toInsert = itemToMap(item); PutRequest putRequest = new PutRequest().withItem(toInsert); WriteRequest writeRequest = new WriteRequest().withPutRequest(putRequest); requests.add(writeRequest); } sendBatch(tableName, requests); } catch (Exception t) { logger.error("Error adding results: " + t.getMessage(), t); throw new RuntimeException(t); } mt.endAndLog(); } }; if (asynch) { EXECUTOR.execute(task); } else { task.run(); } } } /** * @{inheritDoc */ @Override public void deleteForJob(final String tableName, final String jobId, final boolean asynch) { Runnable task = new Runnable() { public void run() { MethodTimer mt = new MethodTimer(logger, this.getClass(), "deleteForJob (" + jobId + ")"); List<Item> items = getItems(tableName, null, null, null, jobId); if (!items.isEmpty()) { List<WriteRequest> requests = new ArrayList<WriteRequest>(); try { for (Item item : items) { String id = null; for (Attribute attr : item.getAttributes()) { if (DatabaseKeys.REQUEST_NAME_KEY.getShortKey().equals(attr.getName())) { id = attr.getValue(); break; } } if (id != null) { Map<String, AttributeValue> keyMap = new HashMap<String, AttributeValue>(); keyMap.put(DatabaseKeys.REQUEST_NAME_KEY.getShortKey(), new AttributeValue().withS(id)); keyMap.put(DatabaseKeys.JOB_ID_KEY.getShortKey(), new AttributeValue().withS(jobId)); DeleteRequest deleteRequest = new DeleteRequest().withKey(keyMap); WriteRequest writeRequest = new WriteRequest().withDeleteRequest(deleteRequest); requests.add(writeRequest); } } sendBatch(tableName, requests); } catch (Exception t) { logger.error("Error adding results: " + t.getMessage(), t); throw new RuntimeException(t); } } mt.endAndLog(); } }; if (asynch) { EXECUTOR.execute(task); } else { task.run(); } } @Override public String getDatabaseName(TankDatabaseType type, String jobId) { return type.name() + "_" + new TankConfig().getInstanceName(); } /** * @{inheritDoc */ @Override public boolean hasJobData(String tableName, String jobId) { if (hasTable(tableName)) { Map<String, Condition> keyConditions = new HashMap<String, Condition>(); Condition jobIdCondition = new Condition().withComparisonOperator(ComparisonOperator.EQ) .withAttributeValueList(new AttributeValue().withS(jobId)); keyConditions.put(DatabaseKeys.JOB_ID_KEY.getShortKey(), jobIdCondition); QueryRequest queryRequest = new QueryRequest().withTableName(tableName).withKeyConditions(keyConditions) .withLimit(1); return dynamoDb.query(queryRequest).getCount() > 0; } return false; } /** * @param item * @return */ private Item getItemFromResult(Map<String, AttributeValue> attributeMap) { List<Attribute> attrs = new ArrayList<Attribute>(); Item ret = new Item(null, attrs); for (Map.Entry<String, AttributeValue> item : attributeMap.entrySet()) { Attribute a = new Attribute(item.getKey(), item.getValue().getS()); attrs.add(a); if (a.getName().equalsIgnoreCase(DatabaseKeys.LOGGING_KEY_KEY.getShortKey())) { ret.setName(a.getValue()); } } return ret; } private Map<String, AttributeValue> getTimingAttributes(TankResult result) { Map<String, AttributeValue> attributes = new HashMap<String, AttributeValue>(); String timestamp = ReportUtil.getTimestamp(result.getTimeStamp()); addAttribute(attributes, DatabaseKeys.TIMESTAMP_KEY.getShortKey(), timestamp); addAttribute(attributes, DatabaseKeys.REQUEST_NAME_KEY.getShortKey(), timestamp + "-" + UUID.randomUUID().toString()); addAttribute(attributes, DatabaseKeys.JOB_ID_KEY.getShortKey(), result.getJobId()); addAttribute(attributes, DatabaseKeys.LOGGING_KEY_KEY.getShortKey(), result.getRequestName()); addAttribute(attributes, DatabaseKeys.STATUS_CODE_KEY.getShortKey(), String.valueOf(result.getStatusCode())); addAttribute(attributes, DatabaseKeys.RESPONSE_TIME_KEY.getShortKey(), String.valueOf(result.getResponseTime())); addAttribute(attributes, DatabaseKeys.RESPONSE_SIZE_KEY.getShortKey(), String.valueOf(result.getResponseSize())); addAttribute(attributes, DatabaseKeys.INSTANCE_ID_KEY.getShortKey(), String.valueOf(result.getInstanceId())); addAttribute(attributes, DatabaseKeys.IS_ERROR_KEY.getShortKey(), String.valueOf(result.isError())); return attributes; } private void addAttribute(Map<String, AttributeValue> attributes, String key, String value) { if (value == null) { value = ""; } attributes.put(key, new AttributeValue().withS(value)); } private void addItemsToTable(String tableName, final BatchWriteItemRequest request) { boolean shouldRetry; int retries = 0; do { shouldRetry = false; try { BatchWriteItemResult result = dynamoDb.batchWriteItem(request); if (result != null) { try { List<ConsumedCapacity> consumedCapacity = result.getConsumedCapacity(); for (ConsumedCapacity cap : consumedCapacity) { logger.info(cap.getCapacityUnits()); } } catch (Exception e) { // ignore this } } } catch (AmazonServiceException e) { if (e instanceof ProvisionedThroughputExceededException) { try { DynamoDB db = new DynamoDB(dynamoDb); Table table = db.getTable(tableName); ProvisionedThroughputDescription oldThroughput = table.getDescription() .getProvisionedThroughput(); logger.info("ProvisionedThroughputExceeded throughput = " + oldThroughput); ProvisionedThroughput newThroughput = new ProvisionedThroughput() .withReadCapacityUnits( table.getDescription().getProvisionedThroughput().getReadCapacityUnits()) .withWriteCapacityUnits( getIncreasedThroughput(table.getDescription().getProvisionedThroughput() .getReadCapacityUnits())); if (!oldThroughput.equals(newThroughput)) { logger.info("Updating throughput to " + newThroughput); table.updateTable(newThroughput); table.waitForActive(); } } catch (Exception e1) { logger.error("Error increasing capacity: " + e, e); } } int status = e.getStatusCode(); if (status == HttpStatus.SC_INTERNAL_SERVER_ERROR || status == HttpStatus.SC_SERVICE_UNAVAILABLE) { shouldRetry = true; long delay = (long) (Math.random() * (Math.pow(4, retries++) * 100L)); try { Thread.sleep(delay); } catch (InterruptedException iex) { logger.error("Caught InterruptedException exception", iex); } } else { logger.error("Error writing to DB: " + e.getMessage()); throw new RuntimeException(e); } } } while (shouldRetry && retries < MAX_NUMBER_OF_RETRIES); } private Long getIncreasedThroughput(Long readCapacityUnits) { long ret = readCapacityUnits * 2; if (ret > MAX_WRITE_UNITS) { ret = MAX_WRITE_UNITS; } return ret; } private void waitForStatus(String tableName, TableStatus status) { logger.info("Waiting for " + tableName + " to become " + status.toString() + "..."); long startTime = System.currentTimeMillis(); long endTime = startTime + (10 * 60 * 1000); while (System.currentTimeMillis() < endTime) { try { Thread.sleep(1000 * 2); } catch (Exception e) { } try { DescribeTableRequest request = new DescribeTableRequest().withTableName(tableName); TableDescription tableDescription = dynamoDb.describeTable(request).getTable(); String tableStatus = tableDescription.getTableStatus(); logger.debug(" - current state: " + tableStatus); if (tableStatus.equals(status.toString())) return; } catch (AmazonServiceException ase) { if (ase.getErrorCode().equalsIgnoreCase("ResourceNotFoundException") == false) throw ase; } } throw new RuntimeException("Table " + tableName + " never went " + status.toString()); } private void waitForDelete(String tableName) { logger.info("Waiting for " + tableName + " to become deleted..."); long startTime = System.currentTimeMillis(); long endTime = startTime + (10 * 60 * 1000); while (System.currentTimeMillis() < endTime) { try { Thread.sleep(1000 * 2); } catch (Exception e) { } try { if (!hasTable(tableName)) { return; } } catch (AmazonServiceException ase) { if (ase.getErrorCode().equalsIgnoreCase("ResourceNotFoundException") == false) throw ase; } } throw new RuntimeException("Table " + tableName + " never deleted"); } /** * @param item * @return */ private Map<String, AttributeValue> itemToMap(Item item) { Map<String, AttributeValue> attributes = new HashMap<String, AttributeValue>(); for (Attribute attr : item.getAttributes()) { addAttribute(attributes, attr.getName(), attr.getValue()); } if (!attributes.containsKey(DatabaseKeys.JOB_ID_KEY.getShortKey())) { throw new RuntimeException("Item does not contain a job ID"); } else if (!attributes.containsKey(DatabaseKeys.REQUEST_NAME_KEY.getShortKey())) { AttributeValue attVal = attributes.get(DatabaseKeys.TIMESTAMP_KEY.getShortKey()); String timestamp = attVal != null ? attVal.getS() : ReportUtil.getTimestamp(new Date()); if (attVal == null) { addAttribute(attributes, DatabaseKeys.TIMESTAMP_KEY.getShortKey(), timestamp); } addAttribute(attributes, DatabaseKeys.REQUEST_NAME_KEY.getShortKey(), timestamp + "-" + UUID.randomUUID().toString()); } return attributes; } /** * @param tableName * @param requests */ private void sendBatch(final String tableName, List<WriteRequest> requests) { int numBatches = (int) Math.ceil(requests.size() / (BATCH_SIZE * 1D)); for (int i = 0; i < numBatches; i++) { Map<String, List<WriteRequest>> requestItems = new HashMap<String, List<WriteRequest>>(); List<WriteRequest> batch = requests.subList(i * BATCH_SIZE, Math.min(i * BATCH_SIZE + BATCH_SIZE, requests.size())); requestItems.put(tableName, batch); addItemsToTable(tableName, new BatchWriteItemRequest().withRequestItems(requestItems)); } } }