Regression: sleuth.propagation-keys aren't available in filter with WebFlux

[anatoliy-balakirev]

Bug report:

Propagation keys are not available in the filter with WebFlux.

Spring Boot version: 2.1.2.RELEASE
Spring Cloud version: Greenwich.RELEASE. No issue with Finchley.RELEASE
Sample project: https://github.com/anatoliy-balakirev/sleuth-test

So basically if I set propagation key:

spring.sleuth.propagation-keys=x-b3-user-id

Within my first controller:

@Slf4j
@RestController
@RequiredArgsConstructor
public class DemoController {

    private final WebClient webClient;

    @GetMapping("/demo")
    public Mono<String> getMessageFromContext() {
        ExtraFieldPropagation.set(USER_ID_KEY, "123");
        log.info("Sending request to another controller. Set userId as an extra field: {}", USER_ID_KEY);

        return webClient.get().uri("/userId").retrieve().bodyToMono(String.class).doOnNext(userId -> {
            log.info("User id from another controller is: {}", userId);
        });
    }

}

And then try to access it within filter, which I register in the WebClient:

    @Bean
    public WebClient myWebClient(final WebClient.Builder webClientBuilder) {
        return webClientBuilder.baseUrl("http://localhost:8080")
            .filter((request, exchangeFunction) -> {
                final String userId = request.headers().getFirst(USER_ID_KEY);
                log.info("User id in filter is: {}", userId);
                return exchangeFunction.exchange(request);
            })
            .build();
    }

I'm getting null there. Still that key is propagated to my next controller, I can access it there and send back:

@Slf4j
@RestController
@RequiredArgsConstructor
public class UserIdController {


    @GetMapping("/userId")
    public Mono<String> getMessage() {
        return Mono.fromCallable(() -> {
            log.info("Got into userId controller");
            return ExtraFieldPropagation.get(USER_ID_KEY);
        });
    }

}

Here is the log with Greenwich.RELEASE:

2019-01-28 13:37:02.098  INFO [-,a2058a4c791cf166,a2058a4c791cf166,false] 87989 --- [ctor-http-nio-3] com.sleuth.test.demo.DemoController      : Sending request to another controller. Set userId as an extra field: x-b3-user-id
2019-01-28 13:37:02.140  INFO [-,a2058a4c791cf166,a2058a4c791cf166,false] 87989 --- [ctor-http-nio-3] com.sleuth.test.demo.DemoApplication     : User id in filter is: null
2019-01-28 13:37:02.229  INFO [-,a2058a4c791cf166,6f2a3567a3746bcb,false] 87989 --- [ctor-http-nio-4] com.sleuth.test.demo.UserIdController    : Got into userId controller
2019-01-28 13:37:02.277  INFO [-,a2058a4c791cf166,a2058a4c791cf166,false] 87989 --- [ctor-http-nio-3] com.sleuth.test.demo.DemoController      : User id from another controller is: 123

