Fix: Spring 7 Webflux NoSuchElementError on HttpHeaders#entrySet()

CHANGELOG.next-release.md

@@ -11,6 +11,7 @@ This file contains all changes which are not released yet.
 <!--FIXES-START-->
 
 - avoid caching when reading version from jar to prevent side effects - [#4543](https://github.com/elastic/apm-agent-java/pull/4543)
+- fix Spring Webflux 7 NoSuchMethodError on HttpHeaders#entrySet()
 
 <!--FIXES-END-->
 # Features and enhancements

apm-agent-core/src/test/java/co/elastic/apm/agent/testutils/Java17OnlyTest.java

@@ -19,7 +19,7 @@
 package co.elastic.apm.agent.testutils;
 
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.condition.EnabledForJreRange;
+import org.junit.jupiter.api.condition.EnabledOnJre;
 import org.junit.jupiter.api.condition.JRE;
 import org.junit.platform.launcher.Launcher;
 import org.junit.platform.launcher.LauncherDiscoveryRequest;
@@ -32,24 +32,24 @@
 import static org.junit.platform.engine.discovery.DiscoverySelectors.selectClass;
 
 /**
- * Why not just put @EnabledForJreRange directly on the test classes?
+ * Why not just put @EnabledOnJre directly on the test classes?
  * JUnit reflectively loads the target class when discovering tests.
  * Because of spring, the tests contain references/annotations compiled with Java 17.
- * This in turn leads to UnsupportedClassVersionErrors before JUnit can evaluate the @EnableForJRERange when running on older java versions (e.g. 11).
+ * This in turn leads to UnsupportedClassVersionErrors before JUnit can evaluate the @EnabledOnJre when running on older java versions (e.g. 11).
  * <p>
  * Therefore, this class can be used to wrap tests, as it programmatically triggers the test execution.
  * The actual test implementation should not be named *Test to not be discovered by the maven surefire plugin.
  */
 public abstract class Java17OnlyTest {
 
-    private Class<?> actualTestClass;
+    private final Class<?> actualTestClass;
 
     public Java17OnlyTest(Class<?> testClazz) {
         this.actualTestClass = testClazz;
     }
 
-    @EnabledForJreRange(min = JRE.JAVA_17)
     @Test
+    @EnabledOnJre(JRE.JAVA_17)
     public void runTests() {
         LauncherDiscoveryRequest request = LauncherDiscoveryRequestBuilder.request()
             .selectors(selectClass(actualTestClass))

apm-agent-plugins/apm-reactor-plugin/src/main/java/co/elastic/apm/agent/reactor/TracedSubscriber.java

@@ -32,6 +32,7 @@
 import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
 import reactor.core.Fuseable;
+import reactor.core.Fuseable.QueueSubscription;
 import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Operators;
 import reactor.util.context.Context;
@@ -82,7 +83,7 @@ public void onSubscribe(Subscription s) {
         boolean hasActivated = doEnter("onSubscribe", context);
         Throwable thrown = null;
         try {
-            subscriber.onSubscribe(s);
+            subscriber.onSubscribe(wrapSubscription(s));
         } catch (Throwable e) {
             thrown = e;
             throw e;
@@ -187,6 +188,78 @@ private void doExit(boolean deactivate, String method, @Nullable TraceState<?> c
         context.deactivate();
     }
 
+    /**
+     * Cancellation does not emit a terminal signal, so observe it explicitly instead of
+     * retaining the context until the TracedSubscriber can be garbage-collected.
+     */
+    @SuppressWarnings("unchecked")
+    private Subscription wrapSubscription(Subscription subscription) {
+
+        if (subscription instanceof QueueSubscription) {
+            return new CancellationAwareQueueSubscription((QueueSubscription<T>) subscription);
+        }
+        return new CancellationAwareSubscription(subscription);
+    }
+
+    private class CancellationAwareSubscription implements Subscription {
+
+        protected final Subscription delegate;
+
+        private CancellationAwareSubscription(Subscription delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public void request(long n) {
+            delegate.request(n);
+        }
+
+        @Override
+        public void cancel() {
+            try {
+                delegate.cancel();
+            } finally {
+                discardIf(true);
+            }
+        }
+    }
+
+    private class CancellationAwareQueueSubscription extends CancellationAwareSubscription
+        implements QueueSubscription<T> {
+
+        private final QueueSubscription<T> delegate;
+
+        private CancellationAwareQueueSubscription(QueueSubscription<T> delegate) {
+            super(delegate);
+            this.delegate = delegate;
+        }
+
+        @Override
+        public int requestFusion(int requestedMode) {
+            return delegate.requestFusion(requestedMode);
+        }
+
+        @Override
+        public T poll() {
+            return delegate.poll();
+        }
+
+        @Override
+        public boolean isEmpty() {
+            return delegate.isEmpty();
+        }
+
+        @Override
+        public int size() {
+            return delegate.size();
+        }
+
+        @Override
+        public void clear() {
+            delegate.clear();
+        }
+    }
+
     private void discardIf(boolean condition) {
         if (!condition) {
             return;

apm-agent-plugins/apm-reactor-plugin/src/test/java/co/elastic/apm/agent/reactor/TracedSubscriberTest.java

@@ -23,18 +23,24 @@
 import co.elastic.apm.agent.impl.transaction.AbstractSpanImpl;
 import co.elastic.apm.agent.impl.transaction.TransactionImpl;
 import co.elastic.apm.agent.tracer.TraceState;
+import org.reactivestreams.Subscription;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import reactor.core.CoreSubscriber;
+import reactor.core.Fuseable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
+import reactor.util.context.Context;
 
 import javax.annotation.Nullable;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Predicate;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -223,6 +229,87 @@ void contextPropagation_Flux_error() {
             .verifyErrorMatches(t -> t == error);
     }
 
+    @Test
+    void cancelledSubscription_releasedContext() {
+        transaction = startTestRootTransaction("root");
+        int initialReferenceCount = transaction.getReferenceCount();
+        AtomicReference<Subscription> downstreamSubscription = new AtomicReference<>();
+
+        CoreSubscriber<Integer> subscriber = new CoreSubscriber<>() {
+            @Override
+            public void onSubscribe(Subscription subscription) {
+                downstreamSubscription.set(subscription);
+            }
+
+            @Override
+            public void onNext(Integer value) {
+            }
+
+            @Override
+            public void onError(Throwable throwable) {
+            }
+
+            @Override
+            public void onComplete() {
+            }
+
+            @Override
+            public Context currentContext() {
+                return Context.empty();
+            }
+        };
+
+        TracedSubscriber<Integer> tracedSubscriber = new TracedSubscriber<>(subscriber, tracer, transaction);
+        assertThat(transaction.getReferenceCount()).isEqualTo(initialReferenceCount + 1);
+
+        AtomicBoolean cancelled = new AtomicBoolean();
+        Fuseable.QueueSubscription<Integer> upstreamSubscription = new Fuseable.QueueSubscription<>() {
+            @Override
+            public int requestFusion(int requestedMode) {
+                return Fuseable.SYNC;
+            }
+
+            @Override
+            public Integer poll() {
+                return null;
+            }
+
+            @Override
+            public boolean isEmpty() {
+                return true;
+            }
+
+            @Override
+            public int size() {
+                return 0;
+            }
+
+            @Override
+            public void clear() {
+            }
+
+            @Override
+            public void request(long n) {
+            }
+
+            @Override
+            public void cancel() {
+                cancelled.set(true);
+            }
+        };
+
+        tracedSubscriber.onSubscribe(upstreamSubscription);
+
+        Subscription wrappedSubscription = downstreamSubscription.get();
+        assertThat(wrappedSubscription).isInstanceOf(Fuseable.QueueSubscription.class);
+        assertThat(((Fuseable.QueueSubscription<?>) wrappedSubscription).requestFusion(Fuseable.ANY)).isEqualTo(Fuseable.SYNC);
+
+        wrappedSubscription.cancel();
+
+        assertThat(cancelled).isTrue();
+        assertThat(transaction.getReferenceCount()).isEqualTo(initialReferenceCount);
+    }
+
     @Test
     void ignoreNoActiveContext() {
         assertThat(tracer.getActive()).isNull();

apm-agent-plugins/apm-spring-webflux/apm-spring-webflux-plugin-spring7/pom.xml

@@ -0,0 +1,96 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <parent>
+        <groupId>co.elastic.apm</groupId>
+        <artifactId>apm-spring-webflux</artifactId>
+        <version>1.56.1-SNAPSHOT</version>
+    </parent>
+
+    <artifactId>apm-spring-webflux-plugin-spring7</artifactId>
+    <name>${project.groupId}:${project.artifactId}</name>
+
+    <properties>
+        <!-- for licence header plugin -->
+        <apm-agent-parent.base.dir>${project.basedir}/../../..</apm-agent-parent.base.dir>
+
+        <animal.sniffer.skip>true</animal.sniffer.skip>
+    </properties>
+
+    <dependencyManagement>
+        <dependencies>
+            <dependency>
+                <groupId>org.springframework.boot</groupId>
+                <artifactId>spring-boot-dependencies</artifactId>
+                <version>${version.spring-boot-4}</version>
+                <type>pom</type>
+                <scope>import</scope>
+            </dependency>
+        </dependencies>
+    </dependencyManagement>
+
+    <dependencies>
+        <dependency>
+            <groupId>${project.groupId}</groupId>
+            <artifactId>apm-spring-webflux-spring5</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+
+        <dependency>
+            <!-- through the spring-boot dependency management, the test app will be "updated" to spring boot 4 -->
+            <groupId>${project.groupId}</groupId>
+            <artifactId>apm-spring-webflux-testapp-spring7</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework</groupId>
+            <artifactId>spring-web</artifactId>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.projectreactor</groupId>
+            <artifactId>reactor-test</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>co.elastic.apm</groupId>
+            <artifactId>apm-spring-webflux-spring5</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+            <type>test-jar</type>
+        </dependency>
+
+        <!-- required for context-propagation during tests, but only at runtime -->
+        <dependency>
+            <groupId>co.elastic.apm</groupId>
+            <artifactId>apm-reactor-plugin</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>co.elastic.apm</groupId>
+            <artifactId>apm-reactor-plugin</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+            <type>test-jar</type>
+        </dependency>
+
+        <dependency>
+            <groupId>org.apache.ivy</groupId>
+            <artifactId>ivy</artifactId>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-dependency-plugin</artifactId>
+            </plugin>
+        </plugins>
+    </build>
+
+</project>

apm-agent-plugins/apm-spring-webflux/apm-spring-webflux-plugin-spring7/src/test/java/co/elastic/apm/agent/springwebflux/Spring7HeaderGetterTest.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ */
+package co.elastic.apm.agent.springwebflux;
+
+import co.elastic.apm.agent.testutils.Java17OnlyTest;
+
+public class Spring7HeaderGetterTest extends Java17OnlyTest {
+
+    public Spring7HeaderGetterTest() {
+        super(Impl.class);
+    }
+
+    public static class Impl extends HeaderGetterTest {
+
+    }
+}