From fe1f3e856aba981111fb4c5d53361f28de612c78 Mon Sep 17 00:00:00 2001
From: Julien Hoarau <madgnome@gmail.com>
Date: Wed, 5 Dec 2018 08:23:44 +1100
Subject: [PATCH] Fixes #1138 Report trace in case of cancellation when calling
 webclient. (#1140)

Report trace in case of cancellation when calling webclient.

Fixes gh-1138
---
 .../TraceWebClientBeanPostProcessor.java      | 25 ++++++++++++++++++-
 .../client/integration/WebClientTests.java    | 25 ++++++++++++++++++-
 2 files changed, 48 insertions(+), 2 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 c30d91e7c..ad9af9fb8 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
@@ -102,6 +102,8 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 	private static final String CLIENT_SPAN_KEY = "sleuth.webclient.clientSpan";
 
+	private static final String CANCELLED_SUBSCRIPTION_ERROR = "CANCELLED";
+
 	static final Propagation.Setter<ClientRequest.Builder, String> SETTER = new Propagation.Setter<ClientRequest.Builder, String>() {
 		@Override
 		public void put(ClientRequest.Builder carrier, String key, String value) {
@@ -238,7 +240,18 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 			@Override
 			public void onSubscribe(Subscription subscription) {
-				this.actual.onSubscribe(subscription);
+				this.actual.onSubscribe(new Subscription() {
+					@Override
+					public void request(long n) {
+						subscription.request(n);
+					}
+
+					@Override
+					public void cancel() {
+						terminateSpanOnCancel();
+						subscription.cancel();
+					}
+				});
 			}
 
 			@Override
@@ -289,6 +302,16 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 				ws.close();
 			}
 
+			void terminateSpanOnCancel() {
+				if (log.isDebugEnabled()) {
+					log.debug("Subscription was cancelled. Will close the span [" + span
+							+ "]");
+				}
+
+				span.tag("error", CANCELLED_SUBSCRIPTION_ERROR);
+				handleReceive(span, ws, null, null);
+			}
+
 			void terminateSpan(@Nullable ClientResponse clientResponse,
 					@Nullable Throwable throwable) {
 				if (clientResponse == null || clientResponse.statusCode() == null) {
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java
index 4720cdd2d..69715ea0d 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java
@@ -17,6 +17,7 @@
 package org.springframework.cloud.sleuth.instrument.web.client.integration;
 
 import java.lang.invoke.MethodHandles;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -51,7 +52,6 @@ import org.awaitility.Awaitility;
 import org.junit.After;
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -378,6 +378,29 @@ public class WebClientTests {
 				.contains("CLIENT");
 	}
 
+	@Test
+	public void shouldReportTraceForCancelledRequestViaWebClient() {
+		Span span = this.tracer.nextSpan().name("foo").start();
+
+		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) {
+			this.webClient.get().uri("http://localhost:" + this.port + "/noresponse")
+					.retrieve().bodyToMono(String.class)
+					.timeout(Duration.ofMillis(0))
+					.block();
+		}
+		catch (Exception e) {
+
+		}
+		finally {
+			span.finish();
+		}
+
+		final Optional<zipkin2.Span> clientSpan = this.reporter.getSpans().stream()
+				.filter(s -> s.kind() == zipkin2.Span.Kind.CLIENT).findFirst();
+		then(clientSpan).isPresent();
+		then(clientSpan.get().tags()).containsEntry("error", "CANCELLED");
+	}
+
 	Object[] parametersForShouldAttachTraceIdWhenCallingAnotherService() {
 		return new Object[] {
 				(ResponseEntityProvider) (tests) -> tests.testFeignInterface.headers(),