com.nike.wingtips.Tracer Java Examples

protected void unlinkTracingAndMdcFromCurrentThread(ChannelHandlerContext ctx,
                                                    Pair<Deque<Span>, Map<String, String>> origThreadInfo) {
    // Update the state (if we have any) with the current values of the MDC and tracer data
    HttpProcessingState state = ChannelAttributes.getHttpProcessingStateForChannel(ctx).get();
    if (state != null) {
        // Get references to the *current* MDC and tracer data for storing in our ctx
        //      and set them on the state object
        Map<String, String> currentMdcContextMapForState = MDC.getCopyOfContextMap();
        Deque<Span> currentTraceStackForState = Tracer.getInstance().unregisterFromThread();

        state.setLoggerMdcContextMap(currentMdcContextMapForState);
        state.setDistributedTraceStack(currentTraceStackForState);
    }

    // Reset the thread to the way it was before linkTracingAndMdcToCurrentThread was called
    AsyncNettyHelper.unlinkTracingAndMdcFromCurrentThread(origThreadInfo);
}
 

private Pair<ObjectHolder<Span>, ObjectHolder<Span>> setupBeforeAndAfterSpanCaptureForOnThrowable(
    CompletableFuture<String> cfMock) throws Throwable {
    ObjectHolder<Span> before = new ObjectHolder<>();
    ObjectHolder<Span> after = new ObjectHolder<>();

    doAnswer(invocation -> {
        before.setObj(Tracer.getInstance().getCurrentSpan());
        return invocation.callRealMethod();
    }).when(circuitBreakerManualTaskMock).handleException(any(Throwable.class));

    doAnswer(invocation -> {
        after.setObj(Tracer.getInstance().getCurrentSpan());
        return invocation.callRealMethod();
    }).when(cfMock).completeExceptionally(any(Throwable.class));

    return Pair.of(before, after);
}
 

@Before
public void beforeMethod() {
    functionMock = mock(Function.class);

    inObj = new Object();
    outObj = new Object();
    throwExceptionDuringCall = false;
    currentSpanStackWhenFunctionWasCalled = new ArrayList<>();
    currentMdcInfoWhenFunctionWasCalled = new ArrayList<>();
    doAnswer(invocation -> {
        currentSpanStackWhenFunctionWasCalled.add(Tracer.getInstance().getCurrentSpanStackCopy());
        currentMdcInfoWhenFunctionWasCalled.add(MDC.getCopyOfContextMap());
        if (throwExceptionDuringCall)
            throw new RuntimeException("kaboom");
        return outObj;
    }).when(functionMock).apply(inObj);

    resetTracing();
}
 

@Test
public void linkTracingAndMdcToCurrentThread_should_set_tracing_and_mdc_to_state_values_if_available() {
    // given
    Map<String, String> stateMdcInfo = new HashMap<>();
    stateMdcInfo.put("foo", "bar");
    Deque<Span> stateTraceStack = new LinkedList<>();
    Span span = Span.generateRootSpanForNewTrace("fooSpanName", LOCAL_ONLY).withTraceId("fooTraceId").build();
    stateTraceStack.add(span);
    state.setLoggerMdcContextMap(stateMdcInfo);
    state.setDistributedTraceStack(stateTraceStack);

    assertThat(MDC.getCopyOfContextMap().isEmpty(), is(true));
    assertThat(Tracer.getInstance().getCurrentSpan(), nullValue());

    // when
    handler.linkTracingAndMdcToCurrentThread(ctxMock);

    // then
    // Tracer adds some stuff to the MDC
    stateMdcInfo.put(SpanFieldForLoggerMdc.TRACE_ID.mdcKey, span.getTraceId());
    assertThat(MDC.getCopyOfContextMap(), is(stateMdcInfo));
    assertThat(Tracer.getInstance().getCurrentSpanStackCopy(), is(stateTraceStack));
}
 

@Test
public void finalizeAndCompleteOverallRequestSpanAttachedToCurrentThread_does_nothing_if_called_when_current_thread_span_is_already_completed() {
    // given
    Span spanMock = mock(Span.class);
    Tracer.getInstance().registerWithThread(new ArrayDeque<>(singleton(spanMock)));
    reset(spanMock);
    doReturn(true).when(spanMock).isCompleted();

    assertThat(Tracer.getInstance().getCurrentSpan()).isSameAs(spanMock);

    // when
    WingtipsSpringWebfluxWebFilter.finalizeAndCompleteOverallRequestSpanAttachedToCurrentThread(
        exchange, mock(Throwable.class), tagAndNamingStrategy, tagAndNamingAdapterMock,
        Pair.of("foo", "bar")
    );

    // then
    assertThat(strategyResponseTaggingAndFinalSpanNameMethodCalled.get()).isFalse();
    assertThat(spanRecorder.completedSpans).isEmpty();
    verify(spanMock).isCompleted();
    verifyNoMoreInteractions(spanMock);
}
 