And here is with Finchley.RELEASE (you can uncomment it in the sample project's build.gradle to test):

2019-01-28 13:38:42.714  INFO [-,ea901756ef83f1bd,ea901756ef83f1bd,false] 88004 --- [ctor-http-nio-3] com.sleuth.test.demo.DemoController      : Sending request to another controller. Set userId as an extra field: x-b3-user-id
2019-01-28 13:38:42.758  INFO [-,ea901756ef83f1bd,ea901756ef83f1bd,false] 88004 --- [ctor-http-nio-3] com.sleuth.test.demo.DemoApplication     : User id in filter is: 123
2019-01-28 13:38:42.854  INFO [-,ea901756ef83f1bd,8c544927d480688e,false] 88004 --- [ctor-http-nio-4] com.sleuth.test.demo.UserIdController    : Got into userId controller
2019-01-28 13:38:42.897  INFO [-,ea901756ef83f1bd,ea901756ef83f1bd,false] 88004 --- [ctor-http-nio-3] com.sleuth.test.demo.DemoController      : User id from another controller is: 123

Regression seems to be introduced here: #1126

[bsideup]

Hi @anatoliy-balakirev

I tried to play with your example and turned on TRACE logging on both Sleuth's instrumentation and Brave.

I see this:

2019-02-05 15:47:25.902  INFO [,9ff3eec732ea58ff,,9ff3eec732ea58ff,] 58144 --- [reactor-http-nio-2]  c.s.t.d.DemoApplication:32 : User id in filter is: null
2019-02-05 15:47:25.917 DEBUG [,9ff3eec732ea58ff,,9ff3eec732ea58ff,] 58144 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:191 : Instrumenting WebClient call
2019-02-05 15:47:25.917 TRACE [,9ff3eec732ea58ff,,9ff3eec732ea58ff,] 58144 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [X-B3-TraceId] with value [9ff3eec732ea58ff]
2019-02-05 15:47:25.918 TRACE [,9ff3eec732ea58ff,,9ff3eec732ea58ff,] 58144 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [X-B3-SpanId] with value [466907629e99dfc0]
2019-02-05 15:47:25.918 TRACE [,9ff3eec732ea58ff,,9ff3eec732ea58ff,] 58144 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [X-B3-ParentSpanId] with value [9ff3eec732ea58ff]
2019-02-05 15:47:25.918 TRACE [,9ff3eec732ea58ff,,9ff3eec732ea58ff,] 58144 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [X-B3-Sampled] with value [0]
2019-02-05 15:47:25.919 TRACE [,9ff3eec732ea58ff,,9ff3eec732ea58ff,] 58144 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [x-b3-user-id] with value [123]

it seems that the header is set after the filter function. However, if I change it to:

            .filter((request, exchangeFunction) -> {
                return Mono.defer(() -> {
                    final String userId = request.headers().getFirst(USER_ID_KEY);
                    log.info("User id in filter is: {}", userId);
                    return exchangeFunction.exchange(request);
                });
            })

I get the desired result:

2019-02-05 15:52:09.374 TRACE [,a1a40963cbe673c8,,a1a40963cbe673c8,] 58436 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [x-b3-user-id] with value [123]
2019-02-05 15:52:09.375 DEBUG [,a1a40963cbe673c8,,a1a40963cbe673c8,] 58436 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:196 : Handled send of NoopSpan(a1a40963cbe673c8/1a3157b1623f5a41)
2019-02-05 15:52:09.376  INFO [,a1a40963cbe673c8,a1a40963cbe673c8,1a3157b1623f5a41,123] 58436 --- [reactor-http-nio-2]  c.s.t.d.DemoApplication:34 : User id in filter is: 123

/cc @marcingrzejszczak

[bsideup]

So, here is the thing. This code in Spring Framework applies all filters (Including Sleuth):
https://github.com/spring-projects/spring-framework/blob/0b9522c84e28f5784bb8d9b5eaf8d9c5cf46bc37/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClientBuilder.java#L207-L211

After #1126, we set the headers in an appropriate place, with the respect to cold publishers (to not pollute the scope where filters are applied)

Here is the change:
https://github.com/spring-cloud/spring-cloud-sleuth/pull/1126/files#diff-5c3dd86d2005f67c41d439eb2ac01b8b

Now, the question is whether this change is breaking one, or a bugfix, in fact :)

@anatoliy-balakirev would you consider using Mono.defer here?

[bsideup]

Okay, after more research I figured that there was another problem (fixed in #1206)

@anatoliy-balakirev I tried your example with a SNAPSHOT from that branch and now the order is correct:

2019-02-05 16:56:58.905  INFO [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  c.s.t.d.DemoController:23 : Sending request to another controller. Set userId as an extra field: x-b3-user-id
2019-02-05 16:56:58.928 DEBUG [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:151 : Instrumenting WebClient call
2019-02-05 16:56:58.930 TRACE [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [X-B3-TraceId] with value [60d828ee3f1d44f5]
2019-02-05 16:56:58.931 TRACE [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [X-B3-SpanId] with value [87d6bd4080a2a086]
2019-02-05 16:56:58.932 TRACE [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [X-B3-ParentSpanId] with value [60d828ee3f1d44f5]
2019-02-05 16:56:58.933 TRACE [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [X-B3-Sampled] with value [0]
2019-02-05 16:56:58.934 TRACE [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:112 : Replacing [x-b3-user-id] with value [123]
2019-02-05 16:56:58.935 DEBUG [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  o.s.c.s.i.w.c.TraceExchangeFilterFunction:156 : Handled send of NoopSpan(60d828ee3f1d44f5/87d6bd4080a2a086)
2019-02-05 16:56:58.945  INFO [,60d828ee3f1d44f5,,60d828ee3f1d44f5,] 62884 --- [reactor-http-nio-2]  c.s.t.d.DemoApplication:32 : User id in filter is: 123

[anatoliy-balakirev]

Sorry, just got here :)
@bsideup thanks for analysing / fixing it!