From 1c5dab84631a7faf5ddffa6cfcd64ead30c1d776 Mon Sep 17 00:00:00 2001
From: wallezhang <doraemon.zh@gmail.com>
Date: Thu, 10 Jun 2021 13:11:00 +0800
Subject: [PATCH] refactor: EntrySpan and ExitSpan support modifying attributes
 in async mode (#7086)

---
 CHANGES.md                                          |  2 +-
 .../apm/agent/core/context/trace/EntrySpan.java     |  6 +++---
 .../apm/agent/core/context/trace/ExitSpan.java      | 13 ++++++-------
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 596193157..1e7c8799e 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -8,7 +8,7 @@ Release Notes.
 
 
 #### Java Agent
-
+* Supports modifying span attributes in async mode.
 
 #### OAP-Backend
 
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/EntrySpan.java b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/EntrySpan.java
index 67f5a5039..860e7d349 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/EntrySpan.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/EntrySpan.java
@@ -55,7 +55,7 @@ public class EntrySpan extends StackBasedTracingSpan {
 
     @Override
     public EntrySpan tag(String key, String value) {
-        if (stackDepth == currentMaxDepth) {
+        if (stackDepth == currentMaxDepth || isInAsyncMode) {
             super.tag(key, value);
         }
         return this;
@@ -63,7 +63,7 @@ public class EntrySpan extends StackBasedTracingSpan {
 
     @Override
     public AbstractTracingSpan setLayer(SpanLayer layer) {
-        if (stackDepth == currentMaxDepth) {
+        if (stackDepth == currentMaxDepth || isInAsyncMode) {
             return super.setLayer(layer);
         } else {
             return this;
@@ -72,7 +72,7 @@ public class EntrySpan extends StackBasedTracingSpan {
 
     @Override
     public AbstractTracingSpan setComponent(Component component) {
-        if (stackDepth == currentMaxDepth) {
+        if (stackDepth == currentMaxDepth || isInAsyncMode) {
             return super.setComponent(component);
         } else {
             return this;
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/ExitSpan.java b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/ExitSpan.java
index a69b2b23c..e61a8d1ca 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/ExitSpan.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/ExitSpan.java
@@ -35,6 +35,7 @@ import org.apache.skywalking.apm.network.trace.component.Component;
  * the httpcomponent span's info.
  */
 public class ExitSpan extends StackBasedTracingSpan implements ExitTypeSpan {
+
     public ExitSpan(int spanId, int parentSpanId, String operationName, String peer, TracingContext owner) {
         super(spanId, parentSpanId, operationName, peer, owner);
     }
@@ -56,7 +57,7 @@ public class ExitSpan extends StackBasedTracingSpan implements ExitTypeSpan {
 
     @Override
     public ExitSpan tag(String key, String value) {
-        if (stackDepth == 1) {
+        if (stackDepth == 1 || isInAsyncMode) {
             super.tag(key, value);
         }
         return this;
@@ -64,7 +65,7 @@ public class ExitSpan extends StackBasedTracingSpan implements ExitTypeSpan {
 
     @Override
     public AbstractTracingSpan tag(AbstractTag<?> tag, String value) {
-        if (stackDepth == 1 || tag.isCanOverwrite()) {
+        if (stackDepth == 1 || tag.isCanOverwrite() || isInAsyncMode) {
             super.tag(tag, value);
         }
         return this;
@@ -72,7 +73,7 @@ public class ExitSpan extends StackBasedTracingSpan implements ExitTypeSpan {
 
     @Override
     public AbstractTracingSpan setLayer(SpanLayer layer) {
-        if (stackDepth == 1) {
+        if (stackDepth == 1 || isInAsyncMode) {
             return super.setLayer(layer);
         } else {
             return this;
@@ -81,7 +82,7 @@ public class ExitSpan extends StackBasedTracingSpan implements ExitTypeSpan {
 
     @Override
     public AbstractTracingSpan setComponent(Component component) {
-        if (stackDepth == 1) {
+        if (stackDepth == 1 || isInAsyncMode) {
             return super.setComponent(component);
         } else {
             return this;
@@ -90,9 +91,7 @@ public class ExitSpan extends StackBasedTracingSpan implements ExitTypeSpan {
 
     @Override
     public ExitSpan log(Throwable t) {
-        if (stackDepth == 1) {
-            super.log(t);
-        }
+        super.log(t);
         return this;
     }