Merge pull request #67 from newrelic/segment_attribute

Segment attribute
newrelic · Sep 23, 2020 · b1885ab · b1885ab
2 parents ea17334 + 702e813
commit b1885ab
Show file tree

Hide file tree

Showing 9 changed files with 302 additions and 94 deletions.
diff --git a/agent-bridge/src/main/java/com/newrelic/agent/bridge/NoOpSegment.java b/agent-bridge/src/main/java/com/newrelic/agent/bridge/NoOpSegment.java
@@ -10,6 +10,8 @@
 import com.newrelic.api.agent.ExternalParameters;
 import com.newrelic.api.agent.OutboundHeaders;
 
+import java.util.Map;
+
 public class NoOpSegment implements TracedActivity {
 
     public static final NoOpSegment INSTANCE = new NoOpSegment();
@@ -67,4 +69,20 @@ public void end(){
     public void endAsync() {
     }
 
+    @Override
+    public void addCustomAttribute(String key, Number value) {
+    }
+
+    @Override
+    public void addCustomAttribute(String key, String value) {
+    }
+
+    @Override
+    public void addCustomAttribute(String key, boolean value) {
+    }
+
+    @Override
+    public void addCustomAttributes(Map<String, Object> attributes) {
+    }
+
 }
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/Segment.java b/newrelic-agent/src/main/java/com/newrelic/agent/Segment.java
@@ -7,19 +7,19 @@
 
 package com.newrelic.agent;
 
-import com.newrelic.agent.bridge.NoOpDistributedTracePayload;
 import com.newrelic.agent.bridge.NoOpTracedMethod;
 import com.newrelic.agent.bridge.TracedMethod;
 import com.newrelic.agent.service.ServiceFactory;
 import com.newrelic.agent.tracers.Tracer;
-import com.newrelic.api.agent.DistributedTracePayload;
+import com.newrelic.api.agent.AttributeHolder;
 import com.newrelic.api.agent.ExternalParameters;
 import com.newrelic.api.agent.OutboundHeaders;
 import com.newrelic.api.agent.Transaction;
 
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
 
-public class Segment implements com.newrelic.agent.bridge.TracedActivity {
+public class Segment implements com.newrelic.agent.bridge.TracedActivity, AttributeHolder {
     private volatile Tracer underlyingTracer;
     private volatile Tracer parent;
     private volatile WeakRefTransaction weakRefTransaction;
@@ -149,6 +149,34 @@ public void endAsync() {
         finish(null, true);
     }
 
+    @Override
+    public void addCustomAttribute(String key, Number value) {
+        if (underlyingTracer != null) {
+            underlyingTracer.addCustomAttribute(key, value);
+        }
+    }
+
+    @Override
+    public void addCustomAttribute(String key, String value) {
+        if (underlyingTracer != null) {
+            underlyingTracer.addCustomAttribute(key, value);
+        }
+    }
+
+    @Override
+    public void addCustomAttribute(String key, boolean value) {
+        if (underlyingTracer != null) {
+            underlyingTracer.addCustomAttribute(key, value);
+        }
+    }
+
+    @Override
+    public void addCustomAttributes(Map<String, Object> attributes) {
+        if (underlyingTracer != null) {
+            underlyingTracer.addCustomAttributes(attributes);
+        }
+    }
+
     /**
      * {@inheritDoc}
      */
@@ -169,7 +197,9 @@ private void finish(final Throwable t, boolean async) {
                 Runnable expireSegmentRunnable = new Runnable() {
                     @Override
                     public void run() {
-                        tracer.getTransactionActivity().getTransaction().finishSegment(segment, t, parent, endThreadName);
+                        tracer.getTransactionActivity()
+                                .getTransaction()
+                                .finishSegment(segment, t, parent, endThreadName);
 
                         // Remove references to underlying and parent tracer to prevent GC issues
                         underlyingTracer = null;
@@ -195,7 +225,8 @@ public String getInitiatingThread() {
     public void setTruncated() {
         Tracer tracer = underlyingTracer;
         if (tracer != null) {
-            tracer.setMetricNameFormatInfo(tracer.getMetricName(), "Truncated/" + tracer.getMetricName(), tracer.getTransactionSegmentUri());
+            tracer.setMetricNameFormatInfo(tracer.getMetricName(), "Truncated/" + tracer.getMetricName(),
+                    tracer.getTransactionSegmentUri());
         }
     }
 

diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/attributes/AttributeValidator.java b/newrelic-agent/src/main/java/com/newrelic/agent/attributes/AttributeValidator.java
@@ -16,7 +16,11 @@
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.nio.charset.StandardCharsets;
-import java.util.*;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
@@ -34,7 +38,8 @@ private String getAttributeType() {
     private boolean isTransactional = true;
 
     // Set of APIs that can report attributes outside of a transaction
-    private static final Collection<String> sendParametersOutsideOfTxn = Arrays.asList("noticeError", "Span.addCustomParameter", "Span.addCustomParameters");
+    private static final Collection<String> sendParametersOutsideOfTxn = Arrays.asList("noticeError",
+            "Span.addCustomParameter", "Span.addCustomParameters", "TracedMethod.addCustomAttributes");
 
     public AttributeValidator(String attributeType) {
         this.attributeType = attributeType;
@@ -133,7 +138,8 @@ protected Map<String, Object> verifyParametersAndReturnValues(Map<String, Object
     }
 
     private boolean isAcceptableValueType(Object value) {
-        if (!(value instanceof String) && !(value instanceof Number) && !(value instanceof Boolean) && !(value instanceof AtomicBoolean)) {
+        if (!(value instanceof String) && !(value instanceof Number) && !(value instanceof Boolean) &&
+                !(value instanceof AtomicBoolean)) {
             return false;
         }
         return true;
@@ -168,7 +174,7 @@ private String truncateValue(String key, String value, String methodCalled) {
      * This function truncates a Unicode String so that it can be encoded in maxBytes. It uses the UTF-8 encoding to
      * determine where the String has to be truncated.
      *
-     * @param s        String to be truncated
+     * @param s String to be truncated
      * @param maxBytes Maximum number of bytes in UTF-8 charset encoding
      * @return
      */

diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/tracers/AbstractTracer.java b/newrelic-agent/src/main/java/com/newrelic/agent/tracers/AbstractTracer.java
@@ -8,6 +8,7 @@
 package com.newrelic.agent.tracers;
 
 import com.newrelic.agent.Agent;
+import com.newrelic.api.agent.AttributeHolder;
 import com.newrelic.agent.MetricNames;
 import com.newrelic.agent.Transaction;
 import com.newrelic.agent.TransactionActivity;
@@ -36,7 +37,7 @@
 /**
  * Base class for all tracers. This implements {@link InvocationHandler#invoke(Object, Method, Object[])}
  */
-public abstract class AbstractTracer implements Tracer {
+public abstract class AbstractTracer implements Tracer, AttributeHolder {
 
     static final int INITIAL_PARAMETER_MAP_SIZE = 5;
     protected static String ATTRIBUTE_TYPE = "custom";
@@ -54,7 +55,7 @@ public abstract class AbstractTracer implements Tracer {
 
     private final long startTimeInMillis;
     AtomicReference<Long> finishTime = new AtomicReference<>(null);
-    private final String ATTRIBUTE_API_METHOD_NAME = "TracedMethod addCustomAttributes";
+    private final String ATTRIBUTE_API_METHOD_NAME = "TracedMethod.addCustomAttributes";
 
     // Tracers MUST NOT store references to the Transaction. Why: tracers are stored in the TransactionActivity,
     // and Activities can be reparented from one Transaction to another by the public APIs that support async.