From e5308d711175ddaaa95af7a56e1a504ce46157ff Mon Sep 17 00:00:00 2001
From: Adrian Cole <adriancole@users.noreply.github.com>
Date: Fri, 6 Mar 2020 11:45:32 +0800
Subject: [PATCH] Tracks requested but without signal interference (#1581)

---
 .../client/TraceWebClientBeanPostProcessor.java  | 13 ++++++-------
 .../TraceWebClientBeanPostProcessorTest.java     | 16 ++++++++--------
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java
index 924976843..ce95437d0 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java
@@ -18,7 +18,6 @@ package org.springframework.cloud.sleuth.instrument.web.client;
 
 import java.util.List;
 import java.util.concurrent.CancellationException;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -290,8 +289,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 	}
 
-	static class TraceWebClientSubscription extends AtomicBoolean
-			implements Subscription {
+	static class TraceWebClientSubscription implements Subscription {
 
 		static final Exception CANCELLED_ERROR = new CancellationException("CANCELLED") {
 			@Override
@@ -304,6 +302,8 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 		final Subscription delegate;
 
+		volatile boolean requested;
+
 		TraceWebClientSubscription(Subscription delegate,
 				AtomicReference<Span> pendingSpan) {
 			this.delegate = delegate;
@@ -312,9 +312,8 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 		@Override
 		public void request(long n) {
-			if (compareAndSet(false, true)) {
-				delegate.request(n); // Not scoping to save overhead
-			}
+			requested = true;
+			delegate.request(n); // Not scoping to save overhead
 		}
 
 		@Override
@@ -331,7 +330,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 									+ span + "]");
 				}
 
-				if (!get()) { // Subscription.request() not called: Abandon the span.
+				if (!requested) { // Abandon the span.
 					span.abandon();
 				}
 				else { // Request was canceled in-flight
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessorTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessorTest.java
index 2f67fbad8..d99b7f6c2 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessorTest.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessorTest.java
@@ -20,10 +20,10 @@ import java.util.concurrent.atomic.AtomicReference;
 
 import brave.Span;
 import org.assertj.core.api.BDDAssertions;
-import org.junit.Test;
-import org.junit.runner.RunWith;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnitRunner;
+import org.mockito.junit.jupiter.MockitoExtension;
 import org.reactivestreams.Subscription;
 
 import org.springframework.cloud.sleuth.instrument.web.client.TraceExchangeFilterFunction.TraceWebClientSubscription;
@@ -36,7 +36,7 @@ import static org.mockito.Mockito.verify;
 /**
  * @author Marcin Grzejszczak
  */
-@RunWith(MockitoJUnitRunner.class)
+@ExtendWith(MockitoExtension.class)
 public class TraceWebClientBeanPostProcessorTest {
 
 	@Mock
@@ -49,7 +49,7 @@ public class TraceWebClientBeanPostProcessorTest {
 	Span span;
 
 	@Test
-	public void should_add_filter_only_once_to_web_client() {
+	void should_add_filter_only_once_to_web_client() {
 		TraceWebClientBeanPostProcessor processor = new TraceWebClientBeanPostProcessor(
 				this.springContext);
 		WebClient client = WebClient.create();
@@ -65,7 +65,7 @@ public class TraceWebClientBeanPostProcessorTest {
 	}
 
 	@Test
-	public void should_add_filter_only_once_to_web_client_via_builder() {
+	void should_add_filter_only_once_to_web_client_via_builder() {
 		TraceWebClientBeanPostProcessor processor = new TraceWebClientBeanPostProcessor(
 				this.springContext);
 		WebClient.Builder builder = WebClient.builder();
@@ -83,7 +83,7 @@ public class TraceWebClientBeanPostProcessorTest {
 	}
 
 	@Test
-	public void should_close_span_on_cancel() {
+	void should_close_span_on_cancel() {
 		TraceWebClientSubscription traceSubscription = new TraceWebClientSubscription(
 				subscription, new AtomicReference<>(span));
 
@@ -98,7 +98,7 @@ public class TraceWebClientBeanPostProcessorTest {
 	}
 
 	@Test
-	public void should_not_crash_on_cancel_when_span_clear() {
+	void should_not_crash_on_cancel_when_span_clear() {
 		TraceWebClientSubscription traceSubscription = new TraceWebClientSubscription(
 				subscription, new AtomicReference<>());