refactors aerospike module

Author: Abhishekkr3003

instrumentation/aerospike-client/aerospike-client-7.0/javaagent/build.gradle.kts

@@ -12,8 +12,8 @@ muzzle {
 }
 
 dependencies {
-  library("com.aerospike:aerospike-client:7.0.0")
-  implementation("io.opentelemetry:opentelemetry-extension-incubator")
+  implementation("com.aerospike:aerospike-client:7.0.0")
+  implementation("io.opentelemetry:opentelemetry-api-incubator")
 
   compileOnly("com.google.auto.value:auto-value-annotations")
   annotationProcessor("com.google.auto.value:auto-value")

instrumentation/aerospike-client/aerospike-client-7.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/aerospike/v7_0/AerospikeClientAttributeExtractor.java

@@ -10,6 +10,8 @@
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.internal.AerospikeRequest;
+import io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.internal.AerospikeSemanticAttributes;
 import javax.annotation.Nullable;
 
 final class AerospikeClientAttributeExtractor
@@ -18,10 +20,7 @@ final class AerospikeClientAttributeExtractor
   @Override
   public void onStart(
       AttributesBuilder attributes, Context parentContext, AerospikeRequest aerospikeRequest) {
-    attributes.put(
-        AerospikeSemanticAttributes.AEROSPIKE_NAMESPACE, aerospikeRequest.getNamespace());
     attributes.put(AerospikeSemanticAttributes.AEROSPIKE_SET_NAME, aerospikeRequest.getSet());
-    attributes.put(AerospikeSemanticAttributes.AEROSPIKE_USER_KEY, aerospikeRequest.getUserKey());
   }
 
   @Override
@@ -31,22 +30,21 @@ public void onEnd(
       AerospikeRequest aerospikeRequest,
       @Nullable Void unused,
       @Nullable Throwable error) {
-    attributes.put(
-        AerospikeSemanticAttributes.AEROSPIKE_STATUS, aerospikeRequest.getStatus().name());
+    if (aerospikeRequest.getNode() != null) {
+      String nodeName = aerospikeRequest.getNode().getName();
+      attributes.put(AerospikeSemanticAttributes.AEROSPIKE_NODE_NAME, nodeName);
+    }
+
     if (error != null) {
       if (error instanceof AerospikeException) {
         AerospikeException aerospikeException = (AerospikeException) error;
         attributes.put(
-            AerospikeSemanticAttributes.AEROSPIKE_ERROR_CODE, aerospikeException.getResultCode());
+            AerospikeSemanticAttributes.AEROSPIKE_STATUS, aerospikeException.getResultCode());
       } else {
-        attributes.put(AerospikeSemanticAttributes.AEROSPIKE_ERROR_CODE, ResultCode.CLIENT_ERROR);
+        attributes.put(AerospikeSemanticAttributes.AEROSPIKE_STATUS, ResultCode.CLIENT_ERROR);
       }
     } else {
-      attributes.put(AerospikeSemanticAttributes.AEROSPIKE_ERROR_CODE, ResultCode.OK);
-      if (aerospikeRequest.getSize() != null) {
-        attributes.put(
-            AerospikeSemanticAttributes.AEROSPIKE_TRANSFER_SIZE, aerospikeRequest.getSize());
-      }
+      attributes.put(AerospikeSemanticAttributes.AEROSPIKE_STATUS, ResultCode.OK);
     }
   }
 }

instrumentation/aerospike-client/aerospike-client-7.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/aerospike/v7_0/AerospikeClientInstrumentationModule.java

@@ -23,17 +23,14 @@ public AerospikeClientInstrumentationModule() {
 
   @Override
   public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
-    return hasClassesNamed("com.aerospike.client.AerospikeClient")
-        .and(hasClassesNamed("com.aerospike.client.metrics.MetricsPolicy"));
+    return hasClassesNamed("com.aerospike.client.IAerospikeClient");
   }
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
+        new IAerospikeClientInstrumentation(),
         new SyncCommandInstrumentation(),
