Spring-webflux 5.3 instrumentation

instrumentation/spring-webflux-5.3.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/master/LICENSE

instrumentation/spring-webflux-5.3.0/build.gradle

@@ -0,0 +1,40 @@
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation("org.springframework:spring-webflux:5.3.0")
+
+    testImplementation("org.springframework:spring-context:5.3.1")
+    testImplementation("io.projectreactor.netty:reactor-netty:1.0.1")
+
+    testImplementation(project(":instrumentation:spring-webclient-5.0"))
+}
+
+sourceCompatibility = JavaVersion.VERSION_1_8
+targetCompatibility = JavaVersion.VERSION_1_8
+
+compileTestJava {
+    sourceCompatibility = JavaVersion.VERSION_1_8
+    targetCompatibility = JavaVersion.VERSION_1_8
+}
+
+compileJava.options.bootstrapClasspath = null
+
+jar {
+    manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.spring-webflux-5.3.0'}
+}
+
+verifyInstrumentation {
+    passesOnly('org.springframework:spring-webflux:[5.3.0,)')
+    excludeRegex '.*.M[0-9]'
+    excludeRegex '.*.RC[0-9]'
+}
+
+test {
+    onlyIf {
+        !project.hasProperty('test7')
+    }
+}
+
+site {
+    title 'Spring webclient'
+    type 'Messaging'
+}

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

@@ -0,0 +1,140 @@
+/*
+ *
+ *  * 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-5.3.0/src/main/java/org/springframework/http/ReactiveHttpOutputMessage_Instrumentation.java

@@ -0,0 +1,63 @@
+/*
+ *
+ *  * Copyright 2020 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 org.reactivestreams.Publisher;
+import org.springframework.core.io.buffer.DataBuffer;
+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> writeWith(Publisher<? extends DataBuffer> body) {
+        try {
+            if (this.token != null) {
+                this.token.expire();
+                this.token = null;
+            }
+        } catch (Throwable t) {
+            AgentBridge.instrumentation.noticeInstrumentationError(t, Weaver.getImplementationTitle());
+        }
+
+        return Weaver.callOriginal();
+    }
+
+    public Mono<Void> writeAndFlushWith(
+            Publisher<? extends Publisher<? extends DataBuffer>> body) {
+        try {
+            if (this.token != null) {
+                this.token.expire();
+                this.token = null;
+            }
+        } catch (Throwable t) {
+            AgentBridge.instrumentation.noticeInstrumentationError(t, Weaver.getImplementationTitle());
+        }
+        return Weaver.callOriginal();
+    }
+
+    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-5.3.0/src/main/java/org/springframework/web/reactive/DefaultServerRequest_Instrumentation.java

@@ -0,0 +1,30 @@
+/*
+ *
+ *  * Copyright 2020 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-5.3.0/src/main/java/org/springframework/web/reactive/DispatchHandler_Instrumentation.java

@@ -0,0 +1,27 @@
+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/master/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-5.3.0/src/main/java/org/springframework/web/reactive/PathResourceLookupFunction_Instrumentation.java

@@ -0,0 +1,31 @@
+/*
+ *
+ *  * Copyright 2020 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;
+    }
+}