From fa3ae0921fea42a7602694dbfca14892324be2cb Mon Sep 17 00:00:00 2001
From: Anuraag Agrawal <aanuraag@amazon.co.jp>
Date: Fri, 4 Sep 2020 14:59:02 +0900
Subject: [PATCH] Use ClassValue and full name cache in AWS SDK 1.1 (#1153)

* Use ClassValue and full name cache in AWS SDK 1.1

* Mysterious class

* Muzzle
---
 .../v1_11/AWSClientInstrumentation.java       |  1 +
 .../v1_11/AWSHttpClientInstrumentation.java   |  4 ++-
 .../auto/awssdk/v1_11/AwsSdkClientTracer.java | 34 +++++++++++--------
 3 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AWSClientInstrumentation.java b/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AWSClientInstrumentation.java
index ab311ea752..ffc0329f84 100644
--- a/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AWSClientInstrumentation.java
+++ b/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AWSClientInstrumentation.java
@@ -54,6 +54,7 @@ public final class AWSClientInstrumentation extends Instrumenter.Default {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".AwsSdkClientTracer",
+      packageName + ".AwsSdkClientTracer$NamesCache",
       packageName + ".RequestMeta",
       packageName + ".TracingRequestHandler",
     };
diff --git a/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AWSHttpClientInstrumentation.java b/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AWSHttpClientInstrumentation.java
index accf4d9061..0abe33e2dd 100644
--- a/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AWSHttpClientInstrumentation.java
+++ b/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AWSHttpClientInstrumentation.java
@@ -56,7 +56,9 @@ public class AWSHttpClientInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".AwsSdkClientTracer", packageName + ".RequestMeta",
+      packageName + ".AwsSdkClientTracer",
+      packageName + ".AwsSdkClientTracer$NamesCache",
+      packageName + ".RequestMeta",
     };
   }
 
diff --git a/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AwsSdkClientTracer.java b/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AwsSdkClientTracer.java
index 3b29447934..2998532517 100644
--- a/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AwsSdkClientTracer.java
+++ b/instrumentation/aws-sdk/aws-sdk-1.11/src/main/java/io/opentelemetry/instrumentation/auto/awssdk/v1_11/AwsSdkClientTracer.java
@@ -29,7 +29,6 @@ import io.opentelemetry.context.propagation.TextMapPropagator.Setter;
 import io.opentelemetry.instrumentation.api.tracer.HttpClientTracer;
 import io.opentelemetry.trace.Span;
 import java.net.URI;
-import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class AwsSdkClientTracer extends HttpClientTracer<Request<?>, Request<?>, Response<?>> {
@@ -38,8 +37,7 @@ public class AwsSdkClientTracer extends HttpClientTracer<Request<?>, Request<?>,
 
   public static final AwsSdkClientTracer TRACER = new AwsSdkClientTracer();
 
-  private final Map<String, String> serviceNames = new ConcurrentHashMap<>();
-  private final Map<Class, String> operationNames = new ConcurrentHashMap<>();
+  private final NamesCache namesCache = new NamesCache();
 
   public AwsSdkClientTracer() {}
 
@@ -50,7 +48,7 @@ public class AwsSdkClientTracer extends HttpClientTracer<Request<?>, Request<?>,
     }
     String awsServiceName = request.getServiceName();
     Class<?> awsOperation = request.getOriginalRequest().getClass();
-    return remapServiceName(awsServiceName) + "." + remapOperationName(awsOperation);
+    return qualifiedOperation(awsServiceName, awsOperation);
   }
 
   public Span startSpan(Request<?> request, RequestMeta requestMeta) {
@@ -95,18 +93,17 @@ public class AwsSdkClientTracer extends HttpClientTracer<Request<?>, Request<?>,
     return super.onResponse(span, response);
   }
 
-  private String remapServiceName(String serviceName) {
-    if (!serviceNames.containsKey(serviceName)) {
-      serviceNames.put(serviceName, serviceName.replace("Amazon", "").trim());
+  private String qualifiedOperation(String service, Class<?> operation) {
+    ConcurrentHashMap<String, String> cache = namesCache.get(operation);
+    String qualified = cache.get(service);
+    if (qualified == null) {
+      qualified =
+          service.replace("Amazon", "").trim()
+              + '.'
+              + operation.getSimpleName().replace("Request", "");
+      cache.put(service, qualified);
     }
-    return serviceNames.get(serviceName);
-  }
-
-  private String remapOperationName(Class<?> awsOperation) {
-    if (!operationNames.containsKey(awsOperation)) {
-      operationNames.put(awsOperation, awsOperation.getSimpleName().replace("Request", ""));
-    }
-    return operationNames.get(awsOperation);
+    return qualified;
   }
 
   @Override
@@ -143,4 +140,11 @@ public class AwsSdkClientTracer extends HttpClientTracer<Request<?>, Request<?>,
   protected String getInstrumentationName() {
     return "io.opentelemetry.auto.aws-sdk-1.11";
   }
+
+  static final class NamesCache extends ClassValue<ConcurrentHashMap<String, String>> {
+    @Override
+    protected ConcurrentHashMap<String, String> computeValue(Class<?> type) {
+      return new ConcurrentHashMap<>();
+    }
+  }
 }