From 90947a093b3e0dce2034e6df8b84ab2021c8277a Mon Sep 17 00:00:00 2001 From: Oren Ben-Meir Date: Wed, 21 Aug 2024 17:12:04 -0400 Subject: [PATCH 1/8] Add traces for the kinesis 2.0 sdk --- .../aws-java-sdk-kinesis-2.0.6/build.gradle | 14 ++ .../services/lambda/KinesisUtil.java | 6 + .../AsyncClientHandler_Instrumentation.java | 28 ++++ ...ultKinesisAsyncClient_Instrumentation.java | 146 ++++++++++++++++++ .../DefaultKinesisClient_Instrumentation.java | 96 ++++++++++++ .../src/test/resources/dt_enabled.yml | 5 + settings.gradle | 1 + 7 files changed, 296 insertions(+) create mode 100644 instrumentation/aws-java-sdk-kinesis-2.0.6/build.gradle create mode 100644 instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java create mode 100644 instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/core/client/handler/AsyncClientHandler_Instrumentation.java create mode 100644 instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java create mode 100644 instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java create mode 100644 instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/resources/dt_enabled.yml diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/build.gradle b/instrumentation/aws-java-sdk-kinesis-2.0.6/build.gradle new file mode 100644 index 0000000000..bd30d61ce9 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/build.gradle @@ -0,0 +1,14 @@ +dependencies { + implementation(project(":agent-bridge")) + implementation("software.amazon.awssdk:kinesis:2.20.45") +} + +jar { + manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.aws-java-sdk-kinesis-2.0.6' } +} + +verifyInstrumentation { + passes 'software.amazon.awssdk:kinesis:[2.1.0,)' + exclude 'software.amazon.awssdk:kinesis:2.17.200' // this version failed the test, but the next one works again. + excludeRegex '.*-preview-[0-9a-f]+' +} \ No newline at end of file diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java new file mode 100644 index 0000000000..c049f93a80 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java @@ -0,0 +1,6 @@ +package com.agent.instrumentation.awsjavasdk2.services.lambda; + +public class KinesisUtil { + private KinesisUtil() {} + +} diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/core/client/handler/AsyncClientHandler_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/core/client/handler/AsyncClientHandler_Instrumentation.java new file mode 100644 index 0000000000..17f8e249b5 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/core/client/handler/AsyncClientHandler_Instrumentation.java @@ -0,0 +1,28 @@ +package software.amazon.awssdk.core.client.handler; + +import com.newrelic.api.agent.Trace; +import com.newrelic.api.agent.weaver.MatchType; +import com.newrelic.api.agent.weaver.Weave; +import com.newrelic.api.agent.weaver.Weaver; +import software.amazon.awssdk.core.SdkRequest; +import software.amazon.awssdk.core.SdkResponse; +import software.amazon.awssdk.core.async.AsyncResponseTransformer; + +import java.util.concurrent.CompletableFuture; + +@Weave(originalName = "software.amazon.awssdk.core.client.handler.AsyncClientHandler", type = MatchType.Interface) +public class AsyncClientHandler_Instrumentation { + // This prevents further traces from forming when using the async client + @Trace(leaf = true, excludeFromTransactionTrace = true) + public CompletableFuture execute( + ClientExecutionParams executionParams) { + return Weaver.callOriginal(); + } + + @Trace(leaf = true, excludeFromTransactionTrace = true) + public CompletableFuture execute( + ClientExecutionParams executionParams, + AsyncResponseTransformer asyncResponseTransformer) { + return Weaver.callOriginal(); + } +} diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java new file mode 100644 index 0000000000..5ddf53fc7f --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java @@ -0,0 +1,146 @@ +package software.amazon.awssdk.services.kinesis; + +import com.newrelic.agent.bridge.AgentBridge; +import com.newrelic.api.agent.NewRelic; +import com.newrelic.api.agent.Segment; +import com.newrelic.api.agent.Trace; +import com.newrelic.api.agent.weaver.MatchType; +import com.newrelic.api.agent.weaver.Weave; +import com.newrelic.api.agent.weaver.Weaver; +import software.amazon.awssdk.services.kinesis.model.DescribeLimitsRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeLimitsResponse; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamConsumerRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamConsumerResponse; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamResponse; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest; +import software.amazon.awssdk.services.kinesis.model.GetRecordsRequest; +import software.amazon.awssdk.services.kinesis.model.GetRecordsResponse; +import software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest; +import software.amazon.awssdk.services.kinesis.model.GetShardIteratorResponse; +import software.amazon.awssdk.services.kinesis.model.ListShardsRequest; +import software.amazon.awssdk.services.kinesis.model.ListShardsResponse; +import software.amazon.awssdk.services.kinesis.model.PutRecordRequest; +import software.amazon.awssdk.services.kinesis.model.PutRecordResponse; +import software.amazon.awssdk.services.kinesis.model.PutRecordsRequest; +import software.amazon.awssdk.services.kinesis.model.PutRecordsResponse; + +import java.util.concurrent.CompletableFuture; +import java.util.function.BiConsumer; + +@Weave(originalName = "software.amazon.awssdk.services.kinesis.DefaultKinesisAsyncClient", type = MatchType.ExactClass) +class DefaultKinesisAsyncClient_Instrumentation { + // Todo: report external call with the CloudParameters API + + @Trace + public CompletableFuture describeLimits(DescribeLimitsRequest describeLimitsRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("Kinesis","describeLimits"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + @Trace + public CompletableFuture describeStream(DescribeStreamRequest describeStreamRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("Kinesis", "describeStream"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + @Trace + public CompletableFuture describeStreamConsumer(DescribeStreamConsumerRequest describeStreamConsumerRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("Kinesis", "describeStreamConsumer"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + @Trace + public CompletableFuture describeStreamSummary(DescribeStreamSummaryRequest describeStreamSummaryRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("KinesisAsyncClient", "describeStreamSummary"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + @Trace + public CompletableFuture getRecords(GetRecordsRequest getRecordsRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("Kinesis", "getRecords"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + + @Trace + public CompletableFuture getShardIterator(GetShardIteratorRequest getShardIteratorRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("Kinesis", "getShardIterator"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + @Trace + public CompletableFuture listShards(ListShardsRequest listShardsRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("Kinesis", "listShards"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + + @Trace + public CompletableFuture putRecord(PutRecordRequest putRecordRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("Kinesis", "putRecord"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + @Trace + public CompletableFuture putRecords(PutRecordsRequest putRecordsRequest) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment("Kinesis", "putRecords"); + CompletableFuture response = Weaver.callOriginal(); + if (response == null) { + return null; + } + return response.whenComplete(new NrSegmentCompletion<>(segment, Weaver.getImplementationTitle())); + } + + private static class NrSegmentCompletion implements BiConsumer { + + private final Segment segment; + private final String implementationTitle; + + NrSegmentCompletion(Segment segment, String implementationTitle) { + this.segment = segment; + this.implementationTitle = implementationTitle; + } + + @Override + public void accept(R r, T t) { + try { + segment.end(); + } catch (Throwable t1) { + AgentBridge.instrumentation.noticeInstrumentationError(t1, implementationTitle); + } + } + } +} diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java new file mode 100644 index 0000000000..ed81c0891e --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java @@ -0,0 +1,96 @@ +package software.amazon.awssdk.services.kinesis; + +import com.newrelic.api.agent.NewRelic; +import com.newrelic.api.agent.Trace; +import com.newrelic.api.agent.TracedMethod; +import com.newrelic.api.agent.weaver.MatchType; +import com.newrelic.api.agent.weaver.Weave; +import com.newrelic.api.agent.weaver.Weaver; +import software.amazon.awssdk.services.kinesis.model.DescribeLimitsRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeLimitsResponse; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamConsumerRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamConsumerResponse; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamResponse; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryResponse; +import software.amazon.awssdk.services.kinesis.model.GetRecordsRequest; +import software.amazon.awssdk.services.kinesis.model.GetRecordsResponse; +import software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest; +import software.amazon.awssdk.services.kinesis.model.GetShardIteratorResponse; +import software.amazon.awssdk.services.kinesis.model.ListShardsRequest; +import software.amazon.awssdk.services.kinesis.model.ListShardsResponse; +import software.amazon.awssdk.services.kinesis.model.PutRecordRequest; +import software.amazon.awssdk.services.kinesis.model.PutRecordResponse; +import software.amazon.awssdk.services.kinesis.model.PutRecordsRequest; +import software.amazon.awssdk.services.kinesis.model.PutRecordsResponse; + +@Weave(originalName = "software.amazon.awssdk.services.kinesis.DefaultKinesisClient", type = MatchType.ExactClass) +class DefaultKinesisClient_Instrumentation { + // Todo: report external call with the CloudParameters API + + @Trace(leaf=true) + public DescribeLimitsResponse describeLimits(DescribeLimitsRequest describeLimitsRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "describeLimits"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + public DescribeStreamResponse describeStream(DescribeStreamRequest describeStreamRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "describeStream"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + public DescribeStreamConsumerResponse describeStreamConsumer(DescribeStreamConsumerRequest describeStreamConsumerRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "describeStreamConsumer"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + public DescribeStreamSummaryResponse describeStreamSummary(DescribeStreamSummaryRequest describeStreamSummaryRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "describeStreamSummary"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + public GetRecordsResponse getRecords(GetRecordsRequest getRecordsRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "getRecords"); + return Weaver.callOriginal(); + } + + + @Trace(leaf=true) + public GetShardIteratorResponse getShardIterator(GetShardIteratorRequest getShardIteratorRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "getShardIterator"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + public ListShardsResponse listShards(ListShardsRequest listShardsRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "listShards"); + return Weaver.callOriginal(); + } + + + @Trace(leaf=true) + public PutRecordResponse putRecord(PutRecordRequest putRecordRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "putRecord"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + public PutRecordsResponse putRecords(PutRecordsRequest putRecordsRequest) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName("Kinesis", "putRecords"); + return Weaver.callOriginal(); + } +} diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/resources/dt_enabled.yml b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/resources/dt_enabled.yml new file mode 100644 index 0000000000..53b0968002 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/resources/dt_enabled.yml @@ -0,0 +1,5 @@ +common: &default_settings + distributed_tracing: + enabled: true + span_events: + enabled: true \ No newline at end of file diff --git a/settings.gradle b/settings.gradle index 30d5ef5b0a..62a48531c7 100644 --- a/settings.gradle +++ b/settings.gradle @@ -77,6 +77,7 @@ include 'instrumentation:aws-java-sdk-sns-1.11.12' include 'instrumentation:aws-java-sdk-sqs-2.1.0' include 'instrumentation:aws-java-sdk-dynamodb-1.11.106' include 'instrumentation:aws-java-sdk-dynamodb-2.15.34' +include 'instrumentation:aws-java-sdk-kinesis-2.0.6' include 'instrumentation:aws-java-sdk-sns-2.0' include 'instrumentation:aws-wrap-0.7.0' include 'instrumentation:akka-2.2' From 483f097e232472e9485eb93a8156935beeb8523c Mon Sep 17 00:00:00 2001 From: Oren Ben-Meir Date: Tue, 27 Aug 2024 15:01:44 -0400 Subject: [PATCH 2/8] Instrument v1 kinesis sdk --- .../build.gradle | 12 ++ .../services/lambda/KinesisUtil.java | 47 +++++ .../services/lambda/SegmentHandler.java | 31 +++ .../amazonaws/AmazonWebServiceRequest.java | 11 ++ .../AsyncHandler_Instrumentation.java | 42 ++++ ...zonKinesisAsyncClient_Instrumentation.java | 178 +++++++++++++++++ .../AmazonKinesisClient_Instrumentation.java | 183 ++++++++++++++++++ .../src/test/resources/dt_enabled.yml | 5 + .../services/lambda/KinesisUtil.java | 8 +- ...ultKinesisAsyncClient_Instrumentation.java | 54 +++--- .../DefaultKinesisClient_Instrumentation.java | 55 +++--- settings.gradle | 1 + 12 files changed, 569 insertions(+), 58 deletions(-) create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/build.gradle create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest.java create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/resources/dt_enabled.yml diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/build.gradle b/instrumentation/aws-java-sdk-kinesis-1.11.106/build.gradle new file mode 100644 index 0000000000..c3aa9ff565 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/build.gradle @@ -0,0 +1,12 @@ +dependencies { + implementation(project(":agent-bridge")) + implementation("com.amazonaws:aws-java-sdk-kinesis:1.11.106") +} + +jar { + manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.aws-java-sdk-kinesis-1.11.106' } +} + +verifyInstrumentation { + passesOnly 'com.amazonaws:aws-java-sdk-kinesis:[1.11.106,)' +} \ No newline at end of file diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java new file mode 100644 index 0000000000..fc99d130cf --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java @@ -0,0 +1,47 @@ +package com.agent.instrumentation.awsjavasdk2.services.lambda; + +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.handlers.AsyncHandler_Instrumentation; +import com.newrelic.agent.bridge.AgentBridge; +import com.newrelic.api.agent.CloudParameters; +import com.newrelic.api.agent.NewRelic; +import com.newrelic.api.agent.Segment; +import com.newrelic.api.agent.Token; +import com.newrelic.api.agent.TracedMethod; + +public class KinesisUtil { + + public static final String PLATFORM = "aws_kinesis_data_streams"; + public static final String TRACE_CATEGORY = "Kinesis"; + private KinesisUtil() {} + + public static void setupToken(AsyncHandler_Instrumentation asyncHandler, AmazonWebServiceRequest request) { + if (AgentBridge.getAgent().getTransaction(false) != null) { + if (asyncHandler != null) { + asyncHandler.token = NewRelic.getAgent().getTransaction().getToken(); + } + if (request != null) { + request.token = NewRelic.getAgent().getTransaction().getToken(); + } + } + } + + public static void linkAndExpireToken(AmazonWebServiceRequest request) { + if (request != null && request.token != null) { + request.token.linkAndExpire(); + request.token = null; + } + } + + public static void setTraceDetails(String kinesisOperation) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName(TRACE_CATEGORY, kinesisOperation); + tracedMethod.reportAsExternal(createCloudParams()); + } + + public static CloudParameters createCloudParams() { + // Todo: add arn to cloud parameters + return CloudParameters.provider(PLATFORM).build(); + } + +} diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java new file mode 100644 index 0000000000..1a042efd89 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java @@ -0,0 +1,31 @@ +package com.agent.instrumentation.awsjavasdk2.services.lambda; + +import com.newrelic.agent.bridge.AgentBridge; +import com.newrelic.api.agent.Segment; + +import java.util.concurrent.CompletableFuture; + +public class SegmentHandler { + private final CompletableFuture completableFuture; + private final Segment segment; + private final String implementationTitle; + + public SegmentHandler(CompletableFuture completableFuture, Segment segment, String implementationTitle) { + this.completableFuture = completableFuture; + this.segment = segment; + this.implementationTitle = implementationTitle; + } + + public CompletableFuture newSegmentCompletionStage() { + if (completableFuture == null) { + return null; + } + return completableFuture.whenComplete((r, t) -> { + try { + segment.end(); + } catch (Throwable t1) { + AgentBridge.instrumentation.noticeInstrumentationError(t1, implementationTitle); + } + }); + } +} \ No newline at end of file diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest.java new file mode 100644 index 0000000000..63202e7543 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest.java @@ -0,0 +1,11 @@ +package com.amazonaws; + +import com.newrelic.api.agent.Token; +import com.newrelic.api.agent.weaver.NewField; +import com.newrelic.api.agent.weaver.Weave; + +@Weave(originalName = "com.amazonaws.AmazonWebServiceRequest") +public class AmazonWebServiceRequest { + @NewField + public Token token; +} diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java new file mode 100644 index 0000000000..9d826c5ba0 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java @@ -0,0 +1,42 @@ +/* + * + * * Copyright 2020 New Relic Corporation. All rights reserved. + * * SPDX-License-Identifier: Apache-2.0 + * + */ + +package com.amazonaws.handlers; + +import com.amazonaws.AmazonWebServiceRequest; +import com.newrelic.api.agent.Segment; +import com.newrelic.api.agent.Token; +import com.newrelic.api.agent.Trace; +import com.newrelic.api.agent.weaver.MatchType; +import com.newrelic.api.agent.weaver.NewField; +import com.newrelic.api.agent.weaver.Weave; +import com.newrelic.api.agent.weaver.Weaver; + +@Weave(originalName ="com.amazonaws.handlers.AsyncHandler", type = MatchType.Interface) +public class AsyncHandler_Instrumentation { + + @NewField + public Token token; + + @Trace(async = true) + public void onError(Exception exception) { + if (token != null) { + token.linkAndExpire(); + token = null; + } + Weaver.callOriginal(); + } + + @Trace(async = true) + public void onSuccess(REQUEST request, RESULT result) { + if (token != null) { + token.linkAndExpire(); + token = null; + } + Weaver.callOriginal(); + } +} \ No newline at end of file diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java new file mode 100644 index 0000000000..6547a01859 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java @@ -0,0 +1,178 @@ +package com.amazonaws.services.kinesis; + +import com.agent.instrumentation.awsjavasdk2.services.lambda.KinesisUtil; +import com.amazonaws.handlers.AsyncHandler_Instrumentation; +import com.amazonaws.services.kinesis.model.AddTagsToStreamRequest; +import com.amazonaws.services.kinesis.model.AddTagsToStreamResult; +import com.amazonaws.services.kinesis.model.CreateStreamRequest; +import com.amazonaws.services.kinesis.model.CreateStreamResult; +import com.amazonaws.services.kinesis.model.DecreaseStreamRetentionPeriodRequest; +import com.amazonaws.services.kinesis.model.DecreaseStreamRetentionPeriodResult; +import com.amazonaws.services.kinesis.model.DeleteStreamRequest; +import com.amazonaws.services.kinesis.model.DeleteStreamResult; +import com.amazonaws.services.kinesis.model.DescribeLimitsRequest; +import com.amazonaws.services.kinesis.model.DescribeLimitsResult; +import com.amazonaws.services.kinesis.model.DescribeStreamRequest; +import com.amazonaws.services.kinesis.model.DescribeStreamResult; +import com.amazonaws.services.kinesis.model.DisableEnhancedMonitoringRequest; +import com.amazonaws.services.kinesis.model.DisableEnhancedMonitoringResult; +import com.amazonaws.services.kinesis.model.EnableEnhancedMonitoringRequest; +import com.amazonaws.services.kinesis.model.EnableEnhancedMonitoringResult; +import com.amazonaws.services.kinesis.model.GetRecordsRequest; +import com.amazonaws.services.kinesis.model.GetRecordsResult; +import com.amazonaws.services.kinesis.model.GetShardIteratorRequest; +import com.amazonaws.services.kinesis.model.GetShardIteratorResult; +import com.amazonaws.services.kinesis.model.IncreaseStreamRetentionPeriodRequest; +import com.amazonaws.services.kinesis.model.IncreaseStreamRetentionPeriodResult; +import com.amazonaws.services.kinesis.model.ListStreamsRequest; +import com.amazonaws.services.kinesis.model.ListStreamsResult; +import com.amazonaws.services.kinesis.model.ListTagsForStreamRequest; +import com.amazonaws.services.kinesis.model.ListTagsForStreamResult; +import com.amazonaws.services.kinesis.model.MergeShardsRequest; +import com.amazonaws.services.kinesis.model.MergeShardsResult; +import com.amazonaws.services.kinesis.model.PutRecordRequest; +import com.amazonaws.services.kinesis.model.PutRecordResult; +import com.amazonaws.services.kinesis.model.PutRecordsRequest; +import com.amazonaws.services.kinesis.model.PutRecordsResult; +import com.amazonaws.services.kinesis.model.RemoveTagsFromStreamRequest; +import com.amazonaws.services.kinesis.model.RemoveTagsFromStreamResult; +import com.amazonaws.services.kinesis.model.SplitShardRequest; +import com.amazonaws.services.kinesis.model.SplitShardResult; +import com.amazonaws.services.kinesis.model.UpdateShardCountRequest; +import com.amazonaws.services.kinesis.model.UpdateShardCountResult; +import com.newrelic.api.agent.Trace; +import com.newrelic.api.agent.weaver.Weaver; + +import java.util.concurrent.Future; + +public class AmazonKinesisAsyncClient_Instrumentation { + + @Trace + public Future addTagsToStreamAsync(AddTagsToStreamRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future createStreamAsync(CreateStreamRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future decreaseStreamRetentionPeriodAsync( + DecreaseStreamRetentionPeriodRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future deleteStreamAsync(DeleteStreamRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future describeLimitsAsync(DescribeLimitsRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future describeStreamAsync(DescribeStreamRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future disableEnhancedMonitoringAsync(DisableEnhancedMonitoringRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future enableEnhancedMonitoringAsync(EnableEnhancedMonitoringRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future getRecordsAsync(GetRecordsRequest request, AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future getShardIteratorAsync(GetShardIteratorRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future increaseStreamRetentionPeriodAsync( + IncreaseStreamRetentionPeriodRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future listStreamsAsync(ListStreamsRequest request, AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future listTagsForStreamAsync(ListTagsForStreamRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future mergeShardsAsync(MergeShardsRequest request, AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future putRecordAsync(PutRecordRequest request, AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future putRecordsAsync(PutRecordsRequest request, AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future removeTagsFromStreamAsync(RemoveTagsFromStreamRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future splitShardAsync(SplitShardRequest request, AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } + + @Trace + public Future updateShardCountAsync(UpdateShardCountRequest request, + AsyncHandler_Instrumentation asyncHandler) { + KinesisUtil.setupToken(asyncHandler, request); + return Weaver.callOriginal(); + } +} diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java new file mode 100644 index 0000000000..671a058caf --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java @@ -0,0 +1,183 @@ +package com.amazonaws.services.kinesis; + +import com.agent.instrumentation.awsjavasdk2.services.lambda.KinesisUtil; +import com.amazonaws.services.kinesis.model.AddTagsToStreamRequest; +import com.amazonaws.services.kinesis.model.AddTagsToStreamResult; +import com.amazonaws.services.kinesis.model.CreateStreamRequest; +import com.amazonaws.services.kinesis.model.CreateStreamResult; +import com.amazonaws.services.kinesis.model.DecreaseStreamRetentionPeriodRequest; +import com.amazonaws.services.kinesis.model.DecreaseStreamRetentionPeriodResult; +import com.amazonaws.services.kinesis.model.DeleteStreamRequest; +import com.amazonaws.services.kinesis.model.DeleteStreamResult; +import com.amazonaws.services.kinesis.model.DescribeLimitsRequest; +import com.amazonaws.services.kinesis.model.DescribeLimitsResult; +import com.amazonaws.services.kinesis.model.DescribeStreamRequest; +import com.amazonaws.services.kinesis.model.DescribeStreamResult; +import com.amazonaws.services.kinesis.model.DisableEnhancedMonitoringRequest; +import com.amazonaws.services.kinesis.model.DisableEnhancedMonitoringResult; +import com.amazonaws.services.kinesis.model.EnableEnhancedMonitoringRequest; +import com.amazonaws.services.kinesis.model.EnableEnhancedMonitoringResult; +import com.amazonaws.services.kinesis.model.GetRecordsRequest; +import com.amazonaws.services.kinesis.model.GetRecordsResult; +import com.amazonaws.services.kinesis.model.GetShardIteratorRequest; +import com.amazonaws.services.kinesis.model.GetShardIteratorResult; +import com.amazonaws.services.kinesis.model.IncreaseStreamRetentionPeriodRequest; +import com.amazonaws.services.kinesis.model.IncreaseStreamRetentionPeriodResult; +import com.amazonaws.services.kinesis.model.ListStreamsRequest; +import com.amazonaws.services.kinesis.model.ListStreamsResult; +import com.amazonaws.services.kinesis.model.ListTagsForStreamRequest; +import com.amazonaws.services.kinesis.model.ListTagsForStreamResult; +import com.amazonaws.services.kinesis.model.MergeShardsRequest; +import com.amazonaws.services.kinesis.model.MergeShardsResult; +import com.amazonaws.services.kinesis.model.PutRecordRequest; +import com.amazonaws.services.kinesis.model.PutRecordResult; +import com.amazonaws.services.kinesis.model.PutRecordsRequest; +import com.amazonaws.services.kinesis.model.PutRecordsResult; +import com.amazonaws.services.kinesis.model.RemoveTagsFromStreamRequest; +import com.amazonaws.services.kinesis.model.RemoveTagsFromStreamResult; +import com.amazonaws.services.kinesis.model.SplitShardRequest; +import com.amazonaws.services.kinesis.model.SplitShardResult; +import com.amazonaws.services.kinesis.model.UpdateShardCountRequest; +import com.amazonaws.services.kinesis.model.UpdateShardCountResult; +import com.newrelic.api.agent.Trace; +import com.newrelic.api.agent.weaver.MatchType; +import com.newrelic.api.agent.weaver.Weave; +import com.newrelic.api.agent.weaver.Weaver; + +@Weave(originalName = "com.amazonaws.services.kinesis.AmazonKinesisClient", type = MatchType.ExactClass) +public class AmazonKinesisClient_Instrumentation { + + @Trace(leaf=true) + final AddTagsToStreamResult executeAddTagsToStream(AddTagsToStreamRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("addTagsToStream"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final CreateStreamResult executeCreateStream(CreateStreamRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("createStream"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final DecreaseStreamRetentionPeriodResult executeDecreaseStreamRetentionPeriod(DecreaseStreamRetentionPeriodRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("decreaseStreamRetentionPeriod"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final DeleteStreamResult executeDeleteStream(DeleteStreamRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("deleteStream"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final DescribeLimitsResult executeDescribeLimits(DescribeLimitsRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("describeLimits"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final DescribeStreamResult executeDescribeStream(DescribeStreamRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("describeStream"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final DisableEnhancedMonitoringResult executeDisableEnhancedMonitoring(DisableEnhancedMonitoringRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("disableEnhancedMonitoring"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final EnableEnhancedMonitoringResult executeEnableEnhancedMonitoring(EnableEnhancedMonitoringRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("enableEnhancedMonitoring"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final GetRecordsResult executeGetRecords(GetRecordsRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("getRecords"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final GetShardIteratorResult executeGetShardIterator(GetShardIteratorRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("getShardIterator"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final IncreaseStreamRetentionPeriodResult executeIncreaseStreamRetentionPeriod(IncreaseStreamRetentionPeriodRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("increaseStreamRetentionPeriod"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final ListStreamsResult executeListStreams(ListStreamsRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("listStreams"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final ListTagsForStreamResult executeListTagsForStream(ListTagsForStreamRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("listTagsForStream"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final MergeShardsResult executeMergeShards(MergeShardsRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("mergeShards"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final PutRecordResult executePutRecord(PutRecordRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("putRecord"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final PutRecordsResult executePutRecords(PutRecordsRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("putRecords"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final RemoveTagsFromStreamResult executeRemoveTagsFromStream(RemoveTagsFromStreamRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("removeTagsFromStream"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final SplitShardResult executeSplitShard(SplitShardRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("splitShard"); + return Weaver.callOriginal(); + } + + @Trace(leaf=true) + final UpdateShardCountResult executeUpdateShardCount(UpdateShardCountRequest request) { + KinesisUtil.linkAndExpireToken(request); + KinesisUtil.setTraceDetails("updateShardCount"); + return Weaver.callOriginal(); + } + +} diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/resources/dt_enabled.yml b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/resources/dt_enabled.yml new file mode 100644 index 0000000000..53b0968002 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/resources/dt_enabled.yml @@ -0,0 +1,5 @@ +common: &default_settings + distributed_tracing: + enabled: true + span_events: + enabled: true \ No newline at end of file diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java index 8740c94d0a..07bf5ac3bc 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java @@ -11,14 +11,16 @@ public class KinesisUtil { public static final String TRACE_CATEGORY = "Kinesis"; private KinesisUtil() {} - public static Segment startSegment(String kinesisOperation) { - return NewRelic.getAgent().getTransaction().startSegment(TRACE_CATEGORY, kinesisOperation); + public static Segment beginSegment(String kinesisOperation) { + Segment segment = NewRelic.getAgent().getTransaction().startSegment(TRACE_CATEGORY, kinesisOperation); + segment.reportAsExternal(createCloudParams()); + return segment; } public static void setTraceDetails(String kinesisOperation) { TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); tracedMethod.setMetricName(TRACE_CATEGORY, kinesisOperation); - tracedMethod.reportAsExternal(CloudParameters.provider(PLATFORM).build()); + tracedMethod.reportAsExternal(createCloudParams()); } public static CloudParameters createCloudParams() { // Todo: add arn to cloud parameters diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java index bfd485f172..ae1a0acd7c 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java @@ -69,50 +69,50 @@ class DefaultKinesisAsyncClient_Instrumentation { public CompletableFuture addTagsToStream(AddTagsToStreamRequest addTagsToStreamRequest) { - Segment segment = KinesisUtil.startSegment("addTagsToStream"); + Segment segment = KinesisUtil.beginSegment("addTagsToStream"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture createStream(CreateStreamRequest createStreamRequest) { - Segment segment = KinesisUtil.startSegment("createStream"); + Segment segment = KinesisUtil.beginSegment("createStream"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture decreaseStreamRetentionPeriod( DecreaseStreamRetentionPeriodRequest decreaseStreamRetentionPeriodRequest) { - Segment segment = KinesisUtil.startSegment("decreaseStreamRetentionPeriod"); + Segment segment = KinesisUtil.beginSegment("decreaseStreamRetentionPeriod"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture deleteStream(DeleteStreamRequest deleteStreamRequest) { - Segment segment = KinesisUtil.startSegment("deleteStream"); + Segment segment = KinesisUtil.beginSegment("deleteStream"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture deregisterStreamConsumer(DeregisterStreamConsumerRequest deregisterStreamConsumerRequest) { - Segment segment = KinesisUtil.startSegment("deregisterStreamConsumer"); + Segment segment = KinesisUtil.beginSegment("deregisterStreamConsumer"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture describeLimits(DescribeLimitsRequest describeLimitsRequest) { - Segment segment = KinesisUtil.startSegment("describeLimits"); + Segment segment = KinesisUtil.beginSegment("describeLimits"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture describeStream(DescribeStreamRequest describeStreamRequest) { - Segment segment = KinesisUtil.startSegment("describeStream"); + Segment segment = KinesisUtil.beginSegment("describeStream"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture describeStreamConsumer(DescribeStreamConsumerRequest describeStreamConsumerRequest) { - Segment segment = KinesisUtil.startSegment("describeStreamConsumer"); + Segment segment = KinesisUtil.beginSegment("describeStreamConsumer"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } @@ -124,115 +124,115 @@ public CompletableFuture describeStreamSummary(D } public CompletableFuture disableEnhancedMonitoring(DisableEnhancedMonitoringRequest disableEnhancedMonitoringRequest) { - Segment segment = KinesisUtil.startSegment("disableEnhancedMonitoring"); + Segment segment = KinesisUtil.beginSegment("disableEnhancedMonitoring"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture enableEnhancedMonitoring(EnableEnhancedMonitoringRequest enableEnhancedMonitoringRequest) { - Segment segment = KinesisUtil.startSegment("enableEnhancedMonitoring"); + Segment segment = KinesisUtil.beginSegment("enableEnhancedMonitoring"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture getRecords(GetRecordsRequest getRecordsRequest) { - Segment segment = KinesisUtil.startSegment("getRecords"); + Segment segment = KinesisUtil.beginSegment("getRecords"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture getShardIterator(GetShardIteratorRequest getShardIteratorRequest) { - Segment segment = KinesisUtil.startSegment("getShardIterator"); + Segment segment = KinesisUtil.beginSegment("getShardIterator"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture increaseStreamRetentionPeriod( IncreaseStreamRetentionPeriodRequest increaseStreamRetentionPeriodRequest) { - Segment segment = KinesisUtil.startSegment("increaseStreamRetentionPeriod"); + Segment segment = KinesisUtil.beginSegment("increaseStreamRetentionPeriod"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture listShards(ListShardsRequest listShardsRequest) { - Segment segment = KinesisUtil.startSegment("listShards"); + Segment segment = KinesisUtil.beginSegment("listShards"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture listStreamConsumers(ListStreamConsumersRequest listStreamConsumersRequest) { - Segment segment = KinesisUtil.startSegment("listStreamConsumers"); + Segment segment = KinesisUtil.beginSegment("listStreamConsumers"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture listStreams(ListStreamsRequest listStreamsRequest) { - Segment segment = KinesisUtil.startSegment("listStreams"); + Segment segment = KinesisUtil.beginSegment("listStreams"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture listTagsForStream(ListTagsForStreamRequest listTagsForStreamRequest) { - Segment segment = KinesisUtil.startSegment("listTagsForStream"); + Segment segment = KinesisUtil.beginSegment("listTagsForStream"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture mergeShards(MergeShardsRequest mergeShardsRequest) { - Segment segment = KinesisUtil.startSegment("mergeShards"); + Segment segment = KinesisUtil.beginSegment("mergeShards"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture putRecord(PutRecordRequest putRecordRequest) { - Segment segment = KinesisUtil.startSegment("putRecord"); + Segment segment = KinesisUtil.beginSegment("putRecord"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture putRecords(PutRecordsRequest putRecordsRequest) { - Segment segment = KinesisUtil.startSegment("putRecords"); + Segment segment = KinesisUtil.beginSegment("putRecords"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture registerStreamConsumer(RegisterStreamConsumerRequest registerStreamConsumerRequest) { - Segment segment = KinesisUtil.startSegment("registerStreamConsumer"); + Segment segment = KinesisUtil.beginSegment("registerStreamConsumer"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture removeTagsFromStream(RemoveTagsFromStreamRequest removeTagsFromStreamRequest) { - Segment segment = KinesisUtil.startSegment("removeTagsFromStream"); + Segment segment = KinesisUtil.beginSegment("removeTagsFromStream"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture splitShard(SplitShardRequest splitShardRequest) { - Segment segment = KinesisUtil.startSegment("splitShard"); + Segment segment = KinesisUtil.beginSegment("splitShard"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture startStreamEncryption(StartStreamEncryptionRequest startStreamEncryptionRequest) { - Segment segment = KinesisUtil.startSegment("startStreamEncryption"); + Segment segment = KinesisUtil.beginSegment("startStreamEncryption"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture stopStreamEncryption(StopStreamEncryptionRequest stopStreamEncryptionRequest) { - Segment segment = KinesisUtil.startSegment("stopStreamEncryption"); + Segment segment = KinesisUtil.beginSegment("stopStreamEncryption"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture subscribeToShard(SubscribeToShardRequest subscribeToShardRequest, SubscribeToShardResponseHandler asyncResponseHandler) { - Segment segment = KinesisUtil.startSegment("stopStreamEncryption"); + Segment segment = KinesisUtil.beginSegment("stopStreamEncryption"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } public CompletableFuture updateShardCount(UpdateShardCountRequest updateShardCountRequest) { - Segment segment = KinesisUtil.startSegment("updateShardCount"); + Segment segment = KinesisUtil.beginSegment("updateShardCount"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java index e1efa5d48f..02f27a02da 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java @@ -1,7 +1,6 @@ package software.amazon.awssdk.services.kinesis; import com.agent.instrumentation.awsjavasdk2.services.lambda.KinesisUtil; -import com.newrelic.api.agent.NewRelic; import com.newrelic.api.agent.Trace; import com.newrelic.api.agent.weaver.MatchType; import com.newrelic.api.agent.weaver.Weave; @@ -65,163 +64,163 @@ class DefaultKinesisClient_Instrumentation { @Trace(leaf=true) public AddTagsToStreamResponse addTagsToStream(AddTagsToStreamRequest addTagsToStreamRequest) { - KinesisUtil.setTraceDetails( "addTagsToStream"); + KinesisUtil.setTraceDetails("addTagsToStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public CreateStreamResponse createStream(CreateStreamRequest createStreamRequest) { - KinesisUtil.setTraceDetails( "createStream"); + KinesisUtil.setTraceDetails("createStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public DecreaseStreamRetentionPeriodResponse decreaseStreamRetentionPeriod(DecreaseStreamRetentionPeriodRequest decreaseStreamRetentionPeriodRequest) { - KinesisUtil.setTraceDetails( "decreaseStreamRetentionPeriod"); + KinesisUtil.setTraceDetails("decreaseStreamRetentionPeriod"); return Weaver.callOriginal(); } @Trace(leaf=true) public DeleteStreamResponse deleteStream(DeleteStreamRequest deleteStreamRequest) { - KinesisUtil.setTraceDetails( "deleteStream"); + KinesisUtil.setTraceDetails("deleteStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public DeregisterStreamConsumerResponse deregisterStreamConsumer(DeregisterStreamConsumerRequest deregisterStreamConsumerRequest) { - KinesisUtil.setTraceDetails( "deregisterStreamConsumer"); + KinesisUtil.setTraceDetails( "deregisterStreamConsumer"); return Weaver.callOriginal(); } @Trace(leaf=true) public DescribeLimitsResponse describeLimits(DescribeLimitsRequest describeLimitsRequest) { - KinesisUtil.setTraceDetails( "describeLimits"); + KinesisUtil.setTraceDetails( "describeLimits"); return Weaver.callOriginal(); } @Trace(leaf=true) public DescribeStreamResponse describeStream(DescribeStreamRequest describeStreamRequest) { - KinesisUtil.setTraceDetails( "describeStream"); + KinesisUtil.setTraceDetails( "describeStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public DescribeStreamConsumerResponse describeStreamConsumer(DescribeStreamConsumerRequest describeStreamConsumerRequest) { - KinesisUtil.setTraceDetails( "describeStreamConsumer"); + KinesisUtil.setTraceDetails( "describeStreamConsumer"); return Weaver.callOriginal(); } @Trace(leaf=true) public DescribeStreamSummaryResponse describeStreamSummary(DescribeStreamSummaryRequest describeStreamSummaryRequest) { - KinesisUtil.setTraceDetails( "describeStreamSummary"); + KinesisUtil.setTraceDetails( "describeStreamSummary"); return Weaver.callOriginal(); } @Trace(leaf=true) public DisableEnhancedMonitoringResponse disableEnhancedMonitoring(DisableEnhancedMonitoringRequest disableEnhancedMonitoringRequest) { - KinesisUtil.setTraceDetails( "disableEnhancedMonitoring"); + KinesisUtil.setTraceDetails( "disableEnhancedMonitoring"); return Weaver.callOriginal(); } @Trace(leaf=true) public EnableEnhancedMonitoringResponse enableEnhancedMonitoring(EnableEnhancedMonitoringRequest enableEnhancedMonitoringRequest) { - KinesisUtil.setTraceDetails( "enableEnhancedMonitoring"); + KinesisUtil.setTraceDetails( "enableEnhancedMonitoring"); return Weaver.callOriginal(); } @Trace(leaf=true) public GetRecordsResponse getRecords(GetRecordsRequest getRecordsRequest) { - KinesisUtil.setTraceDetails( "getRecords"); + KinesisUtil.setTraceDetails( "getRecords"); return Weaver.callOriginal(); } @Trace(leaf=true) public GetShardIteratorResponse getShardIterator(GetShardIteratorRequest getShardIteratorRequest) { - KinesisUtil.setTraceDetails( "getShardIterator"); + KinesisUtil.setTraceDetails( "getShardIterator"); return Weaver.callOriginal(); } @Trace(leaf=true) public IncreaseStreamRetentionPeriodResponse increaseStreamRetentionPeriod(IncreaseStreamRetentionPeriodRequest increaseStreamRetentionPeriodRequest) { - KinesisUtil.setTraceDetails( "increaseStreamRetentionPeriod"); + KinesisUtil.setTraceDetails( "increaseStreamRetentionPeriod"); return Weaver.callOriginal(); } @Trace(leaf=true) public ListShardsResponse listShards(ListShardsRequest listShardsRequest) { - KinesisUtil.setTraceDetails( "listShards"); + KinesisUtil.setTraceDetails( "listShards"); return Weaver.callOriginal(); } @Trace(leaf=true) public ListStreamConsumersResponse listStreamConsumers(ListStreamConsumersRequest listStreamConsumersRequest) { - KinesisUtil.setTraceDetails( "listStreamConsumers"); + KinesisUtil.setTraceDetails( "listStreamConsumers"); return Weaver.callOriginal(); } @Trace(leaf=true) public ListStreamsResponse listStreams(ListStreamsRequest listStreamsRequest) { - KinesisUtil.setTraceDetails( "listStreams"); + KinesisUtil.setTraceDetails( "listStreams"); return Weaver.callOriginal(); } @Trace(leaf=true) public ListTagsForStreamResponse listTagsForStream(ListTagsForStreamRequest listTagsForStreamRequest) { - KinesisUtil.setTraceDetails( "listTagsForStream"); + KinesisUtil.setTraceDetails( "listTagsForStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public MergeShardsResponse mergeShards(MergeShardsRequest mergeShardsRequest) { - KinesisUtil.setTraceDetails( "mergeShards"); + KinesisUtil.setTraceDetails( "mergeShards"); return Weaver.callOriginal(); } @Trace(leaf=true) public PutRecordResponse putRecord(PutRecordRequest putRecordRequest) { - KinesisUtil.setTraceDetails( "putRecord"); + KinesisUtil.setTraceDetails( "putRecord"); return Weaver.callOriginal(); } @Trace(leaf=true) public PutRecordsResponse putRecords(PutRecordsRequest putRecordsRequest) { - KinesisUtil.setTraceDetails( "putRecords"); + KinesisUtil.setTraceDetails( "putRecords"); return Weaver.callOriginal(); } @Trace(leaf=true) public RegisterStreamConsumerResponse registerStreamConsumer(RegisterStreamConsumerRequest registerStreamConsumerRequest) { - KinesisUtil.setTraceDetails( "registerStreamConsumer"); + KinesisUtil.setTraceDetails( "registerStreamConsumer"); return Weaver.callOriginal(); } @Trace(leaf=true) public RemoveTagsFromStreamResponse removeTagsFromStream(RemoveTagsFromStreamRequest removeTagsFromStreamRequest) { - KinesisUtil.setTraceDetails( "removeTagsFromStream"); + KinesisUtil.setTraceDetails( "removeTagsFromStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public SplitShardResponse splitShard(SplitShardRequest splitShardRequest) { - KinesisUtil.setTraceDetails( "splitShard"); + KinesisUtil.setTraceDetails( "splitShard"); return Weaver.callOriginal(); } @Trace(leaf=true) public StartStreamEncryptionResponse startStreamEncryption(StartStreamEncryptionRequest startStreamEncryptionRequest) { - KinesisUtil.setTraceDetails( "startStreamEncryption"); + KinesisUtil.setTraceDetails( "startStreamEncryption"); return Weaver.callOriginal(); } @Trace(leaf=true) public StopStreamEncryptionResponse stopStreamEncryption(StopStreamEncryptionRequest stopStreamEncryptionRequest) { - KinesisUtil.setTraceDetails( "stopStreamEncryption"); + KinesisUtil.setTraceDetails( "stopStreamEncryption"); return Weaver.callOriginal(); } @Trace(leaf=true) public UpdateShardCountResponse updateShardCount(UpdateShardCountRequest updateShardCountRequest) { - KinesisUtil.setTraceDetails( "updateShardCount"); + KinesisUtil.setTraceDetails( "updateShardCount"); return Weaver.callOriginal(); } diff --git a/settings.gradle b/settings.gradle index 7fc787ca1d..0bcbfdb905 100644 --- a/settings.gradle +++ b/settings.gradle @@ -81,6 +81,7 @@ include 'instrumentation:aws-java-sdk-sns-1.11.12' include 'instrumentation:aws-java-sdk-sqs-2.1.0' include 'instrumentation:aws-java-sdk-dynamodb-1.11.106' include 'instrumentation:aws-java-sdk-dynamodb-2.15.34' +include 'instrumentation:aws-java-sdk-kinesis-1.11.106' include 'instrumentation:aws-java-sdk-kinesis-2.0.6' include 'instrumentation:aws-java-sdk-sns-2.0' include 'instrumentation:aws-wrap-0.7.0' From b3e8a77fb31677f87cd41777c6fc6453e57b3f3b Mon Sep 17 00:00:00 2001 From: Oren Ben-Meir Date: Tue, 27 Aug 2024 16:18:23 -0400 Subject: [PATCH 3/8] Setup v2 kinesis client test --- .../services/lambda/KinesisUtil.java | 19 +- ...zonWebServiceRequest_Instrumentation.java} | 4 +- .../AsyncHandler_Instrumentation.java | 2 +- .../AmazonKinesisClient_Instrumentation.java | 1 + .../DefaultKinesisClient_Instrumentation.java | 46 ++--- .../kinesis/DefaultKinesisClientTest.java | 177 ++++++++++++++++++ 6 files changed, 215 insertions(+), 34 deletions(-) rename instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/{AmazonWebServiceRequest.java => AmazonWebServiceRequest_Instrumentation.java} (52%) create mode 100644 instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClientTest.java diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java index fc99d130cf..ae72f54830 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java @@ -1,12 +1,11 @@ package com.agent.instrumentation.awsjavasdk2.services.lambda; import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceRequest_Instrumentation; import com.amazonaws.handlers.AsyncHandler_Instrumentation; import com.newrelic.agent.bridge.AgentBridge; import com.newrelic.api.agent.CloudParameters; import com.newrelic.api.agent.NewRelic; -import com.newrelic.api.agent.Segment; -import com.newrelic.api.agent.Token; import com.newrelic.api.agent.TracedMethod; public class KinesisUtil { @@ -15,21 +14,25 @@ public class KinesisUtil { public static final String TRACE_CATEGORY = "Kinesis"; private KinesisUtil() {} - public static void setupToken(AsyncHandler_Instrumentation asyncHandler, AmazonWebServiceRequest request) { + public static void setupToken(AsyncHandler_Instrumentation asyncHandler, AmazonWebServiceRequest amazonWebServiceRequest) { if (AgentBridge.getAgent().getTransaction(false) != null) { if (asyncHandler != null) { asyncHandler.token = NewRelic.getAgent().getTransaction().getToken(); } - if (request != null) { + if (amazonWebServiceRequest != null && amazonWebServiceRequest instanceof AmazonWebServiceRequest_Instrumentation) { + AmazonWebServiceRequest_Instrumentation request = (AmazonWebServiceRequest_Instrumentation)amazonWebServiceRequest; request.token = NewRelic.getAgent().getTransaction().getToken(); } } } - public static void linkAndExpireToken(AmazonWebServiceRequest request) { - if (request != null && request.token != null) { - request.token.linkAndExpire(); - request.token = null; + public static void linkAndExpireToken(AmazonWebServiceRequest amazonWebServiceRequest) { + if (amazonWebServiceRequest != null && amazonWebServiceRequest instanceof AmazonWebServiceRequest_Instrumentation) { + AmazonWebServiceRequest_Instrumentation request = (AmazonWebServiceRequest_Instrumentation)amazonWebServiceRequest; + if (request.token != null) { + request.token.linkAndExpire(); + request.token = null; + } } } diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest_Instrumentation.java similarity index 52% rename from instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest.java rename to instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest_Instrumentation.java index 63202e7543..43065e13f4 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest_Instrumentation.java @@ -4,8 +4,8 @@ import com.newrelic.api.agent.weaver.NewField; import com.newrelic.api.agent.weaver.Weave; -@Weave(originalName = "com.amazonaws.AmazonWebServiceRequest") -public class AmazonWebServiceRequest { +@Weave(originalName = "com.amazonaws.AmazonWebServiceRequest_Instrumentation") +public abstract class AmazonWebServiceRequest_Instrumentation extends AmazonWebServiceRequest { @NewField public Token token; } diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java index 9d826c5ba0..d5e44f9b26 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java @@ -8,7 +8,7 @@ package com.amazonaws.handlers; import com.amazonaws.AmazonWebServiceRequest; -import com.newrelic.api.agent.Segment; +import com.amazonaws.AmazonWebServiceRequest_Instrumentation; import com.newrelic.api.agent.Token; import com.newrelic.api.agent.Trace; import com.newrelic.api.agent.weaver.MatchType; diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java index 671a058caf..4c404b3f87 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java @@ -1,6 +1,7 @@ package com.amazonaws.services.kinesis; import com.agent.instrumentation.awsjavasdk2.services.lambda.KinesisUtil; +import com.amazonaws.AmazonWebServiceRequest; import com.amazonaws.services.kinesis.model.AddTagsToStreamRequest; import com.amazonaws.services.kinesis.model.AddTagsToStreamResult; import com.amazonaws.services.kinesis.model.CreateStreamRequest; diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java index 02f27a02da..a9223fcccd 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java @@ -88,139 +88,139 @@ public DeleteStreamResponse deleteStream(DeleteStreamRequest deleteStreamRequest @Trace(leaf=true) public DeregisterStreamConsumerResponse deregisterStreamConsumer(DeregisterStreamConsumerRequest deregisterStreamConsumerRequest) { - KinesisUtil.setTraceDetails( "deregisterStreamConsumer"); + KinesisUtil.setTraceDetails("deregisterStreamConsumer"); return Weaver.callOriginal(); } @Trace(leaf=true) public DescribeLimitsResponse describeLimits(DescribeLimitsRequest describeLimitsRequest) { - KinesisUtil.setTraceDetails( "describeLimits"); + KinesisUtil.setTraceDetails("describeLimits"); return Weaver.callOriginal(); } @Trace(leaf=true) public DescribeStreamResponse describeStream(DescribeStreamRequest describeStreamRequest) { - KinesisUtil.setTraceDetails( "describeStream"); + KinesisUtil.setTraceDetails("describeStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public DescribeStreamConsumerResponse describeStreamConsumer(DescribeStreamConsumerRequest describeStreamConsumerRequest) { - KinesisUtil.setTraceDetails( "describeStreamConsumer"); + KinesisUtil.setTraceDetails("describeStreamConsumer"); return Weaver.callOriginal(); } @Trace(leaf=true) public DescribeStreamSummaryResponse describeStreamSummary(DescribeStreamSummaryRequest describeStreamSummaryRequest) { - KinesisUtil.setTraceDetails( "describeStreamSummary"); + KinesisUtil.setTraceDetails("describeStreamSummary"); return Weaver.callOriginal(); } @Trace(leaf=true) public DisableEnhancedMonitoringResponse disableEnhancedMonitoring(DisableEnhancedMonitoringRequest disableEnhancedMonitoringRequest) { - KinesisUtil.setTraceDetails( "disableEnhancedMonitoring"); + KinesisUtil.setTraceDetails("disableEnhancedMonitoring"); return Weaver.callOriginal(); } @Trace(leaf=true) public EnableEnhancedMonitoringResponse enableEnhancedMonitoring(EnableEnhancedMonitoringRequest enableEnhancedMonitoringRequest) { - KinesisUtil.setTraceDetails( "enableEnhancedMonitoring"); + KinesisUtil.setTraceDetails("enableEnhancedMonitoring"); return Weaver.callOriginal(); } @Trace(leaf=true) public GetRecordsResponse getRecords(GetRecordsRequest getRecordsRequest) { - KinesisUtil.setTraceDetails( "getRecords"); + KinesisUtil.setTraceDetails("getRecords"); return Weaver.callOriginal(); } @Trace(leaf=true) public GetShardIteratorResponse getShardIterator(GetShardIteratorRequest getShardIteratorRequest) { - KinesisUtil.setTraceDetails( "getShardIterator"); + KinesisUtil.setTraceDetails("getShardIterator"); return Weaver.callOriginal(); } @Trace(leaf=true) public IncreaseStreamRetentionPeriodResponse increaseStreamRetentionPeriod(IncreaseStreamRetentionPeriodRequest increaseStreamRetentionPeriodRequest) { - KinesisUtil.setTraceDetails( "increaseStreamRetentionPeriod"); + KinesisUtil.setTraceDetails("increaseStreamRetentionPeriod"); return Weaver.callOriginal(); } @Trace(leaf=true) public ListShardsResponse listShards(ListShardsRequest listShardsRequest) { - KinesisUtil.setTraceDetails( "listShards"); + KinesisUtil.setTraceDetails("listShards"); return Weaver.callOriginal(); } @Trace(leaf=true) public ListStreamConsumersResponse listStreamConsumers(ListStreamConsumersRequest listStreamConsumersRequest) { - KinesisUtil.setTraceDetails( "listStreamConsumers"); + KinesisUtil.setTraceDetails("listStreamConsumers"); return Weaver.callOriginal(); } @Trace(leaf=true) public ListStreamsResponse listStreams(ListStreamsRequest listStreamsRequest) { - KinesisUtil.setTraceDetails( "listStreams"); + KinesisUtil.setTraceDetails("listStreams"); return Weaver.callOriginal(); } @Trace(leaf=true) public ListTagsForStreamResponse listTagsForStream(ListTagsForStreamRequest listTagsForStreamRequest) { - KinesisUtil.setTraceDetails( "listTagsForStream"); + KinesisUtil.setTraceDetails("listTagsForStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public MergeShardsResponse mergeShards(MergeShardsRequest mergeShardsRequest) { - KinesisUtil.setTraceDetails( "mergeShards"); + KinesisUtil.setTraceDetails("mergeShards"); return Weaver.callOriginal(); } @Trace(leaf=true) public PutRecordResponse putRecord(PutRecordRequest putRecordRequest) { - KinesisUtil.setTraceDetails( "putRecord"); + KinesisUtil.setTraceDetails("putRecord"); return Weaver.callOriginal(); } @Trace(leaf=true) public PutRecordsResponse putRecords(PutRecordsRequest putRecordsRequest) { - KinesisUtil.setTraceDetails( "putRecords"); + KinesisUtil.setTraceDetails("putRecords"); return Weaver.callOriginal(); } @Trace(leaf=true) public RegisterStreamConsumerResponse registerStreamConsumer(RegisterStreamConsumerRequest registerStreamConsumerRequest) { - KinesisUtil.setTraceDetails( "registerStreamConsumer"); + KinesisUtil.setTraceDetails("registerStreamConsumer"); return Weaver.callOriginal(); } @Trace(leaf=true) public RemoveTagsFromStreamResponse removeTagsFromStream(RemoveTagsFromStreamRequest removeTagsFromStreamRequest) { - KinesisUtil.setTraceDetails( "removeTagsFromStream"); + KinesisUtil.setTraceDetails("removeTagsFromStream"); return Weaver.callOriginal(); } @Trace(leaf=true) public SplitShardResponse splitShard(SplitShardRequest splitShardRequest) { - KinesisUtil.setTraceDetails( "splitShard"); + KinesisUtil.setTraceDetails("splitShard"); return Weaver.callOriginal(); } @Trace(leaf=true) public StartStreamEncryptionResponse startStreamEncryption(StartStreamEncryptionRequest startStreamEncryptionRequest) { - KinesisUtil.setTraceDetails( "startStreamEncryption"); + KinesisUtil.setTraceDetails("startStreamEncryption"); return Weaver.callOriginal(); } @Trace(leaf=true) public StopStreamEncryptionResponse stopStreamEncryption(StopStreamEncryptionRequest stopStreamEncryptionRequest) { - KinesisUtil.setTraceDetails( "stopStreamEncryption"); + KinesisUtil.setTraceDetails("stopStreamEncryption"); return Weaver.callOriginal(); } @Trace(leaf=true) public UpdateShardCountResponse updateShardCount(UpdateShardCountRequest updateShardCountRequest) { - KinesisUtil.setTraceDetails( "updateShardCount"); + KinesisUtil.setTraceDetails("updateShardCount"); return Weaver.callOriginal(); } diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClientTest.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClientTest.java new file mode 100644 index 0000000000..25ec9c80b7 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClientTest.java @@ -0,0 +1,177 @@ +package software.amazon.awssdk.services.kinesis; + +import com.newrelic.agent.introspec.InstrumentationTestConfig; +import com.newrelic.agent.introspec.InstrumentationTestRunner; +import com.newrelic.agent.introspec.Introspector; +import com.newrelic.agent.introspec.SpanEvent; +import com.newrelic.agent.introspec.TraceSegment; +import com.newrelic.agent.introspec.TransactionTrace; +import com.newrelic.api.agent.Trace; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mockito; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.http.AbortableInputStream; +import software.amazon.awssdk.http.ExecutableHttpRequest; +import software.amazon.awssdk.http.HttpExecuteRequest; +import software.amazon.awssdk.http.HttpExecuteResponse; +import software.amazon.awssdk.http.SdkHttpClient; +import software.amazon.awssdk.http.SdkHttpFullResponse; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.kinesis.model.AddTagsToStreamRequest; +import software.amazon.awssdk.services.kinesis.model.CreateStreamRequest; +import software.amazon.awssdk.services.kinesis.model.DecreaseStreamRetentionPeriodRequest; +import software.amazon.awssdk.services.kinesis.model.DeleteStreamRequest; +import software.amazon.awssdk.utils.StringInputStream; + +import java.io.IOException; +import java.util.Collection; +import java.util.List; +import java.util.Optional; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@RunWith(InstrumentationTestRunner.class) +@InstrumentationTestConfig(includePrefixes = {"software.amazon.awssdk"}, configName = "dt_enabled.yml") +public class DefaultKinesisClientTest { + + public KinesisClient kinesisClient; + public HttpExecuteResponse response; + + @Before + public void setup() { + MockHttpClient mockHttpClient = new MockHttpClient(); + response = mockHttpClient.getResponse(); + kinesisClient = KinesisClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(new CredProvider()) + .region(Region.US_EAST_1) + .build(); + } + + @Test + public void testAddTagsToStream() { + txnAddTagsToStream(); + String txnName = "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txnAddTagsToStream"; + String traceName = "Kinesis/addTagsToStream"; + assertKinesisTrace(txnName, traceName); + } + + @Trace(dispatcher = true) + private void txnAddTagsToStream() { + kinesisClient.addTagsToStream(AddTagsToStreamRequest.builder().build()); + } + + @Test + public void testCreateStream() { + txnCreateStream(); + String txnName = "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txnCreateStream"; + String traceName = "Kinesis/createStream"; + assertKinesisTrace(txnName, traceName); + } + + @Trace(dispatcher = true) + private void txnCreateStream() { + kinesisClient.createStream(CreateStreamRequest.builder().build()); + } + + @Test + public void testDecreaseStreamRetentionPeriod() { + txnDecreaseStreamRetentionPeriod(); + String txnName = "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txnDecreaseStreamRetentionPeriod"; + String traceName = "Kinesis/decreaseStreamRetentionPeriod"; + assertKinesisTrace(txnName, traceName); + } + + @Trace(dispatcher = true) + private void txnDecreaseStreamRetentionPeriod() { + kinesisClient.decreaseStreamRetentionPeriod(DecreaseStreamRetentionPeriodRequest.builder().build()); + } + + @Test + public void testDeleteStream() { + txnDeleteStream(); + String txnName = "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txnDeleteStream"; + String traceName = "Kinesis/deleteStream"; + assertKinesisTrace(txnName, traceName); + } + + @Trace(dispatcher = true) + private void txnDeleteStream() { + kinesisClient.deleteStream(DeleteStreamRequest.builder().build()); + } + + private void assertKinesisTrace(String txnName, String traceName) { + Introspector introspector = InstrumentationTestRunner.getIntrospector(); + SpanEvent kinesisSpan = introspector.getSpanEvents().stream() + .filter(span -> traceName.equals(span.getName())) + .findFirst().orElse(null); + assertNotNull(kinesisSpan); + assertEquals("aws_kinesis_data_streams", kinesisSpan.getAgentAttributes().get("cloud.platform")); + Collection transactionTraces = introspector.getTransactionTracesForTransaction( + txnName); + TransactionTrace transactionTrace = transactionTraces.iterator().next(); + List children = transactionTrace.getInitialTraceSegment().getChildren(); + assertEquals(1, children.size()); + TraceSegment trace = children.get(0); + assertEquals(traceName, trace.getName()); + assertEquals("aws_kinesis_data_streams", trace.getTracerAttributes().get("cloud.platform")); + } + + private static class MockHttpClient implements SdkHttpClient { + private final ExecutableHttpRequest executableMock; + private final HttpExecuteResponse response; + private final SdkHttpFullResponse httpResponse; + + public MockHttpClient() { + executableMock = mock(ExecutableHttpRequest.class); + response = mock(HttpExecuteResponse.class, Mockito.RETURNS_DEEP_STUBS); + httpResponse = mock(SdkHttpFullResponse.class, Mockito.RETURNS_DEEP_STUBS); + when(response.httpResponse()).thenReturn(httpResponse); + when(httpResponse.toBuilder().content(any()).build()).thenReturn(httpResponse); + when(httpResponse.isSuccessful()).thenReturn(true); + AbortableInputStream inputStream = AbortableInputStream.create(new StringInputStream("42")); + when(httpResponse.content()).thenReturn(Optional.of(inputStream)); + try { + when(executableMock.call()).thenReturn(response); + + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Override + public void close() { + } + + @Override + public ExecutableHttpRequest prepareRequest(HttpExecuteRequest httpExecuteRequest) { + return executableMock; + } + + @Override + public String clientName() { + return "MockHttpClient"; + } + + public HttpExecuteResponse getResponse() { + return response; + } + } + + private static class CredProvider implements AwsCredentialsProvider { + @Override + public AwsCredentials resolveCredentials() { + AwsCredentials credentials = mock(AwsCredentials.class); + when(credentials.accessKeyId()).thenReturn("accessKeyId"); + when(credentials.secretAccessKey()).thenReturn("secretAccessKey"); + return credentials; + } + } +} From f1afcb6a072666b800d505da62134740f6b3e2d6 Mon Sep 17 00:00:00 2001 From: Oren Ben-Meir Date: Tue, 27 Aug 2024 17:02:46 -0400 Subject: [PATCH 4/8] Add more test cases for kinesis sync client --- .../kinesis/DefaultKinesisClientTest.java | 208 ++++++++++++++---- 1 file changed, 171 insertions(+), 37 deletions(-) diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClientTest.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClientTest.java index 25ec9c80b7..d70b4274bd 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClientTest.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClientTest.java @@ -24,6 +24,28 @@ import software.amazon.awssdk.services.kinesis.model.CreateStreamRequest; import software.amazon.awssdk.services.kinesis.model.DecreaseStreamRetentionPeriodRequest; import software.amazon.awssdk.services.kinesis.model.DeleteStreamRequest; +import software.amazon.awssdk.services.kinesis.model.DeregisterStreamConsumerRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeLimitsRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamConsumerRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest; +import software.amazon.awssdk.services.kinesis.model.DisableEnhancedMonitoringRequest; +import software.amazon.awssdk.services.kinesis.model.EnableEnhancedMonitoringRequest; +import software.amazon.awssdk.services.kinesis.model.GetRecordsRequest; +import software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest; +import software.amazon.awssdk.services.kinesis.model.IncreaseStreamRetentionPeriodRequest; +import software.amazon.awssdk.services.kinesis.model.ListShardsRequest; +import software.amazon.awssdk.services.kinesis.model.ListStreamsRequest; +import software.amazon.awssdk.services.kinesis.model.ListTagsForStreamRequest; +import software.amazon.awssdk.services.kinesis.model.MergeShardsRequest; +import software.amazon.awssdk.services.kinesis.model.PutRecordRequest; +import software.amazon.awssdk.services.kinesis.model.PutRecordsRequest; +import software.amazon.awssdk.services.kinesis.model.RegisterStreamConsumerRequest; +import software.amazon.awssdk.services.kinesis.model.RemoveTagsFromStreamRequest; +import software.amazon.awssdk.services.kinesis.model.SplitShardRequest; +import software.amazon.awssdk.services.kinesis.model.StartStreamEncryptionRequest; +import software.amazon.awssdk.services.kinesis.model.StopStreamEncryptionRequest; +import software.amazon.awssdk.services.kinesis.model.UpdateShardCountRequest; import software.amazon.awssdk.utils.StringInputStream; import java.io.IOException; @@ -57,65 +79,177 @@ public void setup() { @Test public void testAddTagsToStream() { - txnAddTagsToStream(); - String txnName = "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txnAddTagsToStream"; - String traceName = "Kinesis/addTagsToStream"; - assertKinesisTrace(txnName, traceName); + txn(() -> kinesisClient.addTagsToStream(AddTagsToStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/addTagsToStream", true); } - @Trace(dispatcher = true) - private void txnAddTagsToStream() { - kinesisClient.addTagsToStream(AddTagsToStreamRequest.builder().build()); + @Test + public void testCreateStream() { + txn(() -> kinesisClient.createStream(CreateStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/createStream", false); } @Test - public void testCreateStream() { - txnCreateStream(); - String txnName = "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txnCreateStream"; - String traceName = "Kinesis/createStream"; - assertKinesisTrace(txnName, traceName); + public void testDecreaseStreamRetentionPeriod() { + txn(() -> kinesisClient.decreaseStreamRetentionPeriod(DecreaseStreamRetentionPeriodRequest.builder().build())); + assertKinesisTrace("Kinesis/decreaseStreamRetentionPeriod", false); } - @Trace(dispatcher = true) - private void txnCreateStream() { - kinesisClient.createStream(CreateStreamRequest.builder().build()); + @Test + public void testDeleteStream() { + txn(() -> kinesisClient.deleteStream(DeleteStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/deleteStream", false); } @Test - public void testDecreaseStreamRetentionPeriod() { - txnDecreaseStreamRetentionPeriod(); - String txnName = "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txnDecreaseStreamRetentionPeriod"; - String traceName = "Kinesis/decreaseStreamRetentionPeriod"; - assertKinesisTrace(txnName, traceName); + public void testDeregisterStreamConsumer() { + txn(() -> kinesisClient.deregisterStreamConsumer(DeregisterStreamConsumerRequest.builder().build())); + assertKinesisTrace("Kinesis/deregisterStreamConsumer", false); } - @Trace(dispatcher = true) - private void txnDecreaseStreamRetentionPeriod() { - kinesisClient.decreaseStreamRetentionPeriod(DecreaseStreamRetentionPeriodRequest.builder().build()); + @Test + public void testDescribeLimits() { + txn(() -> kinesisClient.describeLimits(DescribeLimitsRequest.builder().build())); + assertKinesisTrace("Kinesis/describeLimits", false); } @Test - public void testDeleteStream() { - txnDeleteStream(); - String txnName = "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txnDeleteStream"; - String traceName = "Kinesis/deleteStream"; - assertKinesisTrace(txnName, traceName); + public void testDescribeStream() { + txn(() -> kinesisClient.describeStream(DescribeStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/describeStream", false); + } + + @Test + public void testDescribeStreamConsumer() { + txn(() -> kinesisClient.describeStreamConsumer(DescribeStreamConsumerRequest.builder().build())); + assertKinesisTrace("Kinesis/describeStreamConsumer", false); + } + + @Test + public void testDescribeStreamSummary() { + txn(() -> kinesisClient.describeStreamSummary(DescribeStreamSummaryRequest.builder().build())); + assertKinesisTrace("Kinesis/describeStreamSummary", false); + } + + @Test + public void DisableEnhancedMonitoring() { + txn(() -> kinesisClient.disableEnhancedMonitoring(DisableEnhancedMonitoringRequest.builder().build())); + assertKinesisTrace("Kinesis/disableEnhancedMonitoring", false); + } + + @Test + public void testEnableEnhancedMonitoring() { + txn(() -> kinesisClient.enableEnhancedMonitoring(EnableEnhancedMonitoringRequest.builder().build())); + assertKinesisTrace("Kinesis/enableEnhancedMonitoring", false); + } + + @Test + public void testGetRecords() { + txn(() -> kinesisClient.getRecords(GetRecordsRequest.builder().build())); + assertKinesisTrace("Kinesis/getRecords", false); + } + + @Test + public void testGetShardIterator() { + txn(() -> kinesisClient.getShardIterator(GetShardIteratorRequest.builder().build())); + assertKinesisTrace("Kinesis/getShardIterator", false); + } + + @Test + public void testIncreaseStreamRetentionPeriod() { + txn(() -> kinesisClient.increaseStreamRetentionPeriod(IncreaseStreamRetentionPeriodRequest.builder().build())); + assertKinesisTrace("Kinesis/increaseStreamRetentionPeriod", false); + } + + @Test + public void testListShards() { + txn(() -> kinesisClient.listShards(ListShardsRequest.builder().build())); + assertKinesisTrace("Kinesis/listShards", false); + } + + @Test + public void testListStreams() { + txn(() -> kinesisClient.listStreams(ListStreamsRequest.builder().build())); + assertKinesisTrace("Kinesis/listStreams", false); + } + + @Test + public void testListTagsForStream() { + txn(() -> kinesisClient.listTagsForStream(ListTagsForStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/listTagsForStream", false); + } + + @Test + public void testMergeShards() { + txn(() -> kinesisClient.mergeShards(MergeShardsRequest.builder().build())); + assertKinesisTrace("Kinesis/mergeShards", false); + } + + @Test + public void testPutRecord() { + txn(() -> kinesisClient.putRecord(PutRecordRequest.builder().build())); + assertKinesisTrace("Kinesis/putRecord", false); + } + + @Test + public void testPutRecords() { + txn(() -> kinesisClient.putRecords(PutRecordsRequest.builder().build())); + assertKinesisTrace("Kinesis/putRecords", false); + } + + @Test + public void testRegisterStreamConsumer() { + txn(() -> kinesisClient.registerStreamConsumer(RegisterStreamConsumerRequest.builder().build())); + assertKinesisTrace("Kinesis/registerStreamConsumer", false); + } + + @Test + public void testRemoveTagsFromStream() { + txn(() -> kinesisClient.removeTagsFromStream(RemoveTagsFromStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/removeTagsFromStream", false); + } + + @Test + public void testSplitShard() { + txn(() -> kinesisClient.splitShard(SplitShardRequest.builder().build())); + assertKinesisTrace("Kinesis/splitShard", false); + } + + @Test + public void testStartStreamEncryption() { + txn(() -> kinesisClient.startStreamEncryption(StartStreamEncryptionRequest.builder().build())); + assertKinesisTrace("Kinesis/startStreamEncryption", false); + } + + @Test + public void testStopStreamEncryption() { + txn(() -> kinesisClient.stopStreamEncryption(StopStreamEncryptionRequest.builder().build())); + assertKinesisTrace("Kinesis/stopStreamEncryption", false); + } + + @Test + public void testUpdateShardCount() { + txn(() -> kinesisClient.updateShardCount(UpdateShardCountRequest.builder().build())); + assertKinesisTrace("Kinesis/updateShardCount", false); } @Trace(dispatcher = true) - private void txnDeleteStream() { - kinesisClient.deleteStream(DeleteStreamRequest.builder().build()); + private void txn(Runnable runnable) { + runnable.run(); } - private void assertKinesisTrace(String txnName, String traceName) { + private void assertKinesisTrace(String traceName, boolean assertSpan) { Introspector introspector = InstrumentationTestRunner.getIntrospector(); - SpanEvent kinesisSpan = introspector.getSpanEvents().stream() - .filter(span -> traceName.equals(span.getName())) - .findFirst().orElse(null); - assertNotNull(kinesisSpan); - assertEquals("aws_kinesis_data_streams", kinesisSpan.getAgentAttributes().get("cloud.platform")); + if (assertSpan) { + // Span events fail to be generated when enough transactions are done in succession + SpanEvent kinesisSpan = introspector.getSpanEvents().stream() + .filter(span -> traceName.equals(span.getName())) + .findFirst().orElse(null); + assertNotNull(kinesisSpan); + assertEquals("aws_kinesis_data_streams", kinesisSpan.getAgentAttributes().get("cloud.platform")); + } Collection transactionTraces = introspector.getTransactionTracesForTransaction( - txnName); + "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisClientTest/txn"); TransactionTrace transactionTrace = transactionTraces.iterator().next(); List children = transactionTrace.getInitialTraceSegment().getChildren(); assertEquals(1, children.size()); From 9c3e160f08523cd1631207ddea12fa0be7923dc9 Mon Sep 17 00:00:00 2001 From: Oren Ben-Meir Date: Tue, 27 Aug 2024 17:17:46 -0400 Subject: [PATCH 5/8] Add kinesis async client test for v2 --- ...ultKinesisAsyncClient_Instrumentation.java | 2 +- .../DefaultKinesisAsyncClientTest.java | 315 ++++++++++++++++++ 2 files changed, 316 insertions(+), 1 deletion(-) create mode 100644 instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClientTest.java diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java index ae1a0acd7c..e0d6a53873 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java @@ -118,7 +118,7 @@ public CompletableFuture describeStreamConsumer( } public CompletableFuture describeStreamSummary(DescribeStreamSummaryRequest describeStreamSummaryRequest) { - Segment segment = NewRelic.getAgent().getTransaction().startSegment("KinesisAsyncClient", "describeStreamSummary"); + Segment segment = KinesisUtil.beginSegment( "describeStreamSummary"); CompletableFuture response = Weaver.callOriginal(); return new SegmentHandler<>(response, segment, Weaver.getImplementationTitle()).newSegmentCompletionStage(); } diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClientTest.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClientTest.java new file mode 100644 index 0000000000..e64542b05f --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/test/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClientTest.java @@ -0,0 +1,315 @@ +package software.amazon.awssdk.services.kinesis; + +import com.newrelic.agent.introspec.InstrumentationTestConfig; +import com.newrelic.agent.introspec.InstrumentationTestRunner; +import com.newrelic.agent.introspec.Introspector; +import com.newrelic.agent.introspec.SpanEvent; +import com.newrelic.agent.introspec.TraceSegment; +import com.newrelic.agent.introspec.TransactionTrace; +import com.newrelic.api.agent.Trace; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mockito; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.core.async.EmptyPublisher; +import software.amazon.awssdk.http.AbortableInputStream; +import software.amazon.awssdk.http.ExecutableHttpRequest; +import software.amazon.awssdk.http.HttpExecuteResponse; +import software.amazon.awssdk.http.SdkHttpFullResponse; +import software.amazon.awssdk.http.async.AsyncExecuteRequest; +import software.amazon.awssdk.http.async.SdkAsyncHttpClient; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.kinesis.model.AddTagsToStreamRequest; +import software.amazon.awssdk.services.kinesis.model.CreateStreamRequest; +import software.amazon.awssdk.services.kinesis.model.DecreaseStreamRetentionPeriodRequest; +import software.amazon.awssdk.services.kinesis.model.DeleteStreamRequest; +import software.amazon.awssdk.services.kinesis.model.DeregisterStreamConsumerRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeLimitsRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamConsumerRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest; +import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest; +import software.amazon.awssdk.services.kinesis.model.DisableEnhancedMonitoringRequest; +import software.amazon.awssdk.services.kinesis.model.EnableEnhancedMonitoringRequest; +import software.amazon.awssdk.services.kinesis.model.GetRecordsRequest; +import software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest; +import software.amazon.awssdk.services.kinesis.model.IncreaseStreamRetentionPeriodRequest; +import software.amazon.awssdk.services.kinesis.model.ListShardsRequest; +import software.amazon.awssdk.services.kinesis.model.ListStreamsRequest; +import software.amazon.awssdk.services.kinesis.model.ListTagsForStreamRequest; +import software.amazon.awssdk.services.kinesis.model.MergeShardsRequest; +import software.amazon.awssdk.services.kinesis.model.PutRecordRequest; +import software.amazon.awssdk.services.kinesis.model.PutRecordsRequest; +import software.amazon.awssdk.services.kinesis.model.RegisterStreamConsumerRequest; +import software.amazon.awssdk.services.kinesis.model.RemoveTagsFromStreamRequest; +import software.amazon.awssdk.services.kinesis.model.SplitShardRequest; +import software.amazon.awssdk.services.kinesis.model.StartStreamEncryptionRequest; +import software.amazon.awssdk.services.kinesis.model.StopStreamEncryptionRequest; +import software.amazon.awssdk.services.kinesis.model.UpdateShardCountRequest; +import software.amazon.awssdk.utils.StringInputStream; + +import java.io.IOException; +import java.util.Collection; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.function.Supplier; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@RunWith(InstrumentationTestRunner.class) +@InstrumentationTestConfig(includePrefixes = {"software.amazon.awssdk"}, configName = "dt_enabled.yml") +public class DefaultKinesisAsyncClientTest { + public KinesisAsyncClient kinesisAsyncClient; + public HttpExecuteResponse response; + + @Before + public void setup() { + AsyncHttpClient mockHttpClient = new AsyncHttpClient(); + response = mockHttpClient.getResponse(); + kinesisAsyncClient = KinesisAsyncClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(new CredProvider()) + .region(Region.US_EAST_1) + .build(); + } + + @Test + public void testAddTagsToStream() { + txn(() -> kinesisAsyncClient.addTagsToStream(AddTagsToStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/addTagsToStream", true); + } + + @Test + public void testCreateStream() { + txn(() -> kinesisAsyncClient.createStream(CreateStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/createStream", false); + } + + @Test + public void testDecreaseStreamRetentionPeriod() { + txn(() -> kinesisAsyncClient.decreaseStreamRetentionPeriod(DecreaseStreamRetentionPeriodRequest.builder().build())); + assertKinesisTrace("Kinesis/decreaseStreamRetentionPeriod", false); + } + + @Test + public void testDeleteStream() { + txn(() -> kinesisAsyncClient.deleteStream(DeleteStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/deleteStream", false); + } + + @Test + public void testDeregisterStreamConsumer() { + txn(() -> kinesisAsyncClient.deregisterStreamConsumer(DeregisterStreamConsumerRequest.builder().build())); + assertKinesisTrace("Kinesis/deregisterStreamConsumer", false); + } + + @Test + public void testDescribeLimits() { + txn(() -> kinesisAsyncClient.describeLimits(DescribeLimitsRequest.builder().build())); + assertKinesisTrace("Kinesis/describeLimits", false); + } + + @Test + public void testDescribeStream() { + txn(() -> kinesisAsyncClient.describeStream(DescribeStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/describeStream", false); + } + + @Test + public void testDescribeStreamConsumer() { + txn(() -> kinesisAsyncClient.describeStreamConsumer(DescribeStreamConsumerRequest.builder().build())); + assertKinesisTrace("Kinesis/describeStreamConsumer", false); + } + + @Test + public void testDescribeStreamSummary() { + txn(() -> kinesisAsyncClient.describeStreamSummary(DescribeStreamSummaryRequest.builder().build())); + assertKinesisTrace("Kinesis/describeStreamSummary", false); + } + + @Test + public void DisableEnhancedMonitoring() { + txn(() -> kinesisAsyncClient.disableEnhancedMonitoring(DisableEnhancedMonitoringRequest.builder().build())); + assertKinesisTrace("Kinesis/disableEnhancedMonitoring", false); + } + + @Test + public void testEnableEnhancedMonitoring() { + txn(() -> kinesisAsyncClient.enableEnhancedMonitoring(EnableEnhancedMonitoringRequest.builder().build())); + assertKinesisTrace("Kinesis/enableEnhancedMonitoring", false); + } + + @Test + public void testGetRecords() { + txn(() -> kinesisAsyncClient.getRecords(GetRecordsRequest.builder().build())); + assertKinesisTrace("Kinesis/getRecords", false); + } + + @Test + public void testGetShardIterator() { + txn(() -> kinesisAsyncClient.getShardIterator(GetShardIteratorRequest.builder().build())); + assertKinesisTrace("Kinesis/getShardIterator", false); + } + + @Test + public void testIncreaseStreamRetentionPeriod() { + txn(() -> kinesisAsyncClient.increaseStreamRetentionPeriod(IncreaseStreamRetentionPeriodRequest.builder().build())); + assertKinesisTrace("Kinesis/increaseStreamRetentionPeriod", false); + } + + @Test + public void testListShards() { + txn(() -> kinesisAsyncClient.listShards(ListShardsRequest.builder().build())); + assertKinesisTrace("Kinesis/listShards", false); + } + + @Test + public void testListStreams() { + txn(() -> kinesisAsyncClient.listStreams(ListStreamsRequest.builder().build())); + assertKinesisTrace("Kinesis/listStreams", false); + } + + @Test + public void testListTagsForStream() { + txn(() -> kinesisAsyncClient.listTagsForStream(ListTagsForStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/listTagsForStream", false); + } + + @Test + public void testMergeShards() { + txn(() -> kinesisAsyncClient.mergeShards(MergeShardsRequest.builder().build())); + assertKinesisTrace("Kinesis/mergeShards", false); + } + + @Test + public void testPutRecord() { + txn(() -> kinesisAsyncClient.putRecord(PutRecordRequest.builder().build())); + assertKinesisTrace("Kinesis/putRecord", false); + } + + @Test + public void testPutRecords() { + txn(() -> kinesisAsyncClient.putRecords(PutRecordsRequest.builder().build())); + assertKinesisTrace("Kinesis/putRecords", false); + } + + @Test + public void testRegisterStreamConsumer() { + txn(() -> kinesisAsyncClient.registerStreamConsumer(RegisterStreamConsumerRequest.builder().build())); + assertKinesisTrace("Kinesis/registerStreamConsumer", false); + } + + @Test + public void testRemoveTagsFromStream() { + txn(() -> kinesisAsyncClient.removeTagsFromStream(RemoveTagsFromStreamRequest.builder().build())); + assertKinesisTrace("Kinesis/removeTagsFromStream", false); + } + + @Test + public void testSplitShard() { + txn(() -> kinesisAsyncClient.splitShard(SplitShardRequest.builder().build())); + assertKinesisTrace("Kinesis/splitShard", false); + } + + @Test + public void testStartStreamEncryption() { + txn(() -> kinesisAsyncClient.startStreamEncryption(StartStreamEncryptionRequest.builder().build())); + assertKinesisTrace("Kinesis/startStreamEncryption", false); + } + + @Test + public void testStopStreamEncryption() { + txn(() -> kinesisAsyncClient.stopStreamEncryption(StopStreamEncryptionRequest.builder().build())); + assertKinesisTrace("Kinesis/stopStreamEncryption", false); + } + + @Test + public void testUpdateShardCount() { + txn(() -> kinesisAsyncClient.updateShardCount(UpdateShardCountRequest.builder().build())); + assertKinesisTrace("Kinesis/updateShardCount", false); + } + + @Trace(dispatcher = true) + private void txn(Supplier> supplier) { + supplier.get(); + } + + private void assertKinesisTrace(String traceName, boolean assertSpan) { + Introspector introspector = InstrumentationTestRunner.getIntrospector(); + if (assertSpan) { + // Span events fail to be generated when enough transactions are done in succession + SpanEvent kinesisSpan = introspector.getSpanEvents().stream() + .filter(span -> traceName.equals(span.getName())) + .findFirst().orElse(null); + assertNotNull(kinesisSpan); + assertEquals("aws_kinesis_data_streams", kinesisSpan.getAgentAttributes().get("cloud.platform")); + } + Collection transactionTraces = introspector.getTransactionTracesForTransaction( + "OtherTransaction/Custom/software.amazon.awssdk.services.kinesis.DefaultKinesisAsyncClientTest/txn"); + TransactionTrace transactionTrace = transactionTraces.iterator().next(); + List children = transactionTrace.getInitialTraceSegment().getChildren(); + assertEquals(1, children.size()); + TraceSegment trace = children.get(0); + assertEquals(traceName, trace.getName()); + assertEquals("aws_kinesis_data_streams", trace.getTracerAttributes().get("cloud.platform")); + } + + // This mock SdkAsyncHttpClient allows testing the AWS SDK without making actual HTTP requests. + private static class AsyncHttpClient implements SdkAsyncHttpClient { + private ExecutableHttpRequest executableMock; + private HttpExecuteResponse response; + private SdkHttpFullResponse httpResponse; + + public AsyncHttpClient() { + executableMock = mock(ExecutableHttpRequest.class); + response = mock(HttpExecuteResponse.class, Mockito.RETURNS_DEEP_STUBS); + httpResponse = mock(SdkHttpFullResponse.class, Mockito.RETURNS_DEEP_STUBS); + when(response.httpResponse()).thenReturn(httpResponse); + when(httpResponse.toBuilder().content(any()).build()).thenReturn(httpResponse); + when(httpResponse.isSuccessful()).thenReturn(true); + AbortableInputStream inputStream = AbortableInputStream.create(new StringInputStream("Dont panic")); + when(httpResponse.content()).thenReturn(Optional.of(inputStream)); + try { + when(executableMock.call()).thenReturn(response); + + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Override + public void close() { + } + + @Override + public CompletableFuture execute(AsyncExecuteRequest asyncExecuteRequest) { + asyncExecuteRequest.responseHandler().onStream(new EmptyPublisher<>()); + return new CompletableFuture<>(); + } + + @Override + public String clientName() { + return "MockHttpClient"; + } + + public HttpExecuteResponse getResponse() { + return response; + } + } + + private static class CredProvider implements AwsCredentialsProvider { + @Override + public AwsCredentials resolveCredentials() { + AwsCredentials credentials = mock(AwsCredentials.class); + when(credentials.accessKeyId()).thenReturn("accessKeyId"); + when(credentials.secretAccessKey()).thenReturn("secretAccessKey"); + return credentials; + } + } +} From 0ca231ee8cbe39c4453358aff24b8dfa21bc6054 Mon Sep 17 00:00:00 2001 From: Oren Ben-Meir Date: Wed, 28 Aug 2024 15:37:21 -0400 Subject: [PATCH 6/8] Create kinesis v1 sdk test --- .../services/kinesis/KinesisUtil.java | 49 ++++++ .../{lambda => kinesis}/SegmentHandler.java | 2 +- .../services/lambda/KinesisUtil.java | 50 ------ ...azonWebServiceRequest_Instrumentation.java | 11 -- ...zonKinesisAsyncClient_Instrumentation.java | 55 +++--- .../AmazonKinesisClient_Instrumentation.java | 3 +- .../kinesis/AmazonKinesisAPITest.java | 164 ++++++++++++++++++ .../{lambda => kinesis}/KinesisUtil.java | 2 +- .../{lambda => kinesis}/SegmentHandler.java | 2 +- ...ultKinesisAsyncClient_Instrumentation.java | 5 +- .../DefaultKinesisClient_Instrumentation.java | 2 +- 11 files changed, 255 insertions(+), 90 deletions(-) create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java rename instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/{lambda => kinesis}/SegmentHandler.java (93%) delete mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java delete mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest_Instrumentation.java create mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/java/com/amazonaws/services/kinesis/AmazonKinesisAPITest.java rename instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/{lambda => kinesis}/KinesisUtil.java (94%) rename instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/{lambda => kinesis}/SegmentHandler.java (94%) diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java new file mode 100644 index 0000000000..48b5da7b16 --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java @@ -0,0 +1,49 @@ +package com.agent.instrumentation.awsjavasdk2.services.kinesis; + +import com.amazonaws.AmazonWebServiceRequest; +import com.newrelic.agent.bridge.AgentBridge; +import com.newrelic.api.agent.CloudParameters; +import com.newrelic.api.agent.NewRelic; +import com.newrelic.api.agent.Token; +import com.newrelic.api.agent.TracedMethod; + +import java.util.Map; + +public class KinesisUtil { + + public static final String PLATFORM = "aws_kinesis_data_streams"; + public static final String TRACE_CATEGORY = "Kinesis"; + + public static final Map requestTokenMap = AgentBridge.collectionFactory.createConcurrentWeakKeyedMap(); + private KinesisUtil() {} + + public static void setTokenForRequest(AmazonWebServiceRequest request) { + if (AgentBridge.getAgent().getTransaction(false) != null) { + if (request != null) { + requestTokenMap.put(request, NewRelic.getAgent().getTransaction().getToken()); + } + } + } + + public static void linkAndExpireToken(AmazonWebServiceRequest request) { + if (request != null) { + Token token = requestTokenMap.get(request); + if (token != null) { + token.linkAndExpire(); + } + requestTokenMap.remove(request); + } + } + + public static void setTraceDetails(String kinesisOperation) { + TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + tracedMethod.setMetricName(TRACE_CATEGORY, kinesisOperation); + tracedMethod.reportAsExternal(createCloudParams()); + } + + public static CloudParameters createCloudParams() { + // Todo: add arn to cloud parameters + return CloudParameters.provider(PLATFORM).build(); + } + +} diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java similarity index 93% rename from instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java rename to instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java index 1a042efd89..5d30e4483e 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java @@ -1,4 +1,4 @@ -package com.agent.instrumentation.awsjavasdk2.services.lambda; +package com.agent.instrumentation.awsjavasdk2.services.kinesis; import com.newrelic.agent.bridge.AgentBridge; import com.newrelic.api.agent.Segment; diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java deleted file mode 100644 index ae72f54830..0000000000 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java +++ /dev/null @@ -1,50 +0,0 @@ -package com.agent.instrumentation.awsjavasdk2.services.lambda; - -import com.amazonaws.AmazonWebServiceRequest; -import com.amazonaws.AmazonWebServiceRequest_Instrumentation; -import com.amazonaws.handlers.AsyncHandler_Instrumentation; -import com.newrelic.agent.bridge.AgentBridge; -import com.newrelic.api.agent.CloudParameters; -import com.newrelic.api.agent.NewRelic; -import com.newrelic.api.agent.TracedMethod; - -public class KinesisUtil { - - public static final String PLATFORM = "aws_kinesis_data_streams"; - public static final String TRACE_CATEGORY = "Kinesis"; - private KinesisUtil() {} - - public static void setupToken(AsyncHandler_Instrumentation asyncHandler, AmazonWebServiceRequest amazonWebServiceRequest) { - if (AgentBridge.getAgent().getTransaction(false) != null) { - if (asyncHandler != null) { - asyncHandler.token = NewRelic.getAgent().getTransaction().getToken(); - } - if (amazonWebServiceRequest != null && amazonWebServiceRequest instanceof AmazonWebServiceRequest_Instrumentation) { - AmazonWebServiceRequest_Instrumentation request = (AmazonWebServiceRequest_Instrumentation)amazonWebServiceRequest; - request.token = NewRelic.getAgent().getTransaction().getToken(); - } - } - } - - public static void linkAndExpireToken(AmazonWebServiceRequest amazonWebServiceRequest) { - if (amazonWebServiceRequest != null && amazonWebServiceRequest instanceof AmazonWebServiceRequest_Instrumentation) { - AmazonWebServiceRequest_Instrumentation request = (AmazonWebServiceRequest_Instrumentation)amazonWebServiceRequest; - if (request.token != null) { - request.token.linkAndExpire(); - request.token = null; - } - } - } - - public static void setTraceDetails(String kinesisOperation) { - TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); - tracedMethod.setMetricName(TRACE_CATEGORY, kinesisOperation); - tracedMethod.reportAsExternal(createCloudParams()); - } - - public static CloudParameters createCloudParams() { - // Todo: add arn to cloud parameters - return CloudParameters.provider(PLATFORM).build(); - } - -} diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest_Instrumentation.java deleted file mode 100644 index 43065e13f4..0000000000 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/AmazonWebServiceRequest_Instrumentation.java +++ /dev/null @@ -1,11 +0,0 @@ -package com.amazonaws; - -import com.newrelic.api.agent.Token; -import com.newrelic.api.agent.weaver.NewField; -import com.newrelic.api.agent.weaver.Weave; - -@Weave(originalName = "com.amazonaws.AmazonWebServiceRequest_Instrumentation") -public abstract class AmazonWebServiceRequest_Instrumentation extends AmazonWebServiceRequest { - @NewField - public Token token; -} diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java index 6547a01859..e255a28ce3 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java @@ -1,6 +1,7 @@ package com.amazonaws.services.kinesis; -import com.agent.instrumentation.awsjavasdk2.services.lambda.KinesisUtil; +import com.agent.instrumentation.awsjavasdk2.services.kinesis.KinesisUtil; +import com.amazonaws.AmazonWebServiceRequest; import com.amazonaws.handlers.AsyncHandler_Instrumentation; import com.amazonaws.services.kinesis.model.AddTagsToStreamRequest; import com.amazonaws.services.kinesis.model.AddTagsToStreamResult; @@ -40,24 +41,28 @@ import com.amazonaws.services.kinesis.model.SplitShardResult; import com.amazonaws.services.kinesis.model.UpdateShardCountRequest; import com.amazonaws.services.kinesis.model.UpdateShardCountResult; +import com.newrelic.agent.bridge.AgentBridge; +import com.newrelic.api.agent.NewRelic; import com.newrelic.api.agent.Trace; import com.newrelic.api.agent.weaver.Weaver; import java.util.concurrent.Future; +import static com.agent.instrumentation.awsjavasdk2.services.kinesis.KinesisUtil.requestTokenMap; + public class AmazonKinesisAsyncClient_Instrumentation { @Trace public Future addTagsToStreamAsync(AddTagsToStreamRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future createStreamAsync(CreateStreamRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @@ -65,55 +70,55 @@ public Future createStreamAsync(CreateStreamRequest request, public Future decreaseStreamRetentionPeriodAsync( DecreaseStreamRetentionPeriodRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future deleteStreamAsync(DeleteStreamRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future describeLimitsAsync(DescribeLimitsRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future describeStreamAsync(DescribeStreamRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future disableEnhancedMonitoringAsync(DisableEnhancedMonitoringRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future enableEnhancedMonitoringAsync(EnableEnhancedMonitoringRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future getRecordsAsync(GetRecordsRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future getShardIteratorAsync(GetShardIteratorRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @@ -121,58 +126,68 @@ public Future getShardIteratorAsync(GetShardIteratorRequ public Future increaseStreamRetentionPeriodAsync( IncreaseStreamRetentionPeriodRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future listStreamsAsync(ListStreamsRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future listTagsForStreamAsync(ListTagsForStreamRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future mergeShardsAsync(MergeShardsRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future putRecordAsync(PutRecordRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future putRecordsAsync(PutRecordsRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future removeTagsFromStreamAsync(RemoveTagsFromStreamRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future splitShardAsync(SplitShardRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } @Trace public Future updateShardCountAsync(UpdateShardCountRequest request, AsyncHandler_Instrumentation asyncHandler) { - KinesisUtil.setupToken(asyncHandler, request); + setToken(asyncHandler, request); return Weaver.callOriginal(); } + + public static void setToken(AsyncHandler_Instrumentation asyncHandler, AmazonWebServiceRequest request) { + if (AgentBridge.getAgent().getTransaction(false) != null) { + if (asyncHandler != null) { + asyncHandler.token = NewRelic.getAgent().getTransaction().getToken(); + } + KinesisUtil.setTokenForRequest(request); + } + } + } diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java index 4c404b3f87..a759fc82e9 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java @@ -1,7 +1,6 @@ package com.amazonaws.services.kinesis; -import com.agent.instrumentation.awsjavasdk2.services.lambda.KinesisUtil; -import com.amazonaws.AmazonWebServiceRequest; +import com.agent.instrumentation.awsjavasdk2.services.kinesis.KinesisUtil; import com.amazonaws.services.kinesis.model.AddTagsToStreamRequest; import com.amazonaws.services.kinesis.model.AddTagsToStreamResult; import com.amazonaws.services.kinesis.model.CreateStreamRequest; diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/java/com/amazonaws/services/kinesis/AmazonKinesisAPITest.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/java/com/amazonaws/services/kinesis/AmazonKinesisAPITest.java new file mode 100644 index 0000000000..df642e3e3e --- /dev/null +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/java/com/amazonaws/services/kinesis/AmazonKinesisAPITest.java @@ -0,0 +1,164 @@ +package com.amazonaws.services.kinesis; + +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.client.builder.AwsClientBuilder; +import com.amazonaws.handlers.AsyncHandler; +import com.amazonaws.services.kinesis.model.AddTagsToStreamRequest; +import com.amazonaws.services.kinesis.model.AddTagsToStreamResult; +import com.newrelic.agent.introspec.InstrumentationTestConfig; +import com.newrelic.agent.introspec.InstrumentationTestRunner; +import com.newrelic.agent.introspec.Introspector; +import com.newrelic.agent.introspec.SpanEvent; +import com.newrelic.agent.introspec.TraceSegment; +import com.newrelic.agent.introspec.TransactionTrace; +import com.newrelic.agent.introspec.internal.HttpServerRule; +import com.newrelic.api.agent.Trace; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.internal.matchers.GreaterThan; + +import java.net.URISyntaxException; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import java.util.function.Function; +import java.util.function.Supplier; +import java.util.stream.Collectors; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@RunWith(InstrumentationTestRunner.class) +@InstrumentationTestConfig(includePrefixes = {"com.amazonaws"}, configName = "dt_enabled.yml") +public class AmazonKinesisAPITest { + + @Rule + public HttpServerRule server = new HttpServerRule(); + private AmazonKinesis kinesisClient; + private AmazonKinesisAsync kinesisAsyncClient; + + @Before + public void setup() throws URISyntaxException { + AwsClientBuilder.EndpointConfiguration endpoint = new AwsClientBuilder.EndpointConfiguration(server.getEndPoint().toString(), "us-east-1"); + kinesisClient = AmazonKinesisClientBuilder.standard() + .withCredentials(new CredProvider()) + .withEndpointConfiguration(endpoint) + .build(); + kinesisAsyncClient = AmazonKinesisAsyncClientBuilder.standard() + .withCredentials(new CredProvider()) + .withEndpointConfiguration(endpoint) + .build(); + } + + @Test + public void testaddCreateStream() { + txn(() -> kinesisClient.addTagsToStream(new AddTagsToStreamRequest())); +// txnAsync(() -> kinesisAsyncClient.addTagsToStreamAsync(new AddTagsToStreamRequest())); +// txnAsyncWithHandler(handler -> kinesisAsyncClient.addTagsToStreamAsync(new AddTagsToStreamRequest(), handler), +// new AsyncHandlerNoOp()); + assertKinesisTrace("addTagsToStream", true); + } + + @Trace(dispatcher = true) + public void txn(Runnable runnable) { + runnable.run(); + } + + @Trace(dispatcher = true) + public void txnAsync(Supplier> asyncSupplier) { + try { + asyncSupplier.get().get(); + } catch (Exception e) { + } + } + + @Trace(dispatcher = true) + public void txnAsyncWithHandler(Function, Future> function, AsyncHandler handler) { + try { + function.apply(handler).get(); + } catch (Exception e) { + } + } + + private void assertKinesisTrace(String kinesisOperation, boolean assertSpan) { + Introspector introspector = InstrumentationTestRunner.getIntrospector(); + final String traceName = "Kinesis/" + kinesisOperation; + if (assertSpan) { + // Span events fail to be generated when enough transactions are done in succession + List kinesisSpans = introspector.getSpanEvents().stream() + .filter(span -> traceName.equals(span.getName())) + .collect(Collectors.toList()); + assertFalse(kinesisSpans.isEmpty()); + for (SpanEvent kinesisSpan: kinesisSpans) { + assertEquals("aws_kinesis_data_streams", kinesisSpan.getAgentAttributes().get("cloud.platform")); + } + } + assertTxn(kinesisOperation, introspector); +// assertTxnAsync(kinesisOperation, "OtherTransaction/Custom/com.amazonaws.services.kinesis.AmazonKinesisAPITest/txnAsync", introspector); +// assertTxnAsync(kinesisOperation, "OtherTransaction/Custom/com.amazonaws.services.kinesis.AmazonKinesisAPITest/txnAsyncWithHandler", introspector); + } + + private void assertTxn(String kinesisOperation, Introspector introspector) { + String transactionName = "OtherTransaction/Custom/com.amazonaws.services.kinesis.AmazonKinesisAPITest/txn"; + final String traceName = "Kinesis/" + kinesisOperation; + Collection transactionTraces = introspector.getTransactionTracesForTransaction(transactionName); + TransactionTrace transactionTrace = transactionTraces.iterator().next(); + List children = transactionTrace.getInitialTraceSegment().getChildren(); + assertEquals(1, children.size()); + TraceSegment trace = children.get(0); + assertEquals(traceName, trace.getName()); + assertEquals("aws_kinesis_data_streams", trace.getTracerAttributes().get("cloud.platform")); + } + + private void assertTxnAsync(String kinesisOperation, String transactionName, Introspector introspector) { + final String traceName = "Kinesis/" + kinesisOperation; + Collection transactionTraces = introspector.getTransactionTracesForTransaction(transactionName); + TransactionTrace transactionTrace = transactionTraces.iterator().next(); + List children = transactionTrace.getInitialTraceSegment().getChildren(); + assertEquals(2, children.size()); + TraceSegment asyncFunctionTrace = children.get(0); + String expectedFunctionTraceName = "com.amazonaws.services.kinesis.AmazonKinesisAsyncClient/" + kinesisOperation; + assertEquals(expectedFunctionTraceName, asyncFunctionTrace.getName()); + TraceSegment externalTrace = children.get(1); + assertEquals(traceName, externalTrace.getName()); + assertEquals("aws_kinesis_data_streams", externalTrace.getTracerAttributes().get("cloud.platform")); + } + + private static class CredProvider implements AWSCredentialsProvider { + @Override + public AWSCredentials getCredentials() { + AWSCredentials credentials = mock(AWSCredentials.class); + when(credentials.getAWSAccessKeyId()).thenReturn("accessKeyId"); + when(credentials.getAWSSecretKey()).thenReturn("secretAccessKey"); + return credentials; + } + + @Override + public void refresh() { + + } + } + + // To trigger scenarios where the async handler is called. + private static class AsyncHandlerNoOp implements AsyncHandler { + + @Override + public void onError(Exception exception) { + + } + + @Override + public void onSuccess(REQ request, RES res) { + + } + } +} diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java similarity index 94% rename from instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java rename to instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java index 07bf5ac3bc..ce69329bee 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/KinesisUtil.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java @@ -1,4 +1,4 @@ -package com.agent.instrumentation.awsjavasdk2.services.lambda; +package com.agent.instrumentation.awsjavasdk2.services.kinesis; import com.newrelic.api.agent.CloudParameters; import com.newrelic.api.agent.NewRelic; diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java similarity index 94% rename from instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java rename to instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java index 4ba94e4997..d996af803a 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/lambda/SegmentHandler.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java @@ -1,4 +1,4 @@ -package com.agent.instrumentation.awsjavasdk2.services.lambda; +package com.agent.instrumentation.awsjavasdk2.services.kinesis; import com.newrelic.agent.bridge.AgentBridge; import com.newrelic.api.agent.Segment; diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java index e0d6a53873..29763c9dc8 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisAsyncClient_Instrumentation.java @@ -1,8 +1,7 @@ package software.amazon.awssdk.services.kinesis; -import com.agent.instrumentation.awsjavasdk2.services.lambda.KinesisUtil; -import com.agent.instrumentation.awsjavasdk2.services.lambda.SegmentHandler; -import com.newrelic.api.agent.NewRelic; +import com.agent.instrumentation.awsjavasdk2.services.kinesis.KinesisUtil; +import com.agent.instrumentation.awsjavasdk2.services.kinesis.SegmentHandler; import com.newrelic.api.agent.Segment; import com.newrelic.api.agent.weaver.MatchType; import com.newrelic.api.agent.weaver.Weave; diff --git a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java index a9223fcccd..04ae89282a 100644 --- a/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-2.0.6/src/main/java/software/amazon/awssdk/services/kinesis/DefaultKinesisClient_Instrumentation.java @@ -1,6 +1,6 @@ package software.amazon.awssdk.services.kinesis; -import com.agent.instrumentation.awsjavasdk2.services.lambda.KinesisUtil; +import com.agent.instrumentation.awsjavasdk2.services.kinesis.KinesisUtil; import com.newrelic.api.agent.Trace; import com.newrelic.api.agent.weaver.MatchType; import com.newrelic.api.agent.weaver.Weave; From 5376047fa1af7ed90d0131bf389469caf813340b Mon Sep 17 00:00:00 2001 From: Oren Ben-Meir Date: Wed, 28 Aug 2024 17:20:36 -0400 Subject: [PATCH 7/8] Add v1 sdk test cases and fix token linking --- .../services/kinesis/KinesisUtil.java | 30 ++- .../services/kinesis/SegmentHandler.java | 31 --- .../AsyncHandler_Instrumentation.java | 1 - ...zonKinesisAsyncClient_Instrumentation.java | 7 +- .../AmazonKinesisClient_Instrumentation.java | 97 ++++---- .../kinesis/AmazonKinesisAPITest.java | 218 ++++++++++++++---- 6 files changed, 243 insertions(+), 141 deletions(-) delete mode 100644 instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java index 48b5da7b16..88e4a8d97f 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/KinesisUtil.java @@ -1,6 +1,7 @@ package com.agent.instrumentation.awsjavasdk2.services.kinesis; import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.handlers.AsyncHandler_Instrumentation; import com.newrelic.agent.bridge.AgentBridge; import com.newrelic.api.agent.CloudParameters; import com.newrelic.api.agent.NewRelic; @@ -20,23 +21,36 @@ private KinesisUtil() {} public static void setTokenForRequest(AmazonWebServiceRequest request) { if (AgentBridge.getAgent().getTransaction(false) != null) { if (request != null) { - requestTokenMap.put(request, NewRelic.getAgent().getTransaction().getToken()); + Token token = NewRelic.getAgent().getTransaction().getToken(); + requestTokenMap.put(request, token); } } } - public static void linkAndExpireToken(AmazonWebServiceRequest request) { + public static void setTraceInformation(String kinesisOperation, AmazonWebServiceRequest request) { + Token token = KinesisUtil.getToken(request); + if (token != null) { + token.linkAndExpire(); + } + KinesisUtil.cleanToken(request); + TracedMethod tracedMethod = NewRelic.getAgent().getTransaction().getTracedMethod(); + KinesisUtil.setTraceDetails(kinesisOperation, tracedMethod); + } + + public static Token getToken(AmazonWebServiceRequest request) { + if (request != null) { + return requestTokenMap.get(request); + } + return null; + } + + public static void cleanToken(AmazonWebServiceRequest request) { if (request != null) { - Token token = requestTokenMap.get(request); - if (token != null) { - token.linkAndExpire(); - } requestTokenMap.remove(request); } } - public static void setTraceDetails(String kinesisOperation) { - TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod(); + public static void setTraceDetails(String kinesisOperation, TracedMethod tracedMethod) { tracedMethod.setMetricName(TRACE_CATEGORY, kinesisOperation); tracedMethod.reportAsExternal(createCloudParams()); } diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java deleted file mode 100644 index 5d30e4483e..0000000000 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/agent/instrumentation/awsjavasdk2/services/kinesis/SegmentHandler.java +++ /dev/null @@ -1,31 +0,0 @@ -package com.agent.instrumentation.awsjavasdk2.services.kinesis; - -import com.newrelic.agent.bridge.AgentBridge; -import com.newrelic.api.agent.Segment; - -import java.util.concurrent.CompletableFuture; - -public class SegmentHandler { - private final CompletableFuture completableFuture; - private final Segment segment; - private final String implementationTitle; - - public SegmentHandler(CompletableFuture completableFuture, Segment segment, String implementationTitle) { - this.completableFuture = completableFuture; - this.segment = segment; - this.implementationTitle = implementationTitle; - } - - public CompletableFuture newSegmentCompletionStage() { - if (completableFuture == null) { - return null; - } - return completableFuture.whenComplete((r, t) -> { - try { - segment.end(); - } catch (Throwable t1) { - AgentBridge.instrumentation.noticeInstrumentationError(t1, implementationTitle); - } - }); - } -} \ No newline at end of file diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java index d5e44f9b26..6a5e814917 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/handlers/AsyncHandler_Instrumentation.java @@ -8,7 +8,6 @@ package com.amazonaws.handlers; import com.amazonaws.AmazonWebServiceRequest; -import com.amazonaws.AmazonWebServiceRequest_Instrumentation; import com.newrelic.api.agent.Token; import com.newrelic.api.agent.Trace; import com.newrelic.api.agent.weaver.MatchType; diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java index e255a28ce3..eee81f6fb1 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisAsyncClient_Instrumentation.java @@ -44,12 +44,13 @@ import com.newrelic.agent.bridge.AgentBridge; import com.newrelic.api.agent.NewRelic; import com.newrelic.api.agent.Trace; +import com.newrelic.api.agent.weaver.MatchType; +import com.newrelic.api.agent.weaver.Weave; import com.newrelic.api.agent.weaver.Weaver; import java.util.concurrent.Future; -import static com.agent.instrumentation.awsjavasdk2.services.kinesis.KinesisUtil.requestTokenMap; - +@Weave(originalName ="com.amazonaws.services.kinesis.AmazonKinesisAsyncClient", type = MatchType.ExactClass) public class AmazonKinesisAsyncClient_Instrumentation { @Trace @@ -181,7 +182,7 @@ public Future updateShardCountAsync(UpdateShardCountRequ return Weaver.callOriginal(); } - public static void setToken(AsyncHandler_Instrumentation asyncHandler, AmazonWebServiceRequest request) { + private void setToken(AsyncHandler_Instrumentation asyncHandler, AmazonWebServiceRequest request) { if (AgentBridge.getAgent().getTransaction(false) != null) { if (asyncHandler != null) { asyncHandler.token = NewRelic.getAgent().getTransaction().getToken(); diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java index a759fc82e9..cf70acb0d3 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/main/java/com/amazonaws/services/kinesis/AmazonKinesisClient_Instrumentation.java @@ -39,6 +39,8 @@ import com.amazonaws.services.kinesis.model.SplitShardResult; import com.amazonaws.services.kinesis.model.UpdateShardCountRequest; import com.amazonaws.services.kinesis.model.UpdateShardCountResult; +import com.newrelic.api.agent.NewRelic; +import com.newrelic.api.agent.Token; import com.newrelic.api.agent.Trace; import com.newrelic.api.agent.weaver.MatchType; import com.newrelic.api.agent.weaver.Weave; @@ -47,136 +49,117 @@ @Weave(originalName = "com.amazonaws.services.kinesis.AmazonKinesisClient", type = MatchType.ExactClass) public class AmazonKinesisClient_Instrumentation { - @Trace(leaf=true) + @Trace(async = true, leaf = true) final AddTagsToStreamResult executeAddTagsToStream(AddTagsToStreamRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("addTagsToStream"); + KinesisUtil.setTraceInformation("addTagsToStream", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final CreateStreamResult executeCreateStream(CreateStreamRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("createStream"); + KinesisUtil.setTraceInformation("createStream", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final DecreaseStreamRetentionPeriodResult executeDecreaseStreamRetentionPeriod(DecreaseStreamRetentionPeriodRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("decreaseStreamRetentionPeriod"); + KinesisUtil.setTraceInformation("decreaseStreamRetentionPeriod", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final DeleteStreamResult executeDeleteStream(DeleteStreamRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("deleteStream"); + KinesisUtil.setTraceInformation("deleteStream", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final DescribeLimitsResult executeDescribeLimits(DescribeLimitsRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("describeLimits"); + KinesisUtil.setTraceInformation("describeLimits", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final DescribeStreamResult executeDescribeStream(DescribeStreamRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("describeStream"); + KinesisUtil.setTraceInformation("describeStream", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final DisableEnhancedMonitoringResult executeDisableEnhancedMonitoring(DisableEnhancedMonitoringRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("disableEnhancedMonitoring"); + KinesisUtil.setTraceInformation("disableEnhancedMonitoring", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final EnableEnhancedMonitoringResult executeEnableEnhancedMonitoring(EnableEnhancedMonitoringRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("enableEnhancedMonitoring"); + KinesisUtil.setTraceInformation("enableEnhancedMonitoring", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final GetRecordsResult executeGetRecords(GetRecordsRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("getRecords"); + KinesisUtil.setTraceInformation("getRecords", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final GetShardIteratorResult executeGetShardIterator(GetShardIteratorRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("getShardIterator"); + KinesisUtil.setTraceInformation("getShardIterator", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final IncreaseStreamRetentionPeriodResult executeIncreaseStreamRetentionPeriod(IncreaseStreamRetentionPeriodRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("increaseStreamRetentionPeriod"); + KinesisUtil.setTraceInformation("increaseStreamRetentionPeriod", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final ListStreamsResult executeListStreams(ListStreamsRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("listStreams"); + KinesisUtil.setTraceInformation("listStreams", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final ListTagsForStreamResult executeListTagsForStream(ListTagsForStreamRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("listTagsForStream"); + KinesisUtil.setTraceInformation("listTagsForStream", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final MergeShardsResult executeMergeShards(MergeShardsRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("mergeShards"); + KinesisUtil.setTraceInformation("mergeShards", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final PutRecordResult executePutRecord(PutRecordRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("putRecord"); + KinesisUtil.setTraceInformation("putRecord", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final PutRecordsResult executePutRecords(PutRecordsRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("putRecords"); + KinesisUtil.setTraceInformation("putRecords", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final RemoveTagsFromStreamResult executeRemoveTagsFromStream(RemoveTagsFromStreamRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("removeTagsFromStream"); + KinesisUtil.setTraceInformation("removeTagsFromStream", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final SplitShardResult executeSplitShard(SplitShardRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("splitShard"); + KinesisUtil.setTraceInformation("splitShard", request); return Weaver.callOriginal(); } - @Trace(leaf=true) + @Trace(async = true, leaf = true) final UpdateShardCountResult executeUpdateShardCount(UpdateShardCountRequest request) { - KinesisUtil.linkAndExpireToken(request); - KinesisUtil.setTraceDetails("updateShardCount"); + KinesisUtil.setTraceInformation("updateShardCount", request); return Weaver.callOriginal(); } diff --git a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/java/com/amazonaws/services/kinesis/AmazonKinesisAPITest.java b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/java/com/amazonaws/services/kinesis/AmazonKinesisAPITest.java index df642e3e3e..f564ff645d 100644 --- a/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/java/com/amazonaws/services/kinesis/AmazonKinesisAPITest.java +++ b/instrumentation/aws-java-sdk-kinesis-1.11.106/src/test/java/com/amazonaws/services/kinesis/AmazonKinesisAPITest.java @@ -6,7 +6,24 @@ import com.amazonaws.client.builder.AwsClientBuilder; import com.amazonaws.handlers.AsyncHandler; import com.amazonaws.services.kinesis.model.AddTagsToStreamRequest; -import com.amazonaws.services.kinesis.model.AddTagsToStreamResult; +import com.amazonaws.services.kinesis.model.CreateStreamRequest; +import com.amazonaws.services.kinesis.model.DecreaseStreamRetentionPeriodRequest; +import com.amazonaws.services.kinesis.model.DeleteStreamRequest; +import com.amazonaws.services.kinesis.model.DescribeLimitsRequest; +import com.amazonaws.services.kinesis.model.DescribeStreamRequest; +import com.amazonaws.services.kinesis.model.DisableEnhancedMonitoringRequest; +import com.amazonaws.services.kinesis.model.EnableEnhancedMonitoringRequest; +import com.amazonaws.services.kinesis.model.GetRecordsRequest; +import com.amazonaws.services.kinesis.model.GetShardIteratorRequest; +import com.amazonaws.services.kinesis.model.IncreaseStreamRetentionPeriodRequest; +import com.amazonaws.services.kinesis.model.ListStreamsRequest; +import com.amazonaws.services.kinesis.model.ListTagsForStreamRequest; +import com.amazonaws.services.kinesis.model.MergeShardsRequest; +import com.amazonaws.services.kinesis.model.PutRecordRequest; +import com.amazonaws.services.kinesis.model.PutRecordsRequest; +import com.amazonaws.services.kinesis.model.RemoveTagsFromStreamRequest; +import com.amazonaws.services.kinesis.model.SplitShardRequest; +import com.amazonaws.services.kinesis.model.UpdateShardCountRequest; import com.newrelic.agent.introspec.InstrumentationTestConfig; import com.newrelic.agent.introspec.InstrumentationTestRunner; import com.newrelic.agent.introspec.Introspector; @@ -19,12 +36,10 @@ import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; -import org.mockito.internal.matchers.GreaterThan; import java.net.URISyntaxException; import java.util.Collection; import java.util.List; -import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; import java.util.function.Function; import java.util.function.Supplier; @@ -32,8 +47,6 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertThat; -import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -48,7 +61,8 @@ public class AmazonKinesisAPITest { @Before public void setup() throws URISyntaxException { - AwsClientBuilder.EndpointConfiguration endpoint = new AwsClientBuilder.EndpointConfiguration(server.getEndPoint().toString(), "us-east-1"); + String serverUriStr = server.getEndPoint().toString(); + AwsClientBuilder.EndpointConfiguration endpoint = new AwsClientBuilder.EndpointConfiguration(serverUriStr, "us-east-1"); kinesisClient = AmazonKinesisClientBuilder.standard() .withCredentials(new CredProvider()) .withEndpointConfiguration(endpoint) @@ -59,25 +73,157 @@ public void setup() throws URISyntaxException { .build(); } + // HttpServerRule is flaky so only 1 test is run + @Test - public void testaddCreateStream() { + public void testAddTagsToStream() { txn(() -> kinesisClient.addTagsToStream(new AddTagsToStreamRequest())); -// txnAsync(() -> kinesisAsyncClient.addTagsToStreamAsync(new AddTagsToStreamRequest())); -// txnAsyncWithHandler(handler -> kinesisAsyncClient.addTagsToStreamAsync(new AddTagsToStreamRequest(), handler), -// new AsyncHandlerNoOp()); - assertKinesisTrace("addTagsToStream", true); + txnAsync(() -> kinesisAsyncClient.addTagsToStreamAsync(new AddTagsToStreamRequest())); + assertKinesisTrace("addTagsToStream", false); } +// @Test +// public void testCreateStream() { +// txn(() -> kinesisClient.createStream(new CreateStreamRequest())); +// txnAsync(() -> kinesisAsyncClient.createStreamAsync(new CreateStreamRequest())); +// assertKinesisTrace("createStream", false); +// } +// +// @Test +// public void testDecreaseStreamRetentionPeriod() { +// txn(() -> kinesisClient.decreaseStreamRetentionPeriod(new DecreaseStreamRetentionPeriodRequest())); +// txnAsync(() -> kinesisAsyncClient.decreaseStreamRetentionPeriodAsync(new DecreaseStreamRetentionPeriodRequest())); +// assertKinesisTrace("decreaseStreamRetentionPeriod", false); +// } +// +// @Test +// public void testDeleteStream() { +// txn(() -> kinesisClient.deleteStream(new DeleteStreamRequest())); +// txnAsync(() -> kinesisAsyncClient.deleteStreamAsync(new DeleteStreamRequest())); +// assertKinesisTrace("deleteStream", false); +// } +// +// @Test +// public void testDescribeLimits() { +// txn(() -> kinesisClient.describeLimits(new DescribeLimitsRequest())); +// txnAsync(() -> kinesisAsyncClient.describeLimitsAsync(new DescribeLimitsRequest())); +// assertKinesisTrace("describeLimits", false); +// } +// +// @Test +// public void testDescribeStream() { +// txn(() -> kinesisClient.describeStream(new DescribeStreamRequest())); +// txnAsync(() -> kinesisAsyncClient.describeStreamAsync(new DescribeStreamRequest())); +// assertKinesisTrace("describeStream", false); +// } +// +// @Test +// public void testDisableEnhancedMonitoring() { +// txn(() -> kinesisClient.disableEnhancedMonitoring(new DisableEnhancedMonitoringRequest())); +// txnAsync(() -> kinesisAsyncClient.disableEnhancedMonitoringAsync(new DisableEnhancedMonitoringRequest())); +// assertKinesisTrace("disableEnhancedMonitoring", false); +// } +// +// @Test +// public void testEnableEnhancedMonitoring() { +// txn(() -> kinesisClient.enableEnhancedMonitoring(new EnableEnhancedMonitoringRequest())); +// txnAsync(() -> kinesisAsyncClient.enableEnhancedMonitoringAsync(new EnableEnhancedMonitoringRequest())); +// assertKinesisTrace("enableEnhancedMonitoring", false); +// } +// +// @Test +// public void testGetRecords() { +// txn(() -> kinesisClient.getRecords(new GetRecordsRequest())); +// txnAsync(() -> kinesisAsyncClient.getRecordsAsync(new GetRecordsRequest())); +// assertKinesisTrace("getRecords", false); +// } +// +// @Test +// public void testGetShardIterator() { +// txn(() -> kinesisClient.getShardIterator(new GetShardIteratorRequest())); +// txnAsync(() -> kinesisAsyncClient.getShardIteratorAsync(new GetShardIteratorRequest())); +// assertKinesisTrace("getShardIterator", false); +// } +// +// @Test +// public void testIncreaseStreamRetentionPeriod() { +// txn(() -> kinesisClient.increaseStreamRetentionPeriod(new IncreaseStreamRetentionPeriodRequest())); +// txnAsync(() -> kinesisAsyncClient.increaseStreamRetentionPeriodAsync(new IncreaseStreamRetentionPeriodRequest())); +// assertKinesisTrace("increaseStreamRetentionPeriod", false); +// } +// +// @Test +// public void testListStreams() { +// txn(() -> kinesisClient.listStreams(new ListStreamsRequest())); +// txnAsync(() -> kinesisAsyncClient.listStreamsAsync(new ListStreamsRequest())); +// assertKinesisTrace("listStreams", false); +// } +// +// @Test +// public void testListTagsForStream() { +// txn(() -> kinesisClient.listTagsForStream(new ListTagsForStreamRequest())); +// txnAsync(() -> kinesisAsyncClient.listTagsForStreamAsync(new ListTagsForStreamRequest())); +// assertKinesisTrace("listTagsForStream", false); +// } +// +// @Test +// public void testMergeShards() { +// txn(() -> kinesisClient.mergeShards(new MergeShardsRequest())); +// txnAsync(() -> kinesisAsyncClient.mergeShardsAsync(new MergeShardsRequest())); +// assertKinesisTrace("mergeShards", false); +// } +// +// @Test +// public void testPutRecord() { +// txn(() -> kinesisClient.putRecord(new PutRecordRequest())); +// txnAsync(() -> kinesisAsyncClient.putRecordAsync(new PutRecordRequest())); +// assertKinesisTrace("putRecord", false); +// } +// +// @Test +// public void testPutRecords() { +// txn(() -> kinesisClient.putRecords(new PutRecordsRequest())); +// txnAsync(() -> kinesisAsyncClient.putRecordsAsync(new PutRecordsRequest())); +// assertKinesisTrace("putRecords", false); +// } +// +// @Test +// public void testRemoveTagsFromStream() { +// txn(() -> kinesisClient.removeTagsFromStream(new RemoveTagsFromStreamRequest())); +// txnAsync(() -> kinesisAsyncClient.removeTagsFromStreamAsync(new RemoveTagsFromStreamRequest())); +// assertKinesisTrace("removeTagsFromStream", false); +// } +// +// @Test +// public void testSplitShard() { +// txn(() -> kinesisClient.splitShard(new SplitShardRequest())); +// txnAsync(() -> kinesisAsyncClient.splitShardAsync(new SplitShardRequest())); +// assertKinesisTrace("splitShard", false); +// } +// +// @Test +// public void testUpdateShardCount() { +// txn(() -> kinesisClient.updateShardCount(new UpdateShardCountRequest())); +// txnAsync(() -> kinesisAsyncClient.updateShardCountAsync(new UpdateShardCountRequest())); +// assertKinesisTrace("updateShardCount", false); +// } + @Trace(dispatcher = true) public void txn(Runnable runnable) { - runnable.run(); + try { + Thread.sleep(200); + runnable.run(); + } catch (Exception e) { + throw new RuntimeException(e); + } } @Trace(dispatcher = true) - public void txnAsync(Supplier> asyncSupplier) { + public void txnAsync(Supplier> function) { try { - asyncSupplier.get().get(); - } catch (Exception e) { + Thread.sleep(200); + function.get().get(); + } catch (Exception ignored) { } } @@ -85,7 +231,7 @@ public void txnAsync(Supplier> asyncSupplier) { public void txnAsyncWithHandler(Function, Future> function, AsyncHandler handler) { try { function.apply(handler).get(); - } catch (Exception e) { + } catch (Exception ignored) { } } @@ -97,14 +243,13 @@ private void assertKinesisTrace(String kinesisOperation, boolean assertSpan) { List kinesisSpans = introspector.getSpanEvents().stream() .filter(span -> traceName.equals(span.getName())) .collect(Collectors.toList()); - assertFalse(kinesisSpans.isEmpty()); + assertEquals(2, kinesisSpans.size()); for (SpanEvent kinesisSpan: kinesisSpans) { assertEquals("aws_kinesis_data_streams", kinesisSpan.getAgentAttributes().get("cloud.platform")); } } assertTxn(kinesisOperation, introspector); -// assertTxnAsync(kinesisOperation, "OtherTransaction/Custom/com.amazonaws.services.kinesis.AmazonKinesisAPITest/txnAsync", introspector); -// assertTxnAsync(kinesisOperation, "OtherTransaction/Custom/com.amazonaws.services.kinesis.AmazonKinesisAPITest/txnAsyncWithHandler", introspector); + assertTxnAsync(kinesisOperation, "OtherTransaction/Custom/com.amazonaws.services.kinesis.AmazonKinesisAPITest/txnAsync", introspector); } private void assertTxn(String kinesisOperation, Introspector introspector) { @@ -120,17 +265,21 @@ private void assertTxn(String kinesisOperation, Introspector introspector) { } private void assertTxnAsync(String kinesisOperation, String transactionName, Introspector introspector) { - final String traceName = "Kinesis/" + kinesisOperation; + final String asyncClientTraceName = "Java/com.amazonaws.services.kinesis.AmazonKinesisAsyncClient/" + kinesisOperation + "Async"; + final String extTraceName = "Kinesis/" + kinesisOperation; Collection transactionTraces = introspector.getTransactionTracesForTransaction(transactionName); TransactionTrace transactionTrace = transactionTraces.iterator().next(); - List children = transactionTrace.getInitialTraceSegment().getChildren(); - assertEquals(2, children.size()); - TraceSegment asyncFunctionTrace = children.get(0); - String expectedFunctionTraceName = "com.amazonaws.services.kinesis.AmazonKinesisAsyncClient/" + kinesisOperation; - assertEquals(expectedFunctionTraceName, asyncFunctionTrace.getName()); - TraceSegment externalTrace = children.get(1); - assertEquals(traceName, externalTrace.getName()); - assertEquals("aws_kinesis_data_streams", externalTrace.getTracerAttributes().get("cloud.platform")); + + List rootChildren = transactionTrace.getInitialTraceSegment().getChildren(); + assertEquals(1, rootChildren.size()); + TraceSegment asyncClientTrace = rootChildren.get(0); + assertEquals(asyncClientTraceName, asyncClientTrace.getName()); + + List asyncFunctionTraceChildren = asyncClientTrace.getChildren(); + assertEquals(1, asyncFunctionTraceChildren.size()); + TraceSegment extTrace = asyncFunctionTraceChildren.get(0); + assertEquals(extTraceName, extTrace.getName()); + assertEquals("aws_kinesis_data_streams", extTrace.getTracerAttributes().get("cloud.platform")); } private static class CredProvider implements AWSCredentialsProvider { @@ -148,17 +297,4 @@ public void refresh() { } } - // To trigger scenarios where the async handler is called. - private static class AsyncHandlerNoOp implements AsyncHandler { - - @Override - public void onError(Exception exception) { - - } - - @Override - public void onSuccess(REQ request, RES res) { - - } - } } From 3167eb06fc6a66c7fe9fe6b3ae8adcaf4107dcd0 Mon Sep 17 00:00:00 2001 From: Oren Ben-Meir Date: Thu, 29 Aug 2024 16:25:41 -0400 Subject: [PATCH 8/8] Fix settings.gradle repeated module --- settings.gradle | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/settings.gradle b/settings.gradle index 0bcbfdb905..05cade8521 100644 --- a/settings.gradle +++ b/settings.gradle @@ -72,6 +72,8 @@ include 'instrumentation:anorm-2.4' include 'instrumentation:aws-bedrock-runtime-2.20' include 'instrumentation:aws-java-sdk-dynamodb-1.11.106' include 'instrumentation:aws-java-sdk-dynamodb-2.15.34' +include 'instrumentation:aws-java-sdk-kinesis-1.11.106' +include 'instrumentation:aws-java-sdk-kinesis-2.0.6' include 'instrumentation:aws-java-sdk-lambda-1.11.280' include 'instrumentation:aws-java-sdk-lambda-2.1' include 'instrumentation:aws-java-sdk-sqs-1.10.44' @@ -79,10 +81,6 @@ include 'instrumentation:aws-java-sdk-s3-1.2.13' include 'instrumentation:aws-java-sdk-s3-2.0' include 'instrumentation:aws-java-sdk-sns-1.11.12' include 'instrumentation:aws-java-sdk-sqs-2.1.0' -include 'instrumentation:aws-java-sdk-dynamodb-1.11.106' -include 'instrumentation:aws-java-sdk-dynamodb-2.15.34' -include 'instrumentation:aws-java-sdk-kinesis-1.11.106' -include 'instrumentation:aws-java-sdk-kinesis-2.0.6' include 'instrumentation:aws-java-sdk-sns-2.0' include 'instrumentation:aws-wrap-0.7.0' include 'instrumentation:akka-2.2'