From 23b8caa27c0f820b95bddc6b23f7869d5d551220 Mon Sep 17 00:00:00 2001
From: Johan Vandeweerd <johan.v@sparkcentral.com>
Date: Wed, 15 May 2019 10:09:52 +0200
Subject: [PATCH] Add decorator hooks during lifecycle of Publisher

---
 .../springwebflux/client/TracingClientResponseMono.java        | 3 +++
 .../springwebflux/client/TracingClientResponseSubscriber.java  | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseMono.java b/dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseMono.java
index b02c2dd73c..8331659176 100644
--- a/dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseMono.java
+++ b/dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseMono.java
@@ -1,5 +1,7 @@
 package datadog.trace.instrumentation.springwebflux.client;
 
+import static datadog.trace.instrumentation.springwebflux.client.SpringWebfluxHttpClientDecorator.DECORATE;
+
 import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.Tracer;
@@ -38,6 +40,7 @@ public class TracingClientResponseMono extends Mono<ClientResponse> {
             .asChildOf(parentSpan)
             .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
             .start();
+    DECORATE.afterStart(span);
 
     try (final Scope scope = tracer.scopeManager().activate(span, false)) {
       final ClientRequest mutatedRequest =
diff --git a/dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseSubscriber.java b/dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseSubscriber.java
index bbccd9a83c..799e98c6fa 100644
--- a/dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseSubscriber.java
+++ b/dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseSubscriber.java
@@ -44,6 +44,7 @@ public class TracingClientResponseSubscriber implements CoreSubscriber<ClientRes
           @Override
           public void cancel() {
             DECORATE.onCancel(span);
+            DECORATE.beforeFinish(span);
             subscription.cancel();
             span.finish();
           }
@@ -73,6 +74,7 @@ public class TracingClientResponseSubscriber implements CoreSubscriber<ClientRes
       subscriber.onError(throwable);
     } finally {
       DECORATE.onError(span, throwable);
+      DECORATE.beforeFinish(span);
       span.finish();
     }
   }
@@ -82,6 +84,7 @@ public class TracingClientResponseSubscriber implements CoreSubscriber<ClientRes
     try {
       subscriber.onComplete();
     } finally {
+      DECORATE.beforeFinish(span);
       span.finish();
     }
   }