@Before
public void beforeMethod() {
    biFunctionMock = mock(BiFunction.class);

    inObj1 = new Object();
    inObj2 = new Object();
    outObj = new Object();
    throwExceptionDuringCall = false;
    currentSpanStackWhenFunctionWasCalled = new ArrayList<>();
    currentMdcInfoWhenFunctionWasCalled = new ArrayList<>();
    doAnswer(invocation -> {
        currentSpanStackWhenFunctionWasCalled.add(Tracer.getInstance().getCurrentSpanStackCopy());
        currentMdcInfoWhenFunctionWasCalled.add(MDC.getCopyOfContextMap());
        if (throwExceptionDuringCall)
            throw new RuntimeException("kaboom");
        return outObj;
    }).when(biFunctionMock).apply(inObj1, inObj2);

    resetTracing();
}
 

@Before
public void beforeMethod() {
    channelMock = mock(Channel.class);
    ctxMock = mock(ChannelHandlerContext.class);
    stateAttributeMock = mock(Attribute.class);
    state = new HttpProcessingState();
    doReturn(channelMock).when(ctxMock).channel();
    doReturn(stateAttributeMock).when(channelMock).attr(ChannelAttributes.HTTP_PROCESSING_STATE_ATTRIBUTE_KEY);
    doReturn(state).when(stateAttributeMock).get();

    runnableMock = mock(Runnable.class);

    throwExceptionDuringCall = false;
    currentSpanStackWhenRunnableWasCalled = new ArrayList<>();
    currentMdcInfoWhenRunnableWasCalled = new ArrayList<>();
    doAnswer(invocation -> {
        currentSpanStackWhenRunnableWasCalled.add(Tracer.getInstance().getCurrentSpanStackCopy());
        currentMdcInfoWhenRunnableWasCalled.add(MDC.getCopyOfContextMap());
        if (throwExceptionDuringCall)
            throw new RuntimeException("kaboom");
        return null;
    }).when(runnableMock).run();

    resetTracingAndMdc();
}
 

@Before
public void beforeMethod() {
    channelMock = mock(Channel.class);
    ctxMock = mock(ChannelHandlerContext.class);
    stateAttributeMock = mock(Attribute.class);
    state = new HttpProcessingState();
    doReturn(channelMock).when(ctxMock).channel();
    doReturn(stateAttributeMock).when(channelMock).attr(ChannelAttributes.HTTP_PROCESSING_STATE_ATTRIBUTE_KEY);
    doReturn(state).when(stateAttributeMock).get();

    supplierMock = mock(Supplier.class);

    outObj = new Object();
    throwExceptionDuringCall = false;
    currentSpanStackWhenSupplierWasCalled = new ArrayList<>();
    currentMdcInfoWhenSupplierWasCalled = new ArrayList<>();
    doAnswer(invocation -> {
        currentSpanStackWhenSupplierWasCalled.add(Tracer.getInstance().getCurrentSpanStackCopy());
        currentMdcInfoWhenSupplierWasCalled.add(MDC.getCopyOfContextMap());
        if (throwExceptionDuringCall)
            throw new RuntimeException("kaboom");
        return outObj;
    }).when(supplierMock).get();

    resetTracingAndMdc();
}
 

