Added support for netty-http-brave; fixes gh-1690

spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/brave/instrument/web/BraveHttpConfiguration.java

@@ -21,12 +21,15 @@ import java.util.regex.Pattern;
 
 import javax.validation.constraints.NotNull;
 
+import brave.Tracer;
 import brave.Tracing;
 import brave.http.HttpRequest;
 import brave.http.HttpTracing;
 import brave.http.HttpTracingCustomizer;
+import brave.propagation.CurrentTraceContext;
 import brave.sampler.SamplerFunction;
 import brave.sampler.SamplerFunctions;
+import reactor.util.context.Context;
 
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -38,6 +41,7 @@ import org.springframework.cloud.sleuth.autoconfig.instrument.web.SleuthWebPrope
 import org.springframework.cloud.sleuth.brave.bridge.BraveHttpRequestParser;
 import org.springframework.cloud.sleuth.brave.bridge.BraveHttpResponseParser;
 import org.springframework.cloud.sleuth.brave.bridge.BraveSamplerFunction;
+import org.springframework.cloud.sleuth.brave.instrument.web.BraveSpanFromContextRetriever;
 import org.springframework.cloud.sleuth.brave.instrument.web.CompositeHttpSampler;
 import org.springframework.cloud.sleuth.brave.instrument.web.SkipPatternHttpClientSampler;
 import org.springframework.cloud.sleuth.brave.instrument.web.SkipPatternHttpServerSampler;
@@ -50,6 +54,7 @@ import org.springframework.cloud.sleuth.instrument.web.HttpServerRequestParser;
 import org.springframework.cloud.sleuth.instrument.web.HttpServerResponseParser;
 import org.springframework.cloud.sleuth.instrument.web.HttpServerSampler;
 import org.springframework.cloud.sleuth.instrument.web.SkipPatternProvider;
+import org.springframework.cloud.sleuth.instrument.web.SpanFromContextRetriever;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Import;
@@ -205,4 +210,15 @@ public class BraveHttpConfiguration {
 		return new SkipPatternHttpClientSampler(Pattern.compile(skipPattern));
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnClass(Context.class)
+	static class BraveWebFilterConfiguration {
+
+		@Bean
+		SpanFromContextRetriever braveSpanFromContextRetriever(CurrentTraceContext currentTraceContext, Tracer tracer) {
+			return new BraveSpanFromContextRetriever(currentTraceContext, tracer);
+		}
+
+	}
+
 }

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveSpan.java

@@ -27,11 +27,11 @@ import org.springframework.cloud.sleuth.TraceContext;
  * @author Marcin Grzejszczak
  * @since 3.0.0
  */
