From 68497e0c516c57fe9b7474213a71f23873257526 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin.grzejszczak@gmail.com>
Date: Fri, 4 Dec 2020 11:46:14 +0100
Subject: [PATCH] Guarding context retreval; fixes gh-1789

---
 .../instrument/reactor/ReactorSleuth.java     | 18 ++++++++--
 .../ScopePassingSpanSubscriberTests.java      | 36 +++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java
index 968f9532a..345bfbb8a 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java
@@ -87,13 +87,13 @@ public abstract class ReactorSleuth {
 				if (log.isTraceEnabled()) {
 					String message = "Spring Context [" + springContext
 							+ "] is not yet refreshed. This is unexpected. Reactor Context is ["
-							+ sub.currentContext() + "] and name is [" + name(sub) + "]";
+							+ context(sub) + "] and name is [" + name(sub) + "]";
 					log.trace(message);
 				}
 				return sub;
 			}
 
-			Context context = sub.currentContext();
+			Context context = context(sub);
 
 			if (log.isTraceEnabled()) {
 				log.trace("Spring context [" + springContext + "], Reactor context ["
@@ -108,7 +108,7 @@ public abstract class ReactorSleuth {
 				if (log.isTraceEnabled() || assertOn) {
 					String message = "Spring Context [" + springContext
 							+ "] did not return a CurrentTraceContext. Reactor Context is ["
-							+ sub.currentContext() + "] and name is [" + name(sub) + "]";
+							+ context(sub) + "] and name is [" + name(sub) + "]";
 					log.trace(message);
 					assert false : message; // should never happen, but don't break.
 				}
@@ -129,6 +129,18 @@ public abstract class ReactorSleuth {
 		});
 	}
 
+	private static <T> Context context(CoreSubscriber<? super T> sub) {
+		try {
+			return sub.currentContext();
+		}
+		catch (Exception ex) {
+			if (log.isDebugEnabled()) {
+				log.debug("Exception occurred while trying to retrieve the context", ex);
+			}
+		}
+		return Context.empty();
+	}
+
 	static String name(CoreSubscriber<?> sub) {
 		return Scannable.from(sub).name();
 	}
diff --git a/tests/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java b/tests/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java
index 73b03438c..013a15105 100644
--- a/tests/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java
+++ b/tests/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriberTests.java
@@ -113,6 +113,33 @@ public class ScopePassingSpanSubscriberTests {
 		}
 	};
 
+	Subscriber<Object> exceptionThrowingPassingSpanSubscriber = new CoreSubscriber<Object>() {
+		@Override
+		public void onSubscribe(Subscription s) {
+
+		}
+
+		@Override
+		public void onNext(Object o) {
+
+		}
+
+		@Override
+		public void onError(Throwable t) {
+
+		}
+
+		@Override
+		public void onComplete() {
+
+		}
+
+		@Override
+		public Context currentContext() {
+			throw new NullPointerException("Boom!");
+		}
+	};
+
 	AnnotationConfigApplicationContext springContext = new AnnotationConfigApplicationContext();
 
 	@Before
@@ -152,6 +179,15 @@ public class ScopePassingSpanSubscriberTests {
 
 	@Test
 	public void should_set_empty_context_when_context_is_null() {
+		ScopePassingSpanSubscriber<?> subscriber = new ScopePassingSpanSubscriber<>(
+				this.exceptionThrowingPassingSpanSubscriber, Context.empty(),
+				this.currentTraceContext, null);
+
+		then(subscriber.currentContext().isEmpty()).isTrue();
+	}
+
+	@Test
+	public void should_set_empty_context_when_exception_occurs_while_trying_to_retrieve_the_context() {
 		ScopePassingSpanSubscriber<?> subscriber = new ScopePassingSpanSubscriber<>(null,
 				Context.empty(), this.currentTraceContext, null);