From 5fd8c744963a996e0b9ce6c533c3502b811d315e Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <mgrzejszczak@pivotal.io>
Date: Thu, 21 Dec 2017 11:13:14 +0100
Subject: [PATCH] Fixed unnecessary last operator wrapping for webflux

fixes #811
---
 .../TraceReactorAutoConfiguration.java        | 29 +++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java
index db88ade78..b437928cd 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java
@@ -11,11 +11,13 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnNotWebApplication;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
 import org.springframework.cloud.sleuth.SpanNamer;
 import org.springframework.cloud.sleuth.TraceKeys;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.instrument.async.TraceableScheduledExecutorService;
 import org.springframework.cloud.sleuth.instrument.web.TraceWebFluxAutoConfiguration;
+import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Mono;
@@ -37,15 +39,34 @@ public class TraceReactorAutoConfiguration {
 
 	@Configuration
 	@ConditionalOnBean(Tracer.class)
-	@ConditionalOnNotWebApplication
 	static class TraceReactorConfiguration {
 		@Autowired Tracer tracer;
 		@Autowired TraceKeys traceKeys;
 		@Autowired SpanNamer spanNamer;
+		@Autowired LastOperatorWrapper lastOperatorWrapper;
+
+		@Bean
+		@ConditionalOnNotWebApplication
+		LastOperatorWrapper spanOperator() {
+			return new LastOperatorWrapper() {
+				@Override public void wrapLastOperator(Tracer tracer) {
+					Hooks.onLastOperator(ReactorSleuth.spanOperator(tracer));
+				}
+			};
+		}
+
+		@Bean
+		@ConditionalOnWebApplication
+		LastOperatorWrapper noOpLastOperatorWrapper() {
+			return new LastOperatorWrapper() {
+				@Override public void wrapLastOperator(Tracer tracer) {
+				}
+			};
+		}
 
 		@PostConstruct
 		public void setupHooks() {
-			Hooks.onLastOperator(ReactorSleuth.spanOperator(this.tracer));
+			this.lastOperatorWrapper.wrapLastOperator(this.tracer);
 			Schedulers.setFactory(new Schedulers.Factory() {
 				@Override public ScheduledExecutorService decorateExecutorService(String schedulerType,
 						Supplier<? extends ScheduledExecutorService> actual) {
@@ -63,4 +84,8 @@ public class TraceReactorAutoConfiguration {
 			Schedulers.resetFactory();
 		}
 	}
+}
+
+interface LastOperatorWrapper {
+	void wrapLastOperator(Tracer tracer);
 }
\ No newline at end of file