-        new SocketInstrumentation(),
-        new AsyncCommandInstrumentation(),
-        new AsyncHandlerInstrumentation(),
-        new AsyncScanAllCommandInstrumentation());
+        new AsyncCommandInstrumentation());
   }
 }

instrumentation/aerospike-client/aerospike-client-7.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/aerospike/v7_0/AerospikeDbAttributesGetter.java

@@ -6,38 +6,39 @@
 package io.opentelemetry.javaagent.instrumentation.aerospike.v7_0;
 
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.DbClientAttributesGetter;
+import io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.internal.AerospikeRequest;
+import io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.internal.AerospikeSemanticAttributes;
 import javax.annotation.Nullable;
 
 final class AerospikeDbAttributesGetter implements DbClientAttributesGetter<AerospikeRequest> {
 
   @Override
-  public String getSystem(AerospikeRequest request) {
+  public String getDbSystem(AerospikeRequest request) {
     return AerospikeSemanticAttributes.DbSystemValues.AEROSPIKE;
   }
 
   @Override
-  @Nullable
-  public String getUser(AerospikeRequest request) {
-    return null;
+  public String getDbOperationName(AerospikeRequest request) {
+    return request.getOperation();
   }
 
+  @Nullable
   @Override
-  public String getName(AerospikeRequest request) {
-    return null;
+  public String getDbNamespace(AerospikeRequest request) {
+    return request.getNamespace();
   }
 
+  @Deprecated
+  @Nullable
   @Override
-  public String getConnectionString(AerospikeRequest request) {
+  public String getUser(AerospikeRequest aerospikeRequest) {
     return null;
   }
 
+  @Deprecated
+  @Nullable
   @Override
-  public String getStatement(AerospikeRequest request) {
+  public String getConnectionString(AerospikeRequest aerospikeRequest) {
     return null;
   }
-
-  @Override
-  public String getOperation(AerospikeRequest request) {
-    return request.getOperation();
-  }
 }

instrumentation/aerospike-client/aerospike-client-7.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/aerospike/v7_0/AerospikeNetworkAttributesGetter.java

@@ -7,6 +7,7 @@
 
 import com.aerospike.client.cluster.Node;
 import io.opentelemetry.instrumentation.api.semconv.network.NetworkAttributesGetter;
+import io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.internal.AerospikeRequest;
 import java.net.InetSocketAddress;
 import javax.annotation.Nullable;
 

instrumentation/aerospike-client/aerospike-client-7.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/aerospike/v7_0/AersopikeSingletons.java

@@ -6,14 +6,16 @@
 package io.opentelemetry.javaagent.instrumentation.aerospike.v7_0;
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.instrumentation.api.incubator.config.internal.InstrumentationConfig;
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.DbClientAttributesExtractor;
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.DbClientSpanNameExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.InstrumenterBuilder;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
 import io.opentelemetry.instrumentation.api.semconv.network.NetworkAttributesExtractor;
 import io.opentelemetry.instrumentation.api.semconv.network.NetworkAttributesGetter;
-import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
+import io.opentelemetry.javaagent.bootstrap.internal.AgentInstrumentationConfig;
+import io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.internal.AerospikeRequest;
 import io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.metrics.AerospikeMetrics;
 
 public final class AersopikeSingletons {
@@ -33,7 +35,7 @@ public final class AersopikeSingletons {
                 DbClientSpanNameExtractor.create(aerospikeDbAttributesGetter))
             .addAttributesExtractor(DbClientAttributesExtractor.create(aerospikeDbAttributesGetter))
             .addAttributesExtractor(NetworkAttributesExtractor.create(netAttributesGetter));
-    InstrumentationConfig instrumentationConfig = InstrumentationConfig.get();
+    InstrumentationConfig instrumentationConfig = AgentInstrumentationConfig.get();
     if (instrumentationConfig.getBoolean(
         "otel.instrumentation.aerospike.experimental-span-attributes", false)) {
       builder.addAttributesExtractor(new AerospikeClientAttributeExtractor());

instrumentation/aerospike-client/aerospike-client-7.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/aerospike/v7_0/AsyncCommandInstrumentation.java

@@ -5,21 +5,21 @@
 
 package io.opentelemetry.javaagent.instrumentation.aerospike.v7_0;
 
-import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
-import static io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.AersopikeSingletons.instrumenter;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperClass;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesGenericArguments;
 
-import com.aerospike.client.Key;
+import com.aerospike.client.AerospikeException;
+import com.aerospike.client.cluster.Node;
 import com.aerospike.client.command.Command;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import java.util.Locale;
+import io.opentelemetry.javaagent.instrumentation.aerospike.v7_0.internal.AerospikeRequestContext;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -40,39 +40,87 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         isConstructor().and(isPublic()), this.getClass().getName() + "$ConstructorAdvice");
+
+    transformer.applyAdviceToMethod(
+        isMethod().and(named("getNode")), this.getClass().getName() + "$GetNodeAdvice");
+
+    transformer.applyAdviceToMethod(
+        isMethod().and(named("onSuccess")), this.getClass().getName() + "$GetOnSuccessAdvice");
+
+    transformer.applyAdviceToMethod(
+        isMethod().and(named("onFailure").and(takesGenericArguments(AerospikeException.class))),
+        this.getClass().getName() + "$GetOnFailureAdvice");
   }
 
   @SuppressWarnings("unused")
   public static class ConstructorAdvice {
+
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void onExit(@Advice.This Command command, @Advice.AllArguments Object[] objects) {
-      Key key = null;
-      for (Object object : objects) {
-        if (object instanceof Key) {
-          key = (Key) object;
-          break;
-        }
+    public static void getCommand(@Advice.This Command command) {
+      VirtualField<Command, AerospikeRequestContext> virtualField =
+          VirtualField.find(Command.class, AerospikeRequestContext.class);
+      AerospikeRequestContext requestContext = virtualField.get(command);
+      // If the AerospikeRequestContext is already there in VirtualField then we do not need
+      // to override it when constructor of other subclasses of AsyncCommand executes as
+      // AsyncCommand follows multilevel inheritance.
+      if (requestContext != null) {
+        return;
       }
-      if (key == null) {
+      requestContext = AerospikeRequestContext.current();
+      if (requestContext == null) {
         return;
       }
+      virtualField.set(command, requestContext);
+      requestContext.detachContext();
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class GetNodeAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void getNode(@Advice.Return Node node, @Advice.This Command command) {
       VirtualField<Command, AerospikeRequestContext> virtualField =
           VirtualField.find(Command.class, AerospikeRequestContext.class);
       AerospikeRequestContext requestContext = virtualField.get(command);
-      if (requestContext != null) {
+      if (requestContext == null || requestContext.getRequest().getNode() != null) {
         return;
       }
-      Context parentContext = currentContext();
-      AerospikeRequest request =
-          AerospikeRequest.create(command.getClass().getSimpleName().toUpperCase(Locale.ROOT), key);
-      if (!instrumenter().shouldStart(parentContext, request)) {
+      requestContext.getRequest().setNode(node);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class GetOnSuccessAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void stopInstrumentation(@Advice.This Command command) {
+      VirtualField<Command, AerospikeRequestContext> virtualField =
+          VirtualField.find(Command.class, AerospikeRequestContext.class);
+      AerospikeRequestContext requestContext = virtualField.get(command);
+      if (requestContext == null) {
         return;
       }
-      Context context = instrumenter().start(parentContext, request);
-      AerospikeRequestContext aerospikeRequestContext =
-          AerospikeRequestContext.attach(request, context);
+      virtualField.set(command, null);
+      requestContext.endSpan();
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class GetOnFailureAdvice {
 
-      virtualField.set(command, aerospikeRequestContext);
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void stopInstrumentation(
+        @Advice.Argument(0) AerospikeException ae, @Advice.This Command command) {
+      VirtualField<Command, AerospikeRequestContext> virtualField =
+          VirtualField.find(Command.class, AerospikeRequestContext.class);
+      AerospikeRequestContext requestContext = virtualField.get(command);
+      if (requestContext == null) {
+        return;
+      }
+      virtualField.set(command, null);
+      requestContext.setThrowable(ae);
+      requestContext.endSpan();
     }
   }
 }