-
Notifications
You must be signed in to change notification settings - Fork 272
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Clean up and expand akka-http client instrumentation and tests
- Loading branch information
1 parent
0bd8ca0
commit 7eac56a
Showing
4 changed files
with
126 additions
and
65 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
File renamed without changes.
87 changes: 87 additions & 0 deletions
87
...ttp-10.0/src/main/java8/datadog/trace/instrumentation/akkahttp/AkkaHttpClientHelpers.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| package datadog.trace.instrumentation.akkahttp; | ||
|
|
||
| import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientDecorator.DECORATE; | ||
|
|
||
| import akka.http.javadsl.model.headers.RawHeader; | ||
| import akka.http.scaladsl.model.HttpRequest; | ||
| import akka.http.scaladsl.model.HttpResponse; | ||
| import akka.http.scaladsl.model.headers.CustomHeader; | ||
| import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; | ||
| import datadog.trace.bootstrap.instrumentation.api.AgentSpan; | ||
| import scala.runtime.AbstractFunction1; | ||
| import scala.util.Try; | ||
|
|
||
| public final class AkkaHttpClientHelpers { | ||
| public static class OnCompleteHandler extends AbstractFunction1<Try<HttpResponse>, Void> { | ||
| private final AgentSpan span; | ||
|
|
||
| public OnCompleteHandler(final AgentSpan span) { | ||
| this.span = span; | ||
| } | ||
|
|
||
| @Override | ||
| public Void apply(final Try<HttpResponse> result) { | ||
| if (result.isSuccess()) { | ||
| DECORATE.onResponse(span, result.get()); | ||
| } else { | ||
| DECORATE.onError(span, result.failed().get()); | ||
| } | ||
| DECORATE.beforeFinish(span); | ||
| span.finish(); | ||
| return null; | ||
| } | ||
| } | ||
|
|
||
| public static class AkkaHttpHeaders implements AgentPropagation.Setter<HttpRequest> { | ||
| private HttpRequest request; | ||
| // Did this request have a span when the AkkaHttpHeaders object was created? | ||
| private boolean hadSpan; | ||
|
|
||
| public AkkaHttpHeaders(final HttpRequest request) { | ||
| hadSpan = request != null && request.getHeader(HasSpanHeader.class).isPresent(); | ||
| if (hadSpan || request == null) { | ||
| this.request = request; | ||
| } else { | ||
| // Coerce a Scala trait Self type into the correct type | ||
| this.request = (HttpRequest) request.addHeader(new HasSpanHeader()); | ||
| } | ||
| } | ||
|
|
||
| public boolean hadSpan() { | ||
| return hadSpan; | ||
| } | ||
|
|
||
| @Override | ||
| public void set(final HttpRequest carrier, final String key, final String value) { | ||
| // Coerce a Scala trait Self type into the correct type | ||
| request = (HttpRequest) request.addHeader(RawHeader.create(key, value)); | ||
| } | ||
|
|
||
| public HttpRequest getRequest() { | ||
| return request; | ||
| } | ||
| } | ||
|
|
||
| // Custom header to mark that this request has a span associated with it | ||
| public static class HasSpanHeader extends CustomHeader { | ||
| @Override | ||
| public String name() { | ||
| return "x-datadog-request-has-span"; | ||
| } | ||
|
|
||
| @Override | ||
| public String value() { | ||
| return "true"; | ||
| } | ||
|
|
||
| @Override | ||
| public boolean renderInRequests() { | ||
| return false; | ||
| } | ||
|
|
||
| @Override | ||
| public boolean renderInResponses() { | ||
| return false; | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters