Linking DynamoDB queries to infra's entities

agent-bridge/src/main/java/com/newrelic/agent/bridge/CollectionFactory.java

@@ -8,6 +8,7 @@
 package com.newrelic.agent.bridge;
 
 import java.util.Map;
+import java.util.function.Function;
 
 /**
  * Allows instrumentation and bridge API implementations to use collections from third partly libraries without
@@ -27,4 +28,14 @@ public interface CollectionFactory {
      * @return concurrent-safe, weak-keyed map
      */
     <K, V> Map<K, V> createConcurrentWeakKeyedMap();
+
+    /**
+     * Wraps the provided function into one that will cache the results for future calls.
+     * @param loader the function that calculates the value.
+     * @param maxSize the max number of items to be cached.
+     * @return the cached item, or the result of the loader call.
+     * @param <K> the type of key
+     * @param <V> the type of value stored/returned
+     */
+    <K, V> Function<K, V> memoize(Function<K, V> loader, int maxSize);
 }

agent-bridge/src/main/java/com/newrelic/agent/bridge/DefaultCollectionFactory.java

@@ -10,11 +10,31 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 
+/**
+ * This implementation of {@link CollectionFactory} will only be used if the agent-bridge
+ * is being used by an application and the agent is NOT being loaded. Thus, it is unlikely
+ * that the objects created by this implementation are going to receive much use.
+ * So methods in this implementation do not need to implement all functional requirements
+ * of the methods in the interface, but they should not break under low use.
+ */
 public class DefaultCollectionFactory implements CollectionFactory {
 
     @Override
     public <K, V> Map<K, V> createConcurrentWeakKeyedMap() {
         return Collections.synchronizedMap(new WeakHashMap<K, V>());
     }
+
+    @Override
+    public <K, V> Function<K, V> memoize(Function<K, V> loader, int maxSize) {
+        Map<K, V> map = new ConcurrentHashMap<>();
+        return k -> map.computeIfAbsent(k, k1 -> {
+            if (map.size() >= maxSize) {
+                map.remove(map.keySet().iterator().next());
+            }
+            return loader.apply(k1);
+        });
+    }
 }

agent-bridge/src/main/java/com/newrelic/agent/bridge/external/MessageConsumeParameters.java

@@ -26,7 +26,7 @@ public class MessageConsumeParameters extends com.newrelic.api.agent.MessageCons
     @Deprecated
     protected MessageConsumeParameters(String library, DestinationType destinationType, String destinationName,
             InboundHeaders inboundHeaders) {
-        super(library, destinationType.toApiDestinationType(), destinationName, inboundHeaders);
+        super(library, destinationType.toApiDestinationType(), destinationName, inboundHeaders, null);
     }
 
 }

agent-bridge/src/main/java/com/newrelic/agent/bridge/external/MessageProduceParameters.java

