Spring webflux 6.1.0 instrumentation module

instrumentation/spring-webflux-6.1.0/NOTICE.txt

@@ -0,0 +1,17 @@
+This product contains a modified part of OpenTelemetry:
+
+  * License:
+
+Copyright 2019 The OpenTelemetry Authors
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+in compliance with the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License
+is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+or implied. See the License for the specific language governing permissions and limitations under
+the License.
+
+  * Homepage: https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/LICENSE

instrumentation/spring-webflux-6.1.0/README.md

@@ -0,0 +1,38 @@
+Spring Webflux Instrumentation
+===============================
+
+This instrumentation assumes that Spring Webflux usage primarily centered
+around maintaining non-blocking flow from the beginning of a request to
+the response rendering completion. If blocking calls occur during the request
+handling, it is possible that the agent will lose transaction context and
+transaction naming will not be able to work as intended.
+
+This is especially likely to happen when using a `ParallelScheduler`
+and `@RequestBody` with a parameter that is NOT wrapped in a `Mono` or `Flux`.
+
+For instance:
+
+```java
+@PostMapping("/path")
+public Mono<String> submit(@RequestBody String body) {
+        ...
+}
+```
+
+When the `@RequestBody` is passed as a parameter in this manner, the
+transaction naming might show up as `NettyDispatcher`. If so, simply wrap
+the `RequestBody` parameter in a `Mono` as shown:
+
+```java
+@PostMapping("/some/path")
+public Mono<String> submit(@RequestBody Mono<String> body) {
+        ...
+}
+```
+
+This should allow the transaction to be named correctly as `some/path (POST)`.
+
+Please see the Spring Webflux documentation for more details on keeping your
+code fully non-blocking.
+
+https://docs.spring.io/spring-framework/docs/current/reference/html/web-reactive.html

instrumentation/spring-webflux-6.1.0/build.gradle

@@ -0,0 +1,40 @@
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation("org.springframework:spring-webflux:6.1.0")
+
+    testImplementation('org.springframework:spring-context:6.1.0')
+    testImplementation("io.projectreactor.netty:reactor-netty-core:1.1.14")
+    testImplementation("io.projectreactor.netty:reactor-netty-http:1.1.14")
+    testImplementation('org.springframework:spring-test:6.1.0')
+    testRuntimeOnly("io.netty:netty-resolver-dns-native-macos:4.1.87.Final:osx-aarch_64")
+
+    testImplementation(project(":instrumentation:spring-webclient-6.0"))
+}
+
+jar {
+    manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.spring-webflux-6.1.0'}
+}
+
+verifyInstrumentation {
+    passesOnly('org.springframework:spring-webflux:[6.1.0,)')
+    excludeRegex '.*.M[0-9]'
+    excludeRegex '.*.RC[0-9]'
+}
+
+java {
+    toolchain {
+        languageVersion.set(JavaLanguageVersion.of(17))
+    }
+}
+
+test {
+    // These instrumentation tests only run on Java 17+ regardless of the -PtestN gradle property that is set.
+    onlyIf {
+        !project.hasProperty('test8') && !project.hasProperty('test11')
+    }
+}
+
+site {
+    title 'Spring webclient'
+    type 'Messaging'
+}

instrumentation/spring-webflux-6.1.0/src/main/java/com/nr/agent/instrumentation/spring/reactive/Util.java

