Add suggestions

Author: amarziali

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java

@@ -19,6 +19,7 @@
 import datadog.trace.api.time.SystemTimeSource;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
+import datadog.trace.bootstrap.instrumentation.api.NotSampledSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.SpanAttributes;
 import datadog.trace.bootstrap.instrumentation.api.SpanLink;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
@@ -163,7 +164,13 @@ private AgentSpan onFrameStart(
         // the link is not added if the user wants to have receive frames on the same trace as the
         // handshake
         wsSpan.addLink(
-            SpanLink.from(handshakeSpan.context(), (byte) 0, "", linkAttributes, inheritSampling));
+            SpanLink.from(
+                inheritSampling
+                    ? handshakeSpan.context()
+                    : new NotSampledSpanContext(handshakeSpan.context()),
+                SpanLink.DEFAULT_FLAGS,
+                "",
+                linkAttributes));
       }
     }
     return wsSpan;

internal-api/build.gradle

@@ -96,6 +96,7 @@ excludedClassesCoverage += [
   "datadog.trace.bootstrap.instrumentation.api.NoopScope",
   "datadog.trace.bootstrap.instrumentation.api.NoopSpan",
   "datadog.trace.bootstrap.instrumentation.api.NoopSpanContext",
+  "datadog.trace.bootstrap.instrumentation.api.NotSampledSpanContext",
   "datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities",
   "datadog.trace.bootstrap.instrumentation.api.Schema",
   "datadog.trace.bootstrap.instrumentation.api.ScopeSource",

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NotSampledSpanContext.java

@@ -0,0 +1,45 @@
+package datadog.trace.bootstrap.instrumentation.api;
+
+import datadog.trace.api.DDTraceId;
+import datadog.trace.api.datastreams.PathwayContext;
+import datadog.trace.api.sampling.PrioritySampling;
+import java.util.Map;
+
+/** An {@link AgentSpanContext} that hides the sampling priority. */
+public final class NotSampledSpanContext implements AgentSpanContext {
+  private final AgentSpanContext delegate;
+
+  public NotSampledSpanContext(AgentSpanContext delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public DDTraceId getTraceId() {
+    return delegate.getTraceId();
+  }
+
+  @Override
+  public long getSpanId() {
+    return delegate.getSpanId();
+  }
+
+  @Override
+  public AgentTraceCollector getTraceCollector() {
+    return delegate.getTraceCollector();
+  }
+
+  @Override
+  public int getSamplingPriority() {
+    return PrioritySampling.UNSET;
+  }
+
+  @Override
+  public Iterable<Map.Entry<String, String>> baggageItems() {
+    return delegate.baggageItems();
+  }
+
+  @Override
+  public PathwayContext getPathwayContext() {
+    return delegate.getPathwayContext();
+  }
+}

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanLink.java

@@ -48,28 +48,7 @@ public static SpanLink from(AgentSpanContext context) {
    */
   public static SpanLink from(
       AgentSpanContext context, byte traceFlags, String traceState, SpanAttributes attributes) {
-    return from(context, traceFlags, traceState, attributes, true);
-  }
-
-  /**
-   * Creates a span link from a span context with W3C trace state and custom attributes. Gathers the
-   * trace and span identifiers from the given instance.
-   *
-   * @param context The context of the span to get the link to.
-   * @param traceFlags The W3C formatted trace flags.
-   * @param traceState The W3C formatted trace state.
-   * @param attributes The link attributes.
-   * @param inheritSampling if true the sampling flag will be set based on the context's sampling
-   *     priority.
-   * @return A span link to the given context.
-   */
-  public static SpanLink from(
-      AgentSpanContext context,
-      byte traceFlags,
-      String traceState,
-      SpanAttributes attributes,
-      boolean inheritSampling) {
-    if (inheritSampling && context.getSamplingPriority() > 0) {
+    if (context.getSamplingPriority() > 0) {
       traceFlags = (byte) (traceFlags | SAMPLED_FLAG);
     }
     return new SpanLink(