@@ -16,7 +16,7 @@
 public class MessageProduceParameters extends com.newrelic.api.agent.MessageProduceParameters implements ExternalParameters {
 
     /**
-     * @Deprecated Do not use. Use {@link com.newrelic.api.agent.MessageProduceParameters#MessageProduceParameters} instead.
+     * @Deprecated Do not use. Use the fluent builder {@link com.newrelic.api.agent.MessageProduceParameters#library(String)} instead.
      *
      * @param library
      * @param destinationType
@@ -26,7 +26,7 @@ public class MessageProduceParameters extends com.newrelic.api.agent.MessageProd
     @Deprecated
     protected MessageProduceParameters(String library, DestinationType destinationType, String destinationName,
             OutboundHeaders outboundHeaders) {
-        super(library, destinationType.toApiDestinationType(), destinationName, outboundHeaders);
+        super(library, destinationType.toApiDestinationType(), destinationName, outboundHeaders, null);
     }
 
 }

agent-bridge/src/main/java/com/newrelic/agent/util/AwsAccountUtil.java

@@ -0,0 +1,68 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.util;
+
+import com.newrelic.agent.bridge.AgentBridge;
+
+import java.util.function.Function;
+
+
+public class AwsAccountUtil {
+    private static AwsAccountUtil INSTANCE = AwsAccountUtil.create();
+
+    private final Function<String, Long> CACHE = AgentBridge.collectionFactory.memoize(this::doDecodeAccount, 32);
+
+    public static AwsAccountUtil get() {
+        return INSTANCE;
+    }
+
+    public Long decodeAccount(String accessKey) {
+        return CACHE.apply(accessKey);
+    }
+
+    private Long doDecodeAccount(String awsAccessKeyId) {
+        String accessKeyWithoutPrefix = awsAccessKeyId.substring(4).toLowerCase();
+        long encodedAccount = base32Decode(accessKeyWithoutPrefix);
+        // magic number
+        long mask = 140737488355200L;
+        // magic incantation to find out the account
+        return (encodedAccount & mask) >> 7;
+    }
+
+    /**
+     * Character range is A-Z, 2-7. 'A' being 0 and '7', 31.
+     * Characters outside of this range will be considered 0.
+     * @param src the string to be decoded. Must be at least 10 characters.
+     * @return a long containing first 6 bytes of the base 32 decoded data.
+     * @throws ArrayIndexOutOfBoundsException if src has less than 10 characters
+     */
+    private long base32Decode(String src) {
+        long base = 0;
+        char[] chars = src.toCharArray();
+        // each char is 5 bits, we need 48 bits
+        for (int i = 0; i < 10; i++) {
+            char c = chars[i];
+            base <<= 5;
+            if (c >= 'a' && c <= 'z') {
+                base += c - 'a';
+            } else if (c >= '2' && c <= '7') {
+                base += c - '2' + 26;
+            }
+        }
+        // 50 bits were read, dropping the lowest 2
+        return base >> 2;
+    }
+
+    private AwsAccountUtil() {
+        // prevent instantiation of utility class
+    }
+
+    private static AwsAccountUtil create() {
+        return new AwsAccountUtil();
+    }
+}

agent-bridge/src/test/java/com/newrelic/agent/util/AwsAccountUtilTest.java

@@ -0,0 +1,25 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.util;
+
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+
+public class AwsAccountUtilTest {
+
+    @Test
+    public void decodeAccount() {
+        Long accountId = AwsAccountUtil.get().decodeAccount("FKKY6RVFFB77ZZZZZZZZ");
+        assertEquals(999999999999L, accountId.longValue());
+
+        accountId = AwsAccountUtil.get().decodeAccount("FKKYQAAAAAAAZZZZZZZZ");
+        assertEquals(1L, accountId.longValue());
+    }
+
+}

instrumentation/aws-java-sdk-dynamodb-1.11.106/src/main/java/com/amazonaws/services/dynamodbv2/AmazonDynamoDBClient_Instrumentation.java

@@ -10,6 +10,7 @@
 import com.amazonaws.AmazonWebServiceClient;
 import com.amazonaws.AmazonWebServiceRequest;
 import com.amazonaws.ClientConfiguration;
+import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.services.dynamodbv2.model.BatchGetItemRequest;
 import com.amazonaws.services.dynamodbv2.model.BatchGetItemResult;
 import com.amazonaws.services.dynamodbv2.model.BatchWriteItemRequest;
@@ -64,150 +65,151 @@ public abstract class AmazonDynamoDBClient_Instrumentation extends AmazonWebServ
     public AmazonDynamoDBClient_Instrumentation(ClientConfiguration clientConfiguration) {
         super(clientConfiguration);
     }
