Merge pull request #1883 from newrelic/1780_subscribeon_fix

Resolves various thread hop problems with Spring Reactor call chaining.
newrelic · May 15, 2024 · 2e42bc8 · 2e42bc8
2 parents 1b74085 + 0315d52
commit 2e42bc8
Show file tree

Hide file tree

Showing 4 changed files with 157 additions and 1 deletion.
diff --git a/...y-reactor-0.9.0/src/main/java/reactor/core/publisher/FluxMapFuseable_Instrumentation.java b/...y-reactor-0.9.0/src/main/java/reactor/core/publisher/FluxMapFuseable_Instrumentation.java
@@ -1,6 +1,6 @@
 /*
  *
- *  * Copyright 2020 New Relic Corporation. All rights reserved.
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
  *  * SPDX-License-Identifier: Apache-2.0
  *
  */

diff --git a/...y-reactor-0.9.0/src/main/java/reactor/core/publisher/MonoSubscribeOn_Instrumentation.java b/...y-reactor-0.9.0/src/main/java/reactor/core/publisher/MonoSubscribeOn_Instrumentation.java
@@ -0,0 +1,49 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package reactor.core.publisher;
+
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Token;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.NewField;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.WeaveAllConstructors;
+import com.newrelic.api.agent.weaver.Weaver;
+
+@Weave(type = MatchType.ExactClass, originalName = "reactor.core.publisher.MonoSubscribeOn")
+abstract class MonoSubscribeOn_Instrumentation {
+
+    @Weave(type = MatchType.ExactClass, originalName = "reactor.core.publisher.MonoSubscribeOn$SubscribeOnSubscriber")
+    static final class SubscribeOnSubscriber_Instrumentation {
+        @NewField
+        private Token token;
+
+        @WeaveAllConstructors
+        SubscribeOnSubscriber_Instrumentation() {
+            if (NewRelic.getAgent().getTransaction() != null && token == null) {
+                token = NewRelic.getAgent().getTransaction().getToken();
+            }
+        }
+
+        public void run () {
+            if (token != null) {
+                Boolean result = token.linkAndExpire();
+                token = null;
+            }
+            Weaver.callOriginal();
+        }
+
+        public void cancel () {
+            if (token != null) {
+                token.linkAndExpire();
+                token = null;
+            }
+            Weaver.callOriginal();
+        }
+    }
+}
diff --git a/...n/reactor-3.3.0/src/main/java/reactor/core/publisher/FluxMapFuseable_Instrumentation.java b/...n/reactor-3.3.0/src/main/java/reactor/core/publisher/FluxMapFuseable_Instrumentation.java
@@ -0,0 +1,58 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package reactor.core.publisher;
+
+import com.newrelic.agent.bridge.AgentBridge;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Token;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.NewField;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.WeaveAllConstructors;
+import com.newrelic.api.agent.weaver.Weaver;
+
+@Weave(type = MatchType.ExactClass, originalName = "reactor.core.publisher.FluxMapFuseable")
+abstract class FluxMapFuseable_Instrumentation {
+
+    @Weave(type = MatchType.ExactClass, originalName = "reactor.core.publisher.FluxMapFuseable$MapFuseableSubscriber")
+    static final class MapFuseableSubscriber_Instrumentation<T, R> {
+        @NewField
+        private Token token;
+
+        @WeaveAllConstructors
+        MapFuseableSubscriber_Instrumentation() {
+            if (AgentBridge.getAgent().getTransaction(false) != null && token == null) {
+                token = NewRelic.getAgent().getTransaction().getToken();
+            }
+        }
+
+        public void onComplete() {
+            if (token != null) {
+                token.linkAndExpire();
+                token = null;
+            }
+            Weaver.callOriginal();
+        }
+
+        public void onNext(T t) {
+            if (token != null) {
+                token.linkAndExpire();
+                token = null;
+            }
+            Weaver.callOriginal();
+        }
+
+        public void onError(Throwable t) {
+            if (token != null) {
+                token.linkAndExpire();
+                token = null;
+            }
+            Weaver.callOriginal();
+        }
+    }
+}
diff --git a/...n/reactor-3.3.0/src/main/java/reactor/core/publisher/MonoSubscribeOn_Instrumentation.java b/...n/reactor-3.3.0/src/main/java/reactor/core/publisher/MonoSubscribeOn_Instrumentation.java
@@ -0,0 +1,49 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package reactor.core.publisher;
+
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Token;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.NewField;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.WeaveAllConstructors;
+import com.newrelic.api.agent.weaver.Weaver;
+
+@Weave(type = MatchType.ExactClass, originalName = "reactor.core.publisher.MonoSubscribeOn")
+abstract class MonoSubscribeOn_Instrumentation {
+
+    @Weave(type = MatchType.ExactClass, originalName = "reactor.core.publisher.MonoSubscribeOn$SubscribeOnSubscriber")
+    static final class SubscribeOnSubscriber_Instrumentation {
+        @NewField
+        private Token token;
+
+        @WeaveAllConstructors
+        SubscribeOnSubscriber_Instrumentation() {
+            if (NewRelic.getAgent().getTransaction() != null && token == null) {
+                token = NewRelic.getAgent().getTransaction().getToken();
+            }
+        }
+
+        public void run () {
+            if (token != null) {
+                Boolean result = token.linkAndExpire();
+                token = null;
+            }
+            Weaver.callOriginal();
+        }
+
+        public void cancel () {
+            if (token != null) {
+                token.linkAndExpire();
+                token = null;
+            }
+            Weaver.callOriginal();
+        }
+    }
+}