Add Vert.x 4 Redis client instrumentation

[akshaypatidar1999]

What Does This Do

Adds instrumentation for Vert.x4 Redis Client
Fixes #6230

Motivation

We are migrating our applications to Vert.x 4 and are not getting any spans for redis

Additional Notes

Jira ticket: [PROJ-IDENT]

[akshaypatidar1999]

There are some API changes in redis client from v3 to v4

• A new send method returning Future is available along with the one that takes the handler

So we need to instrument this method as well but we cannot wrap the response handler like we were doing in the vertx3 as we do not have the actual handler yet (set later using Future.onComplete)

If we do not wrap the handler and just close the parent continuation, spans if any started in the response handler of redis query are not attached to the parent span due to which the runWithParentAndHandler fails.

Ref: https://github.com/akshaypatidar1999/dd-trace-java/blob/feat/vertx-redis-instrumentation/dd-java-agent/instrumentation/vertx-redis-client-4.0/src/main/java/datadog/trace/instrumentation/vertx_redis_client_4/ResponseHandlerWrapper.java#L36

https://github.com/akshaypatidar1999/dd-trace-java/blob/feat/vertx-redis-instrumentation/dd-java-agent/instrumentation/vertx-redis-client-4.0/src/test/groovy/Vertx4RedisRxForkedTest.groovy#L36

@nayeem-kamal @bm1549 @PerfectSlayer Do you have any suggestions about what can be done for instrumenting methods returning Future? Has this been done for any other framework?

[PerfectSlayer]

Hello @akshaypatidar1999

About handling the Future return type, you will need to compose it to close the span only when the Future is completed.
Here is how you can achieve this in the public static void afterSend() method:

    if (clientScope != null) {
      responseFuture = responseFuture.compose(
          response -> {
            clientScope.close();
            return Future.succeededFuture(response);
          },
          throwable -> {
            DECORATE.onError(clientScope, throwable);
            return Future.failedFuture(throwable);
          });
    }

Make sure to update its responseFuture parameter to allow to update it as this code with return the newly created composed Future:

@Advice.Return(readOnly = false) Future<Response> responseFuture,

[akshaypatidar1999]

Thanks @PerfectSlayer for your response.

I am currently using Futute.onComplete and using Future.compose will also not work.

If we use Future.compose like you mentioned then the scope will be closed as soon as the redis server returns the response. The handlers added to the composed Future will be called after that, and again spans started in the response handler of redis query will not be attached to the parent span

[PerfectSlayer]

Sorry @akshaypatidar1999, you are right. My proposal will only take care of span duration, but any span created from a composed future based on the result won't be related to the parent span.

To achieve such thing, we need to capture a Continuation from the parent span (not the one created at method entry, but its parent). It is currently done and stored on the ContextStore. Then on method exit, after closing the client span, we need to activate the parent span continuation. This will create a new scope, that needs to be closed when the Future (and anything composed from it) ends.

As the vert.x Future API does not provide such API (a compose but run as last), and we don't have the Handler running the Future, I went with creating a new vert.x Promise:

   Pair<Boolean, AgentScope.Continuation> pair =
        InstrumentationContext.get(Request.class, Pair.class).get(request);
    AgentScope.Continuation continuation = pair.getRight();
    DECORATE.logging("Parent continuation", continuation);
    if (clientScope != null) {
      Promise<Response> promise = Promise.promise();
      responseFuture.onComplete(event -> {
        AgentScope scope = null;
        try {
          // Close client scope and span
          if (!event.succeeded()) {
            DECORATE.onError(clientScope, event.cause());
          }
          clientScope.close();
          clientScope.span().finish();
          // Activate parent continuation and trigger promise completion
          if (continuation != null) {
            scope = continuation.activate();
          }
          if (event.succeeded()) {
            promise.complete(event.result());
          } else {
            promise.fail(event.cause());
          }
        } finally {
          // Deactivate parent continuation
          if (scope != null) {
            scope.close();
          }
        }
      });
      responseFuture = promise.future();
    }

Creating a new promise should ensure to properly close the parent span continuation after any composed Future from the return Future<Result>.

Note: I think you can also simplify the ContextStore usage by replacing the Boolean with Continuation instead of Pair<Boolean, Continuation>.

[akshaypatidar1999]

Thanks, @PerfectSlayer, using Promise works. I have updated the PR with the changes and added more tests as well. Can you please review it?

I am using Pair<Boolean, Continuation> instead of Continuation because I do not want to wrap the handler twice (because this instrumentation will be applied to Redis, RedisStandaloneConnection). The Boolean is used to indicate that we have wrapped the handler (if not null then we have wrapped the handler) and Continuation is used for storing parent spans.

If we useContinuation then we do not have a way of knowing whether the handler is wrapped or not, since the Conituation can be null in both cases.

Let me know if I am missing something.

[PerfectSlayer]

Looks like this stub is using Vert.X Redis client 3.9, not 4. Isn’t an issue?
Does the stub project need to be copied or can it be a shared module between 3.9 an 4 instrumentation?

[akshaypatidar1999]

It can be a shared module between 3.9 and 4 instrumentation, vertx 4 just has an extra method

@Override
  public Request arg(boolean arg) {
    return null;
  }

[PerfectSlayer]

pair.hasLeft() already covers pair.getLeft() != null

[PerfectSlayer]

Glad to see you made it work 👍

In order to merge the instrumentation, I think we will need to work on deduplicating code between the 3.9 and 4+ Vert.X Redis client instrumentations. Most of the code is similar so it would be great to create a common module for those parts, maybe be for the tests too.

We might also need to rename the package of the original 3.9 instrumentation, maybe introduce name aliases with version number of both instrumentations.
For all this kind of changes, we are looking for someone internally to help you with them and contribute too, so we can push your PR to the last miles to get it to our quality requirements (code maintenance, tests). cc @bm1549

About using a Pair with the InstrumentationContext, I would rather having two different context instances, one with the Boolean, another with the Continuation if they are not related to each other.

A last question for you could be: are you aware of another methods to instrument in the Redis Client 4 to get it fully covered, or do you think we got everything everyone could need?

[akshaypatidar1999]

A last question for you could be: are you aware of another methods to instrument in the Redis Client 4 to get it fully covered, or do you think we got everything everyone could need?

There are three ways to interact with Redis Client 4 i.e RedisConnection, Redis, RedisAPI, all of these eventually use RedisStandaloneConnection.send(). As we have added instrumentation for these classes I think we have it fully covered.

I will be happy to make the changes to meet the quality requirements.

About using a Pair with the InstrumentationContext, I also would have preferred two different context instances but we need them for the same Request instance, we cannot have more than one value for a key in the context store, so went ahead with Pair.

@Override
public Map<String, String> contextStore() {
  Map<String, String> contextStores = new HashMap<>();
  contextStores.put("io.vertx.redis.client.Command", UTF8BytesString.class.getName());
  contextStores.put("io.vertx.redis.client.Request", "datadog.trace.api.Pair");
  contextStores.put("io.vertx.redis.client.RedisConnection", "io.vertx.core.net.SocketAddress");
  return contextStores;
}

[PerfectSlayer]

About using a Pair with the InstrumentationContext, I also would have preferred two different context instances but we need them for the same Request instance, we cannot have more than one value for a key in the context store, so went ahead with Pair.

Sorry, you were right. You can only have one type according a specific key (per Instrumenter).
So there is a way to only use Continuation instead of Pair to avoid allocating many objects but it gets tricky.

It consists in using a specific instance of Continuation, a noop instance for simplicity, as a marker of already instrumented Request (instead of setting a boolean to true). So:

• If no instance is stored (currently meaning no Pair present), you can set the parent continuation if present, or the noop instance otherwise
• If an instance is store (currently a pair is present),
  • If the instance is the noop instance (currently pair.right is null), it was instrumented but no parent continuation,
  • Otherwise (currently pair.right is not null), it is parent continuation to activate.

You can create a such noop Continuation instance using:

/** A marker to denote already instrumented Request (as a noop continuation to fit into the ContextStore). */
private static AgentScope.Continuation INSTRUMENTED_REQUEST_MARKER = AgentTracer.NoopAgentScope.INSTANCE.capture();

[PerfectSlayer]

Hi @akshaypatidar1999

I had a quick chat with @amarziali Yesterday to explain him the context, what's remaining to get it merge and hand him over the PR. So don't be surprised if you start seeing commits pushed to the PR from him 😉

[akshaypatidar1999]

Thanks @PerfectSlayer

[amarziali]

Hi @akshaypatidar1999 I'm working on refactoring and removing some code duplication. Some of the duplicated advices for the 4.x will probably be removed since the advices from 3.9 can still be applied (with some slight modification). I'm finalising the refactoring (that can take still some little time) and adding tests as well for the new future based methods. So please don't be surprised if you notice that some of the classes on the 4.0 module will be deleted.

[amarziali]

@akshaypatidar1999 I finally managed to shrink everything in a module since the most of the advices were compatible with the 4.x and adapted the advice you provided for the future for the 4.x code. I added test suites for 3.9.0, 4.0.0, and 4.x (not testing latest 3.x since the branch is not alive). Also remove Flaky annotation (hence fixing the tests) since otherwise they were not run as mandatory on the CI. The work is pretty ready. @PerfectSlayer when you can review I'm fine with the PR thanks

[PerfectSlayer]

Looks great! Nice improvement 👍
Just one comment about a FIXME.

[amarziali]

thanks for the contribution @akshaypatidar1999 (and @PerfectSlayer for the review)