+    private final AWSCredentialsProvider awsCredentialsProvider = Weaver.callOriginal();
 
     @Trace(async = true, leaf = true)
     final CreateTableResult executeCreateTable(CreateTableRequest createTableRequest) {
         linkAndExpire(createTableRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "createTable",
-                createTableRequest.getTableName(), endpoint);
+                createTableRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final BatchGetItemResult executeBatchGetItem(BatchGetItemRequest batchGetItemRequest) {
         linkAndExpire(batchGetItemRequest);
-        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "batchGetItem", "batch", endpoint);
+        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "batchGetItem", "batch", endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final BatchWriteItemResult executeBatchWriteItem(BatchWriteItemRequest batchWriteItemRequest) {
         linkAndExpire(batchWriteItemRequest);
-        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "batchWriteItem", "batch", endpoint);
+        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "batchWriteItem", "batch", endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final DeleteItemResult executeDeleteItem(DeleteItemRequest deleteItemRequest) {
         linkAndExpire(deleteItemRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "deleteItem",
-                deleteItemRequest.getTableName(), endpoint);
+                deleteItemRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final DeleteTableResult executeDeleteTable(DeleteTableRequest deleteTableRequest) {
         linkAndExpire(deleteTableRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "deleteTable",
-                deleteTableRequest.getTableName(), endpoint);
+                deleteTableRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final DescribeLimitsResult executeDescribeLimits(DescribeLimitsRequest describeLimitsRequest) {
         linkAndExpire(describeLimitsRequest);
-        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "describeLimits", null, endpoint);
+        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "describeLimits", null, endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final DescribeTableResult executeDescribeTable(DescribeTableRequest describeTableRequest) {
         linkAndExpire(describeTableRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "describeTable",
-                describeTableRequest.getTableName(), endpoint);
+                describeTableRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final DescribeTimeToLiveResult executeDescribeTimeToLive(DescribeTimeToLiveRequest describeTimeToLiveRequest) {
         linkAndExpire(describeTimeToLiveRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "describeTimeToLive",
-                describeTimeToLiveRequest.getTableName(), endpoint);
+                describeTimeToLiveRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final GetItemResult executeGetItem(GetItemRequest getItemRequest) {
         linkAndExpire(getItemRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "getItem", getItemRequest.getTableName(),
-                endpoint);
+                endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final ListTablesResult executeListTables(ListTablesRequest listTablesRequest) {
         linkAndExpire(listTablesRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "listTables",
-                listTablesRequest.getExclusiveStartTableName(), endpoint);
+                listTablesRequest.getExclusiveStartTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final ListTagsOfResourceResult executeListTagsOfResource(ListTagsOfResourceRequest listTagsOfResourceRequest) {
         linkAndExpire(listTagsOfResourceRequest);
-        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "listTagsOfResource", null, endpoint);
+        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "listTagsOfResource", null, endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final PutItemResult executePutItem(PutItemRequest putItemRequest) {
         linkAndExpire(putItemRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "putItem", putItemRequest.getTableName(),
-                endpoint);
+                endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final QueryResult executeQuery(QueryRequest queryRequest) {
         linkAndExpire(queryRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "query", queryRequest.getTableName(),
-                endpoint);
+                endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
 
     }
 
     @Trace(async = true, leaf = true)
     final ScanResult executeScan(ScanRequest scanRequest) {
         linkAndExpire(scanRequest);
-        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "scan", scanRequest.getTableName(), endpoint);
+        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "scan", scanRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final TagResourceResult executeTagResource(TagResourceRequest tagResourceRequest) {
         linkAndExpire(tagResourceRequest);
-        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "tagResource", null, endpoint);
+        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "tagResource", null, endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final UntagResourceResult executeUntagResource(UntagResourceRequest untagResourceRequest) {
         linkAndExpire(untagResourceRequest);
-        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "untagResource", null, endpoint);
+        DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "untagResource", null, endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final UpdateItemResult executeUpdateItem(UpdateItemRequest updateItemRequest) {
         linkAndExpire(updateItemRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "updateItem",
-                updateItemRequest.getTableName(), endpoint);
+                updateItemRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final UpdateTableResult executeUpdateTable(UpdateTableRequest updateTableRequest) {
         linkAndExpire(updateTableRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "updateTable",
-                updateTableRequest.getTableName(), endpoint);
+                updateTableRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }
 
     @Trace(async = true, leaf = true)
     final UpdateTimeToLiveResult executeUpdateTimeToLive(UpdateTimeToLiveRequest updateTimeToLiveRequest) {
        linkAndExpire(updateTimeToLiveRequest);
         DynamoDBMetricUtil.metrics(NewRelic.getAgent().getTracedMethod(), "updateTimeToLive",
-                updateTimeToLiveRequest.getTableName(), endpoint);
+                updateTimeToLiveRequest.getTableName(), endpoint, awsCredentialsProvider);
         return Weaver.callOriginal();
     }