AWS SDK instrumentation for S3 and DynamoDB

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/AbstractSpan.java

@@ -113,6 +113,8 @@ public abstract class AbstractSpan<T extends AbstractSpan<T>> implements Recycla
     @Nullable
     protected volatile String type;
 
+    protected volatile boolean sync = true;
+
     protected final AtomicReference<Span> bufferedSpan = new AtomicReference<>();
 
     @Nullable
@@ -341,11 +343,16 @@ public T withName(@Nullable String name, int priority, boolean overrideIfSamePri
         return thiz();
     }
 
-    public T withType(@Nullable String type){
+    public T withType(@Nullable String type) {
         this.type = normalizeEmpty(type);
         return thiz();
     }
 
+    public T withSync(boolean sync) {
+        this.sync = sync;
+        return thiz();
+    }
+
     @Nullable
     protected static String normalizeEmpty(@Nullable String value) {
         return value == null || value.isEmpty() ? null : value;
@@ -367,6 +374,7 @@ public void resetState() {
         finished = true;
         name.setLength(0);
         type = null;
+        sync = true;
         timestamp.set(0L);
         endTimestamp.set(0L);
         traceContext.resetState();
@@ -730,6 +738,10 @@ public String getType() {
         return type;
     }
 
+    public boolean isSync() {
+        return sync;
+    }
+
     private String normalizeType(@Nullable String type) {
         if (type == null || type.isEmpty()) {
             return "custom";

apm-agent-core/src/main/java/co/elastic/apm/agent/report/serialize/DslJsonSerializer.java

@@ -699,7 +699,9 @@ private void serializeSpan(final Span span) {
         jw.writeByte(OBJECT_START);
         writeField("name", span.getNameForSerialization());
         writeTimestamp(span.getTimestamp());
-
+        if(!span.isSync()){
+            writeField("sync", span.isSync());
+        }
         writeField("outcome", span.getOutcome().toString());
         serializeTraceContext(traceContext, true);
         writeField("duration", span.getDurationMs());

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-1-plugin/pom.xml

@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>apm-aws-sdk</artifactId>
+        <groupId>co.elastic.apm</groupId>
+        <version>1.30.2-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>apm-aws-sdk-1-plugin</artifactId>
+    <name>${project.groupId}:${project.artifactId}</name>
+
+    <properties>
+        <apm-agent-parent.base.dir>${project.basedir}/../../..</apm-agent-parent.base.dir>
+        <version.aws.sdk>[1.12.1,)</version.aws.sdk>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>${project.groupId}</groupId>
+            <artifactId>apm-aws-sdk-common</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-java-sdk-s3</artifactId>
+            <version>${version.aws.sdk}</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-java-sdk-dynamodb</artifactId>
+            <version>${version.aws.sdk}</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>${project.groupId}</groupId>
+            <artifactId>apm-aws-sdk-common</artifactId>
+            <version>${project.version}</version>
+            <type>test-jar</type>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+</project>

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-1-plugin/src/main/java/co/elastic/apm/agent/awssdk/v1/AmazonHttpClientInstrumentation.java

@@ -0,0 +1,101 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package co.elastic.apm.agent.awssdk.v1;
+
+import co.elastic.apm.agent.awssdk.v1.helper.DynamoDbHelper;
+import co.elastic.apm.agent.awssdk.v1.helper.S3Helper;
+import co.elastic.apm.agent.bci.TracerAwareInstrumentation;
+import co.elastic.apm.agent.impl.transaction.Outcome;
+import co.elastic.apm.agent.impl.transaction.Span;
+import com.amazonaws.Request;
+import com.amazonaws.handlers.HandlerContextKey;
+import com.amazonaws.http.ExecutionContext;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import javax.annotation.Nullable;
+import java.util.Collection;
+import java.util.Collections;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+public class AmazonHttpClientInstrumentation extends TracerAwareInstrumentation {
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return named("com.amazonaws.http.AmazonHttpClient");
+    }
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return named("execute")
+            .and(takesArguments(5))
+            .and(takesArgument(0, named("com.amazonaws.Request")))
+            .and(takesArgument(3, named("com.amazonaws.http.ExecutionContext")));
+    }
+
+    @Override
+    public Collection<String> getInstrumentationGroupNames() {
+        return Collections.singleton("aws-sdk");
+    }
+
+
+    public static class AdviceClass {
+
+        @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
+        @Nullable
+        public static Object enterInvoke(@Advice.Argument(value = 0) Request<?> request,
+                                         @Advice.Argument(value = 3) ExecutionContext executionContext) {
+            String awsService = request.getHandlerContext(HandlerContextKey.SERVICE_ID);
+            Span span = null;
+            if (awsService.startsWith("S3")) {
+                span = S3Helper.getInstance().startSpan(request, request.getEndpoint(), executionContext);
+            } else if ("DynamoDB".equalsIgnoreCase(awsService)) {
+                span = DynamoDbHelper.getInstance().startSpan(request, request.getEndpoint(), executionContext);
+                DynamoDbHelper.getInstance().removeTableNameForKey(request.getOriginalRequest());
+            }
+
+            if (span != null) {
+                span.activate();
+            }
+
+            return span;
+        }
+
+        @Advice.OnMethodExit(suppress = Throwable.class, inline = false, onThrowable = Throwable.class)
+        public static void exitInvoke(@Nullable @Advice.Enter Object spanObj,
+                                      @Nullable @Advice.Thrown Throwable thrown) {
+            if (spanObj instanceof Span) {
+                Span span = (Span) spanObj;
+                span.deactivate();
+                if (thrown != null) {
+                    span.captureException(thrown);
+                    span.withOutcome(Outcome.FAILURE);
+                } else {
+                    span.withOutcome(Outcome.SUCCESS);
+                }
+                span.end();
+            }
+        }
+    }
+}

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-1-plugin/src/main/java/co/elastic/apm/agent/awssdk/v1/DynamoDBGetTableNameInstrumentation.java

@@ -0,0 +1,65 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package co.elastic.apm.agent.awssdk.v1;
+
+import co.elastic.apm.agent.awssdk.v1.helper.DynamoDbHelper;
+import co.elastic.apm.agent.bci.TracerAwareInstrumentation;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import javax.annotation.Nullable;
+import java.util.Collection;
+import java.util.Collections;
+
+import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
+import static net.bytebuddy.matcher.ElementMatchers.nameEndsWith;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+public class DynamoDBGetTableNameInstrumentation extends TracerAwareInstrumentation {
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return nameStartsWith("com.amazonaws.services.dynamodbv2.model.")
+            .and(nameEndsWith("Request"))
+            .and(declaresMethod(named("getTableName")));
+    }
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return named("getTableName");
+    }
+
+    @Override
+    public Collection<String> getInstrumentationGroupNames() {
+        return Collections.singleton("aws-sdk");
+    }
+
+
+    public static class AdviceClass {
+        @Advice.OnMethodExit(suppress = Throwable.class, inline = false, onThrowable = Throwable.class)
+        public static void exitGetTableName(@Advice.This Object requestObject, @Nullable @Advice.Return String tableName) {
+            if (tableName != null && !DynamoDbHelper.getInstance().hasTableNameForKey(requestObject)) {
+                DynamoDbHelper.getInstance().putTableName(requestObject, tableName);
+            }
+        }
+    }
+}

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-1-plugin/src/main/java/co/elastic/apm/agent/awssdk/v1/helper/DynamoDbHelper.java

@@ -0,0 +1,66 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package co.elastic.apm.agent.awssdk.v1.helper;
+
+import co.elastic.apm.agent.awssdk.common.AbstractDynamoDBInstrumentationHelper;
+import co.elastic.apm.agent.impl.ElasticApmTracer;
+import co.elastic.apm.agent.impl.GlobalTracer;
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
+import com.amazonaws.Request;
+import com.amazonaws.http.ExecutionContext;
+
+import javax.annotation.Nullable;
+
+public class DynamoDbHelper extends AbstractDynamoDBInstrumentationHelper<Request<?>, ExecutionContext> {
+
+    static final WeakMap<Object, String> dynamoDbRequestToTableNameMap = WeakConcurrent.buildMap();
+
+    @Nullable
+    private static DynamoDbHelper INSTANCE;
+
+    public static DynamoDbHelper getInstance() {
+        if (INSTANCE == null) {
+            INSTANCE = new DynamoDbHelper(GlobalTracer.requireTracerImpl());
+        }
+        return INSTANCE;
+    }
+
+
+    public DynamoDbHelper(ElasticApmTracer tracer) {
+        super(tracer, SdkV1DataSource.getInstance());
+    }
+
+    public void putTableName(Object key, String tableName) {
+        dynamoDbRequestToTableNameMap.put(key, tableName);
+    }
+
+    @Nullable
+    public String getTableName(Object key) {
+        return dynamoDbRequestToTableNameMap.get(key);
+    }
+
+    public void removeTableNameForKey(Object key) {
+        dynamoDbRequestToTableNameMap.remove(key);
+    }
+
+    public boolean hasTableNameForKey(Object key) {
+        return dynamoDbRequestToTableNameMap.containsKey(key);
+    }
+}

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-1-plugin/src/main/java/co/elastic/apm/agent/awssdk/v1/helper/S3Helper.java

@@ -0,0 +1,44 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package co.elastic.apm.agent.awssdk.v1.helper;
+
+import co.elastic.apm.agent.awssdk.common.AbstractS3InstrumentationHelper;
+import co.elastic.apm.agent.impl.ElasticApmTracer;
+import co.elastic.apm.agent.impl.GlobalTracer;
+import com.amazonaws.Request;
+import com.amazonaws.http.ExecutionContext;
+
+import javax.annotation.Nullable;
+
+public class S3Helper extends AbstractS3InstrumentationHelper<Request<?>, ExecutionContext> {
+
+    @Nullable
+    private static S3Helper INSTANCE;
+
+    public static S3Helper getInstance() {
+        if (INSTANCE == null) {
+            INSTANCE = new S3Helper(GlobalTracer.requireTracerImpl());
+        }
+        return INSTANCE;
+    }
+
+    public S3Helper(ElasticApmTracer tracer) {
+        super(tracer, SdkV1DataSource.getInstance());
+    }
+}