-class BraveSpan implements Span {
+public class BraveSpan implements Span {
 
 	final brave.Span delegate;
 
-	BraveSpan(brave.Span delegate) {
+	public BraveSpan(brave.Span delegate) {
 		this.delegate = delegate;
 	}
 
@@ -91,11 +91,11 @@ class BraveSpan implements Span {
 		return this.delegate != null ? this.delegate.toString() : "null";
 	}
 
-	static brave.Span toBrave(Span span) {
+	public static brave.Span toBrave(Span span) {
 		return ((BraveSpan) span).delegate;
 	}
 
-	static Span fromBrave(brave.Span span) {
+	public static Span fromBrave(brave.Span span) {
 		return new BraveSpan(span);
 	}
 

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/instrument/web/BraveSpanFromContextRetriever.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2013-2021 the original author or 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
+ *
+ *      https://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 org.springframework.cloud.sleuth.brave.instrument.web;
+
+import brave.Tracer;
+import brave.propagation.CurrentTraceContext;
+import brave.propagation.TraceContext;
+import reactor.util.context.Context;
+
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.brave.bridge.BraveSpan;
+import org.springframework.cloud.sleuth.instrument.web.SpanFromContextRetriever;
+
+/**
+ * Retrieves Brave specific classes from Reactor context.
+ *
+ * @author Marcin Grzejszczak
+ * @since 3.0.2
+ */
+public class BraveSpanFromContextRetriever implements SpanFromContextRetriever {
+
+	private final CurrentTraceContext currentTraceContext;
+
+	private final Tracer tracer;
+
+	public BraveSpanFromContextRetriever(CurrentTraceContext currentTraceContext, Tracer tracer) {
+		this.currentTraceContext = currentTraceContext;
+		this.tracer = tracer;
+	}
+
+	@Override
+	public Span findSpan(Context context) {
+		Object braveSpan = context.getOrDefault(brave.Span.class, null);
+		if (braveSpan != null) {
+			return BraveSpan.fromBrave((brave.Span) braveSpan);
+		}
+		Object braveContext = context.getOrDefault(TraceContext.class, null);
+		if (braveContext != null) {
+			TraceContext traceContext = (TraceContext) braveContext;
+			try (CurrentTraceContext.Scope scope = this.currentTraceContext.maybeScope(traceContext)) {
+				return BraveSpan.fromBrave(this.tracer.currentSpan());
+			}
+		}
+		return null;
+	}
+
+}

spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/BraveSpanFromContextRetrieverTests.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2013-2021 the original author or 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
+ *
+ *      https://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 org.springframework.cloud.sleuth.brave.instrument.web;
+
+import brave.Span;
+import brave.Tracing;
+import brave.propagation.StrictCurrentTraceContext;
+import brave.propagation.TraceContext;
+import brave.sampler.Sampler;
+import brave.test.TestSpanHandler;
+import org.junit.jupiter.api.Test;
+import reactor.util.context.Context;
+
+import org.springframework.cloud.sleuth.brave.bridge.BraveSpan;
+
+import static org.assertj.core.api.BDDAssertions.then;
+
+class BraveSpanFromContextRetrieverTests {
+
+	TestSpanHandler spans = new TestSpanHandler();
+
+	StrictCurrentTraceContext traceContext = StrictCurrentTraceContext.create();
+
+	Tracing tracing = Tracing.newBuilder().currentTraceContext(this.traceContext)
+			.sampler(Sampler.ALWAYS_SAMPLE).addSpanHandler(this.spans).build();
+
+	brave.Tracer tracer = this.tracing.tracer();
+
+	BraveSpanFromContextRetriever retriever = new BraveSpanFromContextRetriever(this.traceContext, this.tracer);
+
+	@Test
+	void should_return_null_when_no_brave_specific_entries_are_present_in_context() {
+		then(retriever.findSpan(Context.empty())).isNull();
+	}
+
+	@Test
+	void should_return_span_when_brave_span_present_in_context() {
+		Span span = this.tracer.nextSpan();
+
+		then(BraveSpan.toBrave(retriever.findSpan(Context.of(Span.class, span)))).isSameAs(span);
+	}
+
+	@Test
+	void should_return_span_when_brave_trace_context_present_in_context() {
+		Span span = this.tracer.nextSpan();
+
+		then(BraveSpan.toBrave(retriever.findSpan(Context.of(TraceContext.class, span.context())))).isEqualTo(span);
+	}
+}

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/SpanFromContextRetriever.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2013-2021 the original author or 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
+ *
+ *      https://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 org.springframework.cloud.sleuth.instrument.web;
+
+import java.util.regex.Pattern;
+
+import reactor.util.context.Context;
+
+import org.springframework.cloud.sleuth.Span;
+
+/**
+ * Provides a URL {@link Pattern} for spans that should be not sampled.
+ *
+ * @author Marcin Grzejszczak
+ * @since 3.0.2
+ */
+public interface SpanFromContextRetriever {
+
+	/**
+	 * @param context - Reactor context
+	 * @return span or {@code null} if no span present
+	 */
+	default Span findSpan(Context context) {
+		return null;
+	};
+
+}

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java

@@ -81,6 +81,8 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa
 
 	private int order;
 
+	private SpanFromContextRetriever spanFromContextRetriever;
+
 	@Deprecated
 	public TraceWebFilter(Tracer tracer, HttpServerHandler handler) {
 		this.tracer = tracer;
@@ -102,7 +104,7 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa
 		if (log.isDebugEnabled()) {
 			log.debug("Received a request to uri [" + uri + "]");
 		}
-		return new MonoWebFilterTrace(source, exchange, tracePresent, this);
+		return new MonoWebFilterTrace(source, exchange, tracePresent, this, spanFromContextRetriever());
 	}
 
 	private boolean isTracePresent() {
@@ -135,6 +137,15 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa
 		return this.currentTraceContext;
 	}
 
+	private SpanFromContextRetriever spanFromContextRetriever() {
+		if (this.spanFromContextRetriever == null) {
+			this.spanFromContextRetriever = this.applicationContext.getBeanProvider(SpanFromContextRetriever.class)
+					.getIfAvailable(() -> new SpanFromContextRetriever() {
+					});
+		}
+		return this.spanFromContextRetriever;
+	}
+
 	private static class MonoWebFilterTrace extends MonoOperator<Void, Void> implements TraceContextPropagator {
 
 		final ServerWebExchange exchange;
@@ -151,8 +162,10 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa
 
 		final CurrentTraceContext currentTraceContext;
 
+		final SpanFromContextRetriever spanFromContextRetriever;
+
 		MonoWebFilterTrace(Mono<? extends Void> source, ServerWebExchange exchange, boolean initialTracePresent,
-				TraceWebFilter parent) {
+				TraceWebFilter parent, SpanFromContextRetriever spanFromContextRetriever) {
 			super(source);
 			this.tracer = parent.tracer;
 			this.handler = parent.handler;
@@ -160,6 +173,7 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa
 			this.exchange = exchange;
 			this.span = exchange.getAttribute(TRACE_REQUEST_ATTR);
 			this.initialTracePresent = initialTracePresent;
+			this.spanFromContextRetriever = spanFromContextRetriever;
 		}
 
 		@Override
@@ -207,12 +221,16 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa
 						log.debug("Found span in attribute " + span);
 					}
 				}
-				else {
+				span = this.spanFromContextRetriever.findSpan(c);
+				if (this.span == null && span == null) {
 					span = this.handler.handleReceive(new WrappedRequest(this.exchange.getRequest()));
 					if (log.isDebugEnabled()) {
 						log.debug("Handled receive of span " + span);
 					}
 				}
+				else if (log.isDebugEnabled()) {
+					log.debug("Found tracer specific span in reactor context [" + span + "]");
+				}
 				this.exchange.getAttributes().put(TRACE_REQUEST_ATTR, span);
 			}
 			return span;