@@ -0,0 +1,141 @@
+/*
+ *
+ *  * Copyright 2020 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.nr.agent.instrumentation.spring.reactive;
+
+import com.newrelic.agent.bridge.Token;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Trace;
+import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscription;
+import org.springframework.web.reactive.function.server.RequestPredicate;
+import org.springframework.web.reactive.function.server.ServerRequest;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.Operators;
+import reactor.util.context.Context;
+
+import java.util.function.Function;
+
+public class Util {
+
+    public static final String NR_TXN_NAME = "newrelic-transaction-name";
+    public static final String NR_TOKEN = "newrelic-token";
+
+    public static <T> Mono<T> setTransactionToken(Mono<T> mono, Token token) {
+        return mono.<T>transform(tokenLift(token));
+    }
+
+    public static <T> Function<? super Mono<T>, ? extends Publisher<T>> tokenLift(
+            Token token) {
+        return Operators.lift(
+                (scannable, subscriber) -> new TokenLinkingSubscriber<T>(subscriber, token));
+    }
+
+    public static class TokenLinkingSubscriber<T> implements CoreSubscriber<T> {
+
+        private final CoreSubscriber<? super T> subscriber;
+        private final Context context;
+
+        public TokenLinkingSubscriber(
+                CoreSubscriber<? super T> subscriber, Token token) {
+            this.subscriber = subscriber;
+            this.context = subscriber.currentContext().put(NR_TOKEN, token);
+        }
+
+        @Override
+        public void onSubscribe(Subscription s) {
+            subscriber.onSubscribe(s);
+        }
+
+        @Override
+        public void onNext(T t) {
+            subscriber.onNext(t);
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            withNRError(() -> subscriber.onError(t), t);
+        }
+
+        @Override
+        public void onComplete() {
+            withNRToken(() -> subscriber.onComplete());
+        }
+
+        @Override
+        public Context currentContext() {
+            return context;
+        }
+
+        @Trace(async = true, excludeFromTransactionTrace = true)
+        private void withNRToken(Runnable runnable) {
+            Token token = currentContext().get(NR_TOKEN);
+            if (token != null) {
+                token.linkAndExpire();
+            }
+            runnable.run();
+        }
+
+        @Trace(async = true, excludeFromTransactionTrace = true)
+        private void withNRError(Runnable runnable, Throwable throwable) {
+            Token token = currentContext().get(NR_TOKEN);
+            if (token != null && token.isActive()) {
+                token.linkAndExpire();
+                NewRelic.noticeError(throwable);
+            }
+            runnable.run();
+        }
+    }
+
+    public static RequestPredicate createRequestPredicate(final String name,
+                                                          final RequestPredicate originalRequestPredicate) {
+        return new RequestPredicate() {
+            @Override
+            public boolean test(ServerRequest request) {
+                final boolean matched = originalRequestPredicate.test(request);
+                if (matched) {
+                    Util.addPath(request, "QueryParameter/" + name);
+                }
+                return matched;
+            }
+
+            @Override
+            public String toString() {
+                return "";
+            }
+        };
+    }
+
+    public static RequestPredicate createPathExtensionPredicate(String extension,
+                                                                RequestPredicate originalRequestPredicate) {
+        return new RequestPredicate() {
+            @Override
+            public boolean test(ServerRequest request) {
+                final boolean matched = originalRequestPredicate.test(request);
+                if (matched) {
+                    Util.addPath(request, "PathExtension/" + extension);
+                }
+                return matched;
+            }
+
+            @Override
+            public String toString() {
+                return "";
+            }
+        };
+    }
+
+    public static void addPath(ServerRequest request, String name) {
+        Token token = (Token) request.attributes().get(NR_TOKEN);
+        if (token != null && !name.isEmpty()) {
+            request.attributes().computeIfAbsent(NR_TXN_NAME, k -> "");
+            String existingName = (String) request.attributes().get(NR_TXN_NAME);
+            request.attributes().put(NR_TXN_NAME, existingName + name);
+        }
+    }
+}

instrumentation/spring-webflux-6.1.0/src/main/java/org/springframework/http/ReactiveHttpOutputMessage_Instrumentation.java

@@ -0,0 +1,34 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package org.springframework.http;
+
+import com.newrelic.agent.bridge.AgentBridge;
+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.Weaver;
+import reactor.core.publisher.Mono;
+
+@Weave(type = MatchType.Interface, originalName = "org.springframework.http.ReactiveHttpOutputMessage")
+public class ReactiveHttpOutputMessage_Instrumentation {
+
+    @NewField
+    public Token token;
+
+    public Mono<Void> setComplete() {
+        try {
+            if (this.token != null) {
+                this.token.expire();
+                this.token = null;
+            }
+        } catch (Throwable t) {
+            AgentBridge.instrumentation.noticeInstrumentationError(t, Weaver.getImplementationTitle());
+        }
+        return Weaver.callOriginal();
+    }
+}

instrumentation/spring-webflux-6.1.0/src/main/java/org/springframework/web/reactive/DefaultServerRequest_Instrumentation.java

@@ -0,0 +1,30 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.springframework.web.reactive;
+
+import com.newrelic.api.agent.Token;
+import com.newrelic.api.agent.weaver.Weave;
+import com.nr.agent.instrumentation.spring.reactive.Util;
+import org.springframework.http.codec.HttpMessageReader;
+import org.springframework.web.server.ServerWebExchange;
+
+import java.util.List;
+import java.util.Map;
+
+@Weave(originalName = "org.springframework.web.reactive.function.server.DefaultServerRequest")
+abstract class DefaultServerRequest_Instrumentation {
+
+    DefaultServerRequest_Instrumentation(ServerWebExchange exchange, List<HttpMessageReader<?>> messageReaders) {
+        final Token token = exchange == null ? null : exchange.getAttribute(Util.NR_TOKEN);
+        if (token != null) {
+            attributes().put(Util.NR_TOKEN, token);
+        }
+    }
+
+    public abstract Map<String, Object> attributes();
+}

instrumentation/spring-webflux-6.1.0/src/main/java/org/springframework/web/reactive/DispatchHandler_Instrumentation.java

@@ -0,0 +1,33 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package org.springframework.web.reactive;
+
+import com.newrelic.agent.bridge.Token;
+import com.newrelic.api.agent.Trace;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import com.nr.agent.instrumentation.spring.reactive.Util;
+import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Mono;
+
+// Based on OpenTelemetry instrumentation
+// https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/spring/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/DispatcherHandlerAdvice.java
+@Weave(type =MatchType.ExactClass, originalName = "org.springframework.web.reactive.DispatcherHandler")
+public class DispatchHandler_Instrumentation {
+    @Trace
+    public Mono<Void> handle(ServerWebExchange exchange) {
+        final Token token = exchange == null ? null : exchange.getAttribute(Util.NR_TOKEN);
+
+        Mono<Void> original = Weaver.callOriginal();
+
+        if (token != null) {
+            return Util.setTransactionToken(original, token);
+        }
+        return original;
+    }
+}

instrumentation/spring-webflux-6.1.0/src/main/java/org/springframework/web/reactive/PathResourceLookupFunction_Instrumentation.java

@@ -0,0 +1,31 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.springframework.web.reactive;
+
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import com.nr.agent.instrumentation.spring.reactive.Util;
+import org.springframework.core.io.Resource;
+import org.springframework.web.reactive.function.server.ServerRequest;
+import org.springframework.web.util.pattern.PathPattern;
+import reactor.core.publisher.Mono;
+
+@Weave(originalName = "org.springframework.web.reactive.function.server.PathResourceLookupFunction")
+class PathResourceLookupFunction_Instrumentation {
+    private final PathPattern pattern = Weaver.callOriginal();
+
+    private final Resource location = Weaver.callOriginal();
+
+    public Mono<Resource> apply(ServerRequest request) {
+        Mono<Resource> result = Weaver.callOriginal();
+        if (!Mono.empty().equals(result)) {
+            Util.addPath(request, pattern.getPatternString());
+        }
+        return result;
+    }
+}