Unify future operation handling

Author: steverao

instrumentation/spymemcached-2.12/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spymemcached/BulkGetCompletionListener.java

@@ -9,6 +9,7 @@
 
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
+import java.util.Collection;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nullable;
 import net.spy.memcached.MemcachedConnection;
@@ -39,13 +40,12 @@ private static MemcachedNode getRepresentativeNodeFromConnection(MemcachedConnec
       // Strategy: Get the "most representative" node for bulk operations
       // We choose the last active node in the list, which often represents
       // the most recently added or most stable node in the cluster
-      java.util.Collection<net.spy.memcached.MemcachedNode> allNodes =
-          connection.getLocator().getAll();
+      Collection<MemcachedNode> allNodes = connection.getLocator().getAll();
 
       MemcachedNode lastActiveNode = null;
       MemcachedNode fallbackNode = null;
 
-      for (net.spy.memcached.MemcachedNode node : allNodes) {
+      for (MemcachedNode node : allNodes) {
         if (fallbackNode == null) {
           fallbackNode = node;
         }

instrumentation/spymemcached-2.12/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spymemcached/GetCompletionListener.java

@@ -5,7 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.spymemcached;
 
-import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.GET_FUTURE_OPERATION;
+import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.FUTURE_OPERATION;
 import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.instrumenter;
 
 import io.opentelemetry.api.trace.Span;
@@ -30,7 +30,6 @@ public static GetCompletionListener create(
       MemcachedConnection connection,
       String methodName,
       GetFuture<?> future) {
-    // Extract handling node from future before creating span
     MemcachedNode handlingNode = extractHandlingNodeFromFuture(future);
     SpymemcachedRequest request = SpymemcachedRequest.create(connection, methodName, handlingNode);
     if (!instrumenter().shouldStart(parentContext, request)) {
@@ -41,7 +40,7 @@ public static GetCompletionListener create(
 
   @Nullable
   private static MemcachedNode extractHandlingNodeFromFuture(GetFuture<?> future) {
-    Operation operation = GET_FUTURE_OPERATION.get(future);
+    Operation operation = FUTURE_OPERATION.get(future);
     if (operation != null) {
       return operation.getHandlingNode();
     }

instrumentation/spymemcached-2.12/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spymemcached/OperationCompletionListener.java

@@ -5,7 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.spymemcached;
 
-import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.OPERATION_FUTURE_OPERATION;
+import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.FUTURE_OPERATION;
 import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.instrumenter;
 
 import io.opentelemetry.api.trace.Span;
@@ -30,7 +30,6 @@ public static OperationCompletionListener create(
       MemcachedConnection connection,
       String methodName,
       OperationFuture<?> future) {
-    // Extract handling node from future before creating span
     MemcachedNode handlingNode = extractHandlingNodeFromFuture(future);
     SpymemcachedRequest request = SpymemcachedRequest.create(connection, methodName, handlingNode);
     if (!instrumenter().shouldStart(parentContext, request)) {
@@ -41,7 +40,7 @@ public static OperationCompletionListener create(
 
   @Nullable
   private static MemcachedNode extractHandlingNodeFromFuture(OperationFuture<?> future) {
-    Operation operation = OPERATION_FUTURE_OPERATION.get(future);
+    Operation operation = FUTURE_OPERATION.get(future);
     if (operation != null) {
       return operation.getHandlingNode();
     }

instrumentation/spymemcached-2.12/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spymemcached/SetOperationInstrumentation.java

@@ -5,15 +5,15 @@
 
 package io.opentelemetry.javaagent.instrumentation.spymemcached;
 
-import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.GET_FUTURE_OPERATION;
-import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.OPERATION_FUTURE_OPERATION;
+import static io.opentelemetry.javaagent.instrumentation.spymemcached.SpymemcachedSingletons.FUTURE_OPERATION;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.util.concurrent.Future;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -42,10 +42,8 @@ public static class SetOperationAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void onExit(@Advice.This Object future, @Advice.Argument(0) Operation operation) {
 
-      if (future instanceof OperationFuture) {
-        OPERATION_FUTURE_OPERATION.set((OperationFuture<?>) future, operation);
-      } else if (future instanceof GetFuture) {
-        GET_FUTURE_OPERATION.set((GetFuture<?>) future, operation);
+      if (future instanceof OperationFuture || future instanceof GetFuture) {
+        FUTURE_OPERATION.set((Future<?>) future, operation);
       }
     }
   }

instrumentation/spymemcached-2.12/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spymemcached/SpymemcachedSingletons.java

@@ -11,21 +11,17 @@
 import io.opentelemetry.instrumentation.api.incubator.semconv.db.DbClientSpanNameExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
-import io.opentelemetry.instrumentation.api.semconv.network.NetworkAttributesExtractor;
 import io.opentelemetry.instrumentation.api.semconv.network.ServerAttributesExtractor;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
-import net.spy.memcached.internal.GetFuture;
-import net.spy.memcached.internal.OperationFuture;
+import java.util.concurrent.Future;
 import net.spy.memcached.ops.Operation;
 
 public final class SpymemcachedSingletons {
   private static final String INSTRUMENTATION_NAME = "io.opentelemetry.spymemcached-2.12";
 
   private static final Instrumenter<SpymemcachedRequest, Object> INSTRUMENTER;
 
-  public static final VirtualField<OperationFuture<?>, Operation> OPERATION_FUTURE_OPERATION;
-
-  public static final VirtualField<GetFuture<?>, Operation> GET_FUTURE_OPERATION;
+  public static final VirtualField<Future<?>, Operation> FUTURE_OPERATION;
 
   static {
     SpymemcachedAttributesGetter dbAttributesGetter = new SpymemcachedAttributesGetter();
@@ -39,12 +35,10 @@ public final class SpymemcachedSingletons {
                 DbClientSpanNameExtractor.create(dbAttributesGetter))
             .addAttributesExtractor(DbClientAttributesExtractor.create(dbAttributesGetter))
             .addAttributesExtractor(ServerAttributesExtractor.create(netAttributesGetter))
-            .addAttributesExtractor(NetworkAttributesExtractor.create(netAttributesGetter))
             .addOperationMetrics(DbClientMetrics.get())
             .buildInstrumenter(SpanKindExtractor.alwaysClient());
 
-    OPERATION_FUTURE_OPERATION = VirtualField.find(OperationFuture.class, Operation.class);
-    GET_FUTURE_OPERATION = VirtualField.find(GetFuture.class, Operation.class);
+    FUTURE_OPERATION = VirtualField.find(Future.class, Operation.class);
   }
 
   public static Instrumenter<SpymemcachedRequest, Object> instrumenter() {

instrumentation/spymemcached-2.12/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spymemcached/SyncCompletionListener.java

@@ -40,7 +40,6 @@ private static MemcachedNode getHandlingNodeForKey(MemcachedConnection connectio
       // Use the connection's locator to find the primary node for this key
       return connection.getLocator().getPrimary(key);
     } catch (RuntimeException e) {
-      // If we can't determine the node, return null
       return null;
     }
   }