From d090ce325545f4109080680785c2d5e23571594a Mon Sep 17 00:00:00 2001
From: blake-bauman <blake_bauman@apple.com>
Date: Mon, 15 Mar 2021 11:52:23 -0700
Subject: [PATCH] Issue 1860 - WebFluxSleuthOperators should use ContextView
 (#1881)

* Adds variant of WebFluxSleuthOperators methods which use ContextView
* Have existing methods using Context invoke the new methods

Fixes gh-1860
---
 .../web/WebFluxSleuthOperators.java           | 23 ++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/WebFluxSleuthOperators.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/WebFluxSleuthOperators.java
index 023917eea..b18f58a86 100644
--- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/WebFluxSleuthOperators.java
+++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/WebFluxSleuthOperators.java
@@ -24,6 +24,7 @@ import org.apache.commons.logging.LogFactory;
 import reactor.core.publisher.Signal;
 import reactor.core.publisher.SignalType;
 import reactor.util.context.Context;
+import reactor.util.context.ContextView;
 
 import org.springframework.cloud.sleuth.CurrentTraceContext;
 import org.springframework.cloud.sleuth.Span;
@@ -94,6 +95,15 @@ public final class WebFluxSleuthOperators {
 	 * @param runnable - lambda to execute within the tracing context
 	 */
 	public static void withSpanInScope(Context context, Runnable runnable) {
+		withSpanInScope((ContextView) context, runnable);
+	}
+
+	/**
+	 * Wraps a runnable with a span.
+	 * @param context - Reactor context that contains the {@link TraceContext}
+	 * @param runnable - lambda to execute within the tracing context
+	 */
+	public static void withSpanInScope(ContextView context, Runnable runnable) {
 		CurrentTraceContext currentTraceContext = context.get(CurrentTraceContext.class);
 		TraceContext traceContext = traceContextOrNew(context);
 		try (CurrentTraceContext.Scope scope = currentTraceContext.maybeScope(traceContext)) {
@@ -109,12 +119,23 @@ public final class WebFluxSleuthOperators {
 	 * @return value from the callable
 	 */
 	public static <T> T withSpanInScope(Context context, Callable<T> callable) {
+		return withSpanInScope((ContextView) context, callable);
+	}
+
+	/**
+	 * Wraps a callable with a span.
+	 * @param context - Reactor context that contains the {@link TraceContext}
+	 * @param callable - lambda to execute within the tracing context
+	 * @param <T> callable's return type
+	 * @return value from the callable
+	 */
+	public static <T> T withSpanInScope(ContextView context, Callable<T> callable) {
 		CurrentTraceContext currentTraceContext = context.get(CurrentTraceContext.class);
 		TraceContext traceContext = traceContextOrNew(context);
 		return withContext(callable, currentTraceContext, traceContext);
 	}
 
-	private static TraceContext traceContextOrNew(Context context) {
+	private static TraceContext traceContextOrNew(ContextView context) {
 		Tracer tracer = context.get(Tracer.class);
 		if (!context.hasKey(TraceContext.class)) {
 			if (log.isDebugEnabled()) {