@Test
public void completeSubspan_works_as_expected() {
    // given
    Tracer.getInstance().startRequestWithRootSpan("fooRootSpan");
    Span currentSpanForCompletion = Tracer.getInstance().startSubSpan("fooSubspan", SpanPurpose.LOCAL_ONLY);
    TracingState expectedTracingStateForCompletion = TracingState.getCurrentThreadTracingState();

    // Now that the tracing-state-for-completion has been setup, setup the current thread tracing state to
    //      something completely different.
    Tracer.getInstance().unregisterFromThread();
    Tracer.getInstance().startRequestWithRootSpan("someCompletelyDifferentRootSpan");
    TracingState baseTracingState = TracingState.getCurrentThreadTracingState();

    Throwable error = mock(Throwable.class);

    // when
    filterSpy.completeSubspan(expectedTracingStateForCompletion, request, responseMock, error);

    // then
    assertThat(strategyResponseTaggingAndFinalSpanNameMethodCalled.get()).isTrue();
    strategyResponseTaggingArgs.get().verifyArgs(
        currentSpanForCompletion, request, responseMock, error, tagAndNamingAdapterMock
    );

    assertThat(currentSpanForCompletion.isCompleted()).isTrue();
    assertThat(spanRecorder.completedSpans).isEqualTo(singletonList(currentSpanForCompletion));

    // The base tracing state should have been restored right before the method call ended.
    assertThat(TracingState.getCurrentThreadTracingState()).isEqualTo(baseTracingState);
}
 

@DataProvider(value = {
    "true",
    "false"
})
@Test
public void onSuccess_handles_tracing_and_mdc_info_as_expected(boolean throwException) {
    // given
    throwExceptionDuringCall = throwException;
    Tracer.getInstance().startRequestWithRootSpan("foo");
    Deque<Span> spanStack = Tracer.getInstance().getCurrentSpanStackCopy();
    Map<String, String> mdcInfo = MDC.getCopyOfContextMap();
    ListenableFutureCallbackWithTracing instance = new ListenableFutureCallbackWithTracing(
        listenableFutureCallbackMock, spanStack, mdcInfo
    );
    resetTracing();
    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isNull();
    assertThat(MDC.getCopyOfContextMap()).isNullOrEmpty();

    // when
    Throwable ex = catchThrowable(() -> instance.onSuccess(successInObj));

    // then
    verify(listenableFutureCallbackMock).onSuccess(successInObj);
    if (throwException) {
        assertThat(ex).isNotNull();
    }
    else {
        assertThat(ex).isNull();
    }

    assertThat(currentSpanStackWhenListenableFutureCallbackWasCalled.get(0)).isEqualTo(spanStack);
    assertThat(currentMdcInfoWhenListenableFutureCallbackWasCalled.get(0)).isEqualTo(mdcInfo);

    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isNull();
    assertThat(MDC.getCopyOfContextMap()).isNullOrEmpty();
}
 

@Test
public void completeSubspanAttachedToCurrentThread_does_nothing_if_called_when_current_thread_span_is_null() {
    // given
    assertThat(Tracer.getInstance().getCurrentSpan()).isNull();

    // when
    WingtipsSpringWebfluxExchangeFilterFunction.completeSubspanAttachedToCurrentThread(
        request, responseMock, mock(Throwable.class), tagAndNamingStrategy, tagAndNamingAdapterMock,
        Pair.of("foo", "bar")
    );

    // then
    assertThat(strategyResponseTaggingAndFinalSpanNameMethodCalled.get()).isFalse();
    assertThat(spanRecorder.completedSpans).isEmpty();
}
 

private WingtipsExchangeFilterFunctionTracingCompletionSubscriber setupSubscriberWrapper() {
    Tracer.getInstance().startRequestWithRootSpan("TracingState for Subscriber wrapper root span");
    Tracer.getInstance().startSubSpan("TracingState for Subscriber wrapper subspan", SpanPurpose.LOCAL_ONLY);

    @SuppressWarnings("unchecked")
    WingtipsExchangeFilterFunctionTracingCompletionSubscriber result =
        new WingtipsExchangeFilterFunctionTracingCompletionSubscriber(
            mock(CoreSubscriber.class), request, mock(Context.class), TracingState.getCurrentThreadTracingState(),
            tagAndNamingStrategy, tagAndNamingAdapterMock
        );

    Tracer.getInstance().unregisterFromThread();

    return result;
}
 

@DataProvider(value = {
    "true",
    "false"
})
@Test
public void onSuccess_handles_tracing_and_mdc_info_as_expected(boolean throwException) {
    // given
    throwExceptionDuringCall = throwException;
    Tracer.getInstance().startRequestWithRootSpan("foo");
    Deque<Span> spanStack = Tracer.getInstance().getCurrentSpanStackCopy();
    Map<String, String> mdcInfo = MDC.getCopyOfContextMap();
    SuccessCallbackWithTracing instance = new SuccessCallbackWithTracing(
        successCallbackMock, spanStack, mdcInfo
    );
    resetTracing();
    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isNull();
    assertThat(MDC.getCopyOfContextMap()).isNullOrEmpty();

    // when
    Throwable ex = catchThrowable(() -> instance.onSuccess(inObj));

    // then
    verify(successCallbackMock).onSuccess(inObj);
    if (throwException) {
        assertThat(ex).isNotNull();
    }
    else {
        assertThat(ex).isNull();
    }

    assertThat(currentSpanStackWhenSuccessCallbackWasCalled.get(0)).isEqualTo(spanStack);
    assertThat(currentMdcInfoWhenSuccessCallbackWasCalled.get(0)).isEqualTo(mdcInfo);

    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isNull();
    assertThat(MDC.getCopyOfContextMap()).isNullOrEmpty();
}
 

@Autowired
public WingtipsSpringBoot2WebfluxConfiguration(WingtipsSpringBoot2WebfluxProperties wingtipsProperties) {
    this.wingtipsProperties = wingtipsProperties;
    // Set the span logging representation if specified in the wingtips properties.
    if (wingtipsProperties.getSpanLoggingFormat() != null) {
        Tracer.getInstance().setSpanLoggingRepresentation(wingtipsProperties.getSpanLoggingFormat());
    }
}
 

@Test
public void constructor_sets_values_exactly_as_given_when_subtracing_is_off() {
    // given
    CompletableFuture cfResponse = mock(CompletableFuture.class);
    AsyncResponseHandler responseHandlerFunc = mock(AsyncResponseHandler.class);
    RequestBuilderWrapper rbwMock = mock(RequestBuilderWrapper.class);
    Optional<CircuitBreaker<Response>> circuitBreaker = Optional.of(mock(CircuitBreaker.class));
    Deque<Span> spanStack = mock(Deque.class);
    Map<String, String> mdcInfo = mock(Map.class);

    Deque<Span> spanStackBeforeCall = Tracer.getInstance().getCurrentSpanStackCopy();
    Map<String, String> mdcInfoBeforeCall = MDC.getCopyOfContextMap();

    // when
    AsyncCompletionHandlerWithTracingAndMdcSupport instance = new AsyncCompletionHandlerWithTracingAndMdcSupport(
        cfResponse, responseHandlerFunc, false, rbwMock, circuitBreaker, spanStack, mdcInfo,
        tagAndNamingStrategy
    );

    // then
    assertThat(instance.completableFutureResponse).isSameAs(cfResponse);
    assertThat(instance.responseHandlerFunction).isSameAs(responseHandlerFunc);
    assertThat(instance.performSubSpanAroundDownstreamCalls).isEqualTo(false);
    assertThat(instance.circuitBreakerManualTask).isSameAs(circuitBreaker);
    assertThat(instance.distributedTraceStackToUse).isSameAs(spanStack);
    assertThat(instance.mdcContextToUse).isSameAs(mdcInfo);

    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isEqualTo(spanStackBeforeCall);
    assertThat(MDC.getCopyOfContextMap()).isEqualTo(mdcInfoBeforeCall);
}
 

@UseDataProvider("createNewSpanForRequestScenarioDataProvider")
@Test
public void createNewSpanForRequest_works_as_expected(CreateNewSpanForRequestScenario scenario) {
    // given
    doReturn(scenario.incomingTraceIdHeader).when(requestHeadersMock).getFirst(TraceHeaders.TRACE_ID);
    doReturn(scenario.incomingSpanIdHeader).when(requestHeadersMock).getFirst(TraceHeaders.SPAN_ID);
    doReturn(scenario.incomingSampledHeader).when(requestHeadersMock).getFirst(TraceHeaders.TRACE_SAMPLED);
    doReturn(scenario.incomingUserIdHeader).when(requestHeadersMock).getFirst("user-id");
    doReturn(scenario.incomingAltUserIdHeader).when(requestHeadersMock).getFirst("alt-user-id");

    assertThat(Tracer.getInstance().getCurrentSpan()).isNull();

    // when
    Span result = filterSpy.createNewSpanForRequest(exchange);

    // then
    if (scenario.incomingTraceIdHeader != null) {
        assertThat(result.getTraceId()).isEqualTo(scenario.incomingTraceIdHeader);
    }
    assertThat(result.getParentSpanId()).isEqualTo(scenario.expectedParentSpanId);
    assertThat(result.isSampleable()).isEqualTo(scenario.expectedSampleableValue);
    assertThat(result.getUserId()).isEqualTo(scenario.expectedUserId);

    if (scenario.incomingTraceIdHeader != null && scenario.incomingSpanIdHeader == null) {
        assertThat(HttpRequestTracingUtils.hasInvalidParentIdBecauseCallerDidNotSendSpanId(result)).isTrue();
        assertThat(result.getTags().get(CHILD_OF_SPAN_FROM_HEADERS_WHERE_CALLER_DID_NOT_SEND_SPAN_ID_TAG_KEY))
            .isEqualTo("true");
    }

    // The current thread tracing state should be setup appropriately.
    assertThat(Tracer.getInstance().getCurrentSpan()).isSameAs(result);
}
 

@Before
public void beforeMethod() {
    resetTracing();
    spanRecorder = new SpanRecorder();
    Tracer.getInstance().addSpanLifecycleListener(spanRecorder);

    dtConfigAdjustments = new DtConfigAdjustments();
    adjustableServerTaggingStrategy.config = dtConfigAdjustments;
    adjustableProxyTaggingStrategy.config = dtConfigAdjustments;
}
 

public static ChannelHandlerContextMocks mockChannelHandlerContextWithTraceInfo(String userId) {
    if (Tracer.getInstance().getCurrentSpan() == null) {
        Tracer.getInstance().startRequestWithRootSpan("mockChannelHandlerContext", userId);
    }

    ChannelHandlerContextMocks channelHandlerMocks = mockChannelHandlerContext();

    when(channelHandlerMocks.mockHttpProcessingState.getLoggerMdcContextMap()).thenReturn(MDC.getCopyOfContextMap());
    when(channelHandlerMocks.mockHttpProcessingState.getDistributedTraceStack()).thenReturn(Tracer.getInstance().getCurrentSpanStackCopy());

    return channelHandlerMocks;
}
 

@DataProvider(value = {
    "true",
    "false"
})
@Test
public void run_handles_tracing_and_mdc_info_as_expected(boolean throwException) {
    // given
    throwExceptionDuringCall = throwException;
    Tracer.getInstance().startRequestWithRootSpan("foo");
    Deque<Span> spanStack = Tracer.getInstance().getCurrentSpanStackCopy();
    Map<String, String> mdcInfo = MDC.getCopyOfContextMap();
    final RunnableWithTracing instance = new RunnableWithTracing(
        runnableMock, spanStack, mdcInfo
    );
    resetTracing();
    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isNull();
    assertThat(MDC.getCopyOfContextMap()).isNullOrEmpty();

    // when
    Throwable ex = catchThrowable(new ThrowingCallable() {
        @Override
        public void call() throws Throwable {
            instance.run();
        }
    });

    // then
    verify(runnableMock).run();
    if (throwException)
        assertThat(ex).isNotNull();
    else
        assertThat(ex).isNull();

    assertThat(currentSpanStackWhenRunnableWasCalled.get(0)).isEqualTo(spanStack);
    assertThat(currentMdcInfoWhenRunnableWasCalled.get(0)).isEqualTo(mdcInfo);

    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isNull();
    assertThat(MDC.getCopyOfContextMap()).isNullOrEmpty();
}
 

private Pair<Deque<Span>, Map<String, String>> generateTracingAndMdcInfo() {
    resetTracingAndMdc();
    Tracer.getInstance().startRequestWithRootSpan("someSpan");
    Pair<Deque<Span>, Map<String, String>> result = Pair.of(
        Tracer.getInstance().getCurrentSpanStackCopy(), new HashMap<>(MDC.getCopyOfContextMap())
    );
    resetTracingAndMdc();
    return result;
}
 

@UseDataProvider("extraCustomTagsScenarioDataProvider")
@Test
public void completeSubspanAttachedToCurrentThread_works_as_expected_happy_path(
    ExtraCustomTagsScenario scenario
) {
    // given
    Span rootSpan = Tracer.getInstance().startRequestWithRootSpan("fooRootSpan");
    Span currentSpan = Tracer.getInstance().startSubSpan("fooSubspan", SpanPurpose.LOCAL_ONLY);
    Throwable error = mock(Throwable.class);

    // when
    WingtipsSpringWebfluxExchangeFilterFunction.completeSubspanAttachedToCurrentThread(
        request, responseMock, error, tagAndNamingStrategy, tagAndNamingAdapterMock, scenario.extraCustomTags
    );

    // then
    assertThat(strategyResponseTaggingAndFinalSpanNameMethodCalled.get()).isTrue();
    strategyResponseTaggingArgs.get().verifyArgs(
        currentSpan, request, responseMock, error, tagAndNamingAdapterMock
    );

    if (scenario.extraCustomTags != null) {
        for (Pair<String, String> expectedTag : scenario.extraCustomTags) {
            assertThat(currentSpan.getTags().get(expectedTag.getKey())).isEqualTo(expectedTag.getValue());
        }
    }

    assertThat(currentSpan.isCompleted()).isTrue();
    assertThat(spanRecorder.completedSpans).isEqualTo(singletonList(currentSpan));
    assertThat(Tracer.getInstance().getCurrentSpan()).isEqualTo(rootSpan);
}
 

@DataProvider(value = {
    "true",
    "false"
})
@Test
public void linkTracingToCurrentThread_pair_works_as_expected_with_non_null_pair_and_null_innards(
    boolean useStaticMethod
) {
    // given
    Pair<Deque<Span>, Map<String, String>> infoForLinking = Pair.of(null, null);
    resetTracing();
    Tracer.getInstance().startRequestWithRootSpan("foo-" + UUID.randomUUID().toString());
    Pair<Deque<Span>, Map<String, String>> expectedPreCallInfo = Pair.of(
        Tracer.getInstance().getCurrentSpanStackCopy(),
        MDC.getCopyOfContextMap()
    );

    // when
    Pair<Deque<Span>, Map<String, String>> preCallInfo =
        (useStaticMethod)
        ? linkTracingToCurrentThread(infoForLinking)
        : DEFAULT_IMPL.linkTracingToCurrentThread(infoForLinking);
    
    Pair<Deque<Span>, Map<String, String>> postCallInfo = Pair.of(
        Tracer.getInstance().getCurrentSpanStackCopy(),
        MDC.getCopyOfContextMap()
    );

    // then
    assertThat(preCallInfo).isEqualTo(expectedPreCallInfo);
    assertThat(postCallInfo.getLeft()).isNull();
    assertThat(postCallInfo.getRight()).isNullOrEmpty();
}
 

@DataProvider(value = {
    "true",
    "false"
})
@Test
public void unlinkTracingAndMdcFromCurrentThread_pair_works_as_expected(boolean useNullPair) {
    // given
    Pair<Deque<Span>, Map<String, String>> infoForLinking = (useNullPair) ? null
                                                                           : setupStateWithTracingAndMdcInfo();
    // Setup the current thread with something that is not ultimately what we expect so that our assertions are
    //      verifying that the unlinkTracingAndMdcFromCurrentThread method actually did something.
    resetTracingAndMdc();
    Tracer.getInstance().startRequestWithRootSpan("foo-" + UUID.randomUUID().toString());

    // when
    AsyncNettyHelper.unlinkTracingAndMdcFromCurrentThread(infoForLinking);
    Pair<Deque<Span>, Map<String, String>> postCallInfo = Pair.of(
        Tracer.getInstance().getCurrentSpanStackCopy(),
        MDC.getCopyOfContextMap()
    );

    // then
    if (useNullPair)
        assertThat(postCallInfo).isEqualTo(Pair.of(null, Collections.emptyMap()));
    else
        assertThat(postCallInfo).isEqualTo(infoForLinking);
}
 

@DataProvider(value = {
    "true",
    "false"
})
@Test
public void apply_handles_tracing_and_mdc_info_as_expected(boolean throwException) {
    // given
    throwExceptionDuringCall = throwException;
    Tracer.getInstance().startRequestWithRootSpan("foo");
    Deque<Span> spanStack = Tracer.getInstance().getCurrentSpanStackCopy();
    Map<String, String> mdcInfo = MDC.getCopyOfContextMap();
    ChannelFutureListenerWithTracingAndMdc instance = new ChannelFutureListenerWithTracingAndMdc(
        consumerMock, spanStack, mdcInfo
    );
    resetTracingAndMdc();
    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isNull();
    assertThat(MDC.getCopyOfContextMap()).isEmpty();

    // when
    Throwable ex = catchThrowable(() -> instance.operationComplete(inObj));

    // then
    verify(consumerMock).accept(inObj);
    if (throwException) {
        assertThat(ex).isNotNull();
    }
    else {
        assertThat(ex).isNull();
    }

    assertThat(currentSpanStackWhenChannelFutureWasCalled.get(0)).isEqualTo(spanStack);
    assertThat(currentMdcInfoWhenChannelFutureWasCalled.get(0)).isEqualTo(mdcInfo);

    assertThat(Tracer.getInstance().getCurrentSpanStackCopy()).isNull();
    assertThat(MDC.getCopyOfContextMap()).isEmpty();
}
 

@Before
public void beforeMethod() {
    channelMock = mock(Channel.class);
    ctxMock = mock(ChannelHandlerContext.class);
    stateAttributeMock = mock(Attribute.class);
    state = new HttpProcessingState();
    doReturn(channelMock).when(ctxMock).channel();
    doReturn(stateAttributeMock).when(channelMock).attr(ChannelAttributes.HTTP_PROCESSING_STATE_ATTRIBUTE_KEY);
    doReturn(state).when(stateAttributeMock).get();

    functionMock = mock(Function.class);

    inObj = new Object();
    outObj = new Object();
    throwExceptionDuringCall = false;
    currentSpanStackWhenFunctionWasCalled = new ArrayList<>();
    currentMdcInfoWhenFunctionWasCalled = new ArrayList<>();
    doAnswer(invocation -> {
        currentSpanStackWhenFunctionWasCalled.add(Tracer.getInstance().getCurrentSpanStackCopy());
        currentMdcInfoWhenFunctionWasCalled.add(MDC.getCopyOfContextMap());
        if (throwExceptionDuringCall)
            throw new RuntimeException("kaboom");
        return outObj;
    }).when(functionMock).apply(inObj);

    resetTracingAndMdc();
}
 

@DataProvider(value = {
    "true   |   true",
    "true   |   false",
    "false  |   true",
    "false  |   false",
}, splitBy = "\\|")
@Test
public void doFilterInternal_should_reset_tracing_info_to_whatever_was_on_the_thread_originally(
    boolean isAsync, boolean throwExceptionInInnerFinallyBlock
) {
    // given
    final RequestTracingFilter filter = getBasicFilter();
    if (isAsync) {
        setupAsyncContextWorkflow();
    }
    RuntimeException exToThrowInInnerFinallyBlock = null;
    if (throwExceptionInInnerFinallyBlock) {
        exToThrowInInnerFinallyBlock = new RuntimeException("kaboom");
        doThrow(exToThrowInInnerFinallyBlock).when(requestMock).isAsyncStarted();
    }
    Tracer.getInstance().startRequestWithRootSpan("someOutsideSpan");
    TracingState originalTracingState = TracingState.getCurrentThreadTracingState();

    // when
    Throwable ex = catchThrowable(
        () -> filter.doFilterInternal(requestMock, responseMock, spanCapturingFilterChain)
    );

    // then
    if (throwExceptionInInnerFinallyBlock) {
        assertThat(ex).isSameAs(exToThrowInInnerFinallyBlock);
    }
    assertThat(TracingState.getCurrentThreadTracingState()).isEqualTo(originalTracingState);
    assertThat(spanCapturingFilterChain.capturedSpan).isNotNull();
    // The original tracing state was replaced on the thread before returning, but the span used by the filter chain
    //      should *not* come from the original tracing state - it should have come from the incoming headers or
    //      a new one generated.
    assertThat(spanCapturingFilterChain.capturedSpan.getTraceId())
        .isNotEqualTo(originalTracingState.spanStack.peek().getTraceId());
}
 

@Test
public void createAsyncSubSpanAndExecute_trigger_null_subspanFinisher_in_catch_block_branch_for_code_coverage() {
    // given
    Tracer.getInstance().startRequestWithRootSpan("someRootSpan");
    TracingState tracingStateBeforeInterceptorCall = TracingState.getCurrentThreadTracingState();

    WingtipsAsyncClientHttpRequestInterceptor interceptorSpy = spy(new WingtipsAsyncClientHttpRequestInterceptor(
        true, tagAndNamingStrategy, tagAndNamingAdapterMock
    ));

    RuntimeException explodingSubspanNameMethodEx =
        new RuntimeException("Intentional exception thrown by getSubspanSpanName()");

    doThrow(explodingSubspanNameMethodEx).when(interceptorSpy).getSubspanSpanName(
        any(HttpRequest.class), any(HttpTagAndSpanNamingStrategy.class), any(HttpTagAndSpanNamingAdapter.class)
    );

    HttpRequestWrapperWithModifiableHeaders wrapperRequest =
        new HttpRequestWrapperWithModifiableHeaders(requestMock);
    byte[] body = new byte[]{42};


    // when
    Throwable ex = catchThrowable(
        () -> interceptorSpy.createAsyncSubSpanAndExecute(wrapperRequest, body, executionMock)
    );

    // then
    assertThat(ex).isSameAs(explodingSubspanNameMethodEx);
    verify(interceptorSpy).getSubspanSpanName(wrapperRequest, tagAndNamingStrategy, tagAndNamingAdapterMock);

    // TracingState should have been reset even though an exception occurred in some unexpected place.
    assertThat(normalizeTracingState(TracingState.getCurrentThreadTracingState()))
        .isEqualTo(normalizeTracingState(tracingStateBeforeInterceptorCall));
}
 

@Before
public void beforeMethod() {
    channelMock = mock(Channel.class);
    ctxMock = mock(ChannelHandlerContext.class);
    stateAttributeMock = mock(Attribute.class);
    state = new HttpProcessingState();
    doReturn(channelMock).when(ctxMock).channel();
    doReturn(stateAttributeMock).when(channelMock).attr(ChannelAttributes.HTTP_PROCESSING_STATE_ATTRIBUTE_KEY);
    doReturn(state).when(stateAttributeMock).get();

    consumerMock = mock(BiConsumer.class);

    inObj1 = new Object();
    inObj2 = new Object();
    throwExceptionDuringCall = false;
    currentSpanStackWhenBiConsumerWasCalled = new ArrayList<>();
    currentMdcInfoWhenBiConsumerWasCalled = new ArrayList<>();
    doAnswer(invocation -> {
        currentSpanStackWhenBiConsumerWasCalled.add(Tracer.getInstance().getCurrentSpanStackCopy());
        currentMdcInfoWhenBiConsumerWasCalled.add(MDC.getCopyOfContextMap());
        if (throwExceptionDuringCall)
            throw new RuntimeException("kaboom");
        return null;
    }).when(consumerMock).accept(inObj1, inObj2);

    resetTracingAndMdc();
}
 

@Before
public void beforeMethod() {
    clearTracerSpanLifecycleListeners();

    spanRecorder = new SpanRecorder();
    Tracer.getInstance().addSpanLifecycleListener(spanRecorder);
}
 

/**
 * Returns the name that should be used for the span surrounding the downstream call. Defaults to {@link
 * ProxyRouterSpanNamingAndTaggingStrategy#getInitialSpanNameOverride(HttpRequest, RequestInfo, String, String)}
 * if that returns a non-null value, then falls back to whatever {@link
 * ProxyRouterSpanNamingAndTaggingStrategy#getInitialSpanName(HttpRequest)} returns, with a ultimate fallback
 * of {@link HttpRequestTracingUtils#getFallbackSpanNameForHttpRequest(String, String)} if the naming strategy
 * returned null or blank string for both the override and initial span name.
 *
 * @param downstreamRequest The Netty {@link HttpRequest} for the downstream call.
 * @param namingStrategy The {@link ProxyRouterSpanNamingAndTaggingStrategy} being used.
 * @return The name that should be used for the span surrounding the downstream call.
 */
protected @NotNull String getSubspanSpanName(
    @NotNull HttpRequest downstreamRequest,
    @NotNull RequestInfo<?> overallRequest,
    @NotNull ProxyRouterSpanNamingAndTaggingStrategy<Span> namingStrategy
) {
    String spanNameFromStrategy = namingStrategy.getInitialSpanName(downstreamRequest);
    Span overallRequestSpan = Tracer.getInstance().getCurrentSpan();
    String overallRequestSpanName = (overallRequestSpan == null) ? null : overallRequestSpan.getSpanName();

    String spanNameOverride = namingStrategy.getInitialSpanNameOverride(
        downstreamRequest, overallRequest, spanNameFromStrategy, overallRequestSpanName
    );

    if (StringUtils.isNotBlank(spanNameFromStrategy)) {
        // We got a span name from the strategy. See if it should be overridden.
        if (StringUtils.isNotBlank(spanNameOverride)) {
            return spanNameOverride;
        }

        // No override, so just use the name from the strategy.
        return spanNameFromStrategy;
    }

    // The naming strategy didn't have anything for us at all. See if there's an override.
    if (StringUtils.isNotBlank(spanNameOverride)) {
        return spanNameOverride;
    }

    // The naming strategy didn't have anything for us and there was no override. Fall back to something reasonable.
    return getFallbackSpanName(downstreamRequest);
}