BaggagedField update value not working in SpringBootTest after migrating from Spring boot 2.7.x to 3.2.3

24 views Asked by At

I have Tracing library which by default propagate the CorrelationId and x-correlation-id into headers of the requests.

That also provides a utility class which is used to set the custom correlationId which can be traced wherever the requests is passed.

The configuration and all other is working fine with actual application but when I am trying to run it with SpringBootTest the BaggagedField is not getting updated. Code is as below.

Configuration:

package com.test.tracing.configuration;

import brave.baggage.BaggageField;
import brave.baggage.BaggagePropagationConfig;
import brave.baggage.BaggagePropagationCustomizer;
import brave.propagation.Propagation;
import com.test.tracing.configuration.EnsureCorrelationIdPropagationFactory;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.lang.NonNull;


@Configuration
public class CorrelationAutoConfiguration1 {

    public static final BaggageField CORRELATION_ID_BAGGAGE_FIELD = BaggageField.create("correlationId");

    @Bean
    BaggagePropagationCustomizer baggagePropagationCustomizer() {
        return fb -> fb.add(BaggagePropagationConfig.SingleBaggageField
            .newBuilder(CORRELATION_ID_BAGGAGE_FIELD)
            .addKeyName("X-MC-Correlation-Id")
            .addKeyName("Correlation-Id")
            .build());
    }

    @Bean
    static BeanPostProcessor ensureCorrelationIdPropagationFactory() {
        return new BeanPostProcessor() {
            @Override
            public Object postProcessAfterInitialization(@NonNull Object bean, String beanName) {
                if (bean instanceof Propagation.Factory) {
                    return new EnsureCorrelationIdPropagationFactory((Propagation.Factory) bean);
                } else {
                    return bean;
                }
            }
        };
    }

}

Propagation Factory:

package com.test.tracing.configuration;

import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import org.springframework.util.StringUtils;

import java.util.UUID;

import static com.test.tracing.configuration.CorrelationAutoConfiguration.CORRELATION_ID_BAGGAGE_FIELD;

public class EnsureCorrelationIdPropagationFactory extends Propagation.Factory {

    private Propagation.Factory delegate;

    EnsureCorrelationIdPropagationFactory(Propagation.Factory delegate) {
        this.delegate = delegate;
    }

    @Override
    public boolean supportsJoin() {
        return delegate.supportsJoin();
    }

    @Override
    public boolean requires128BitTraceId() {
        return delegate.requires128BitTraceId();
    }

    @Override
    public <K> Propagation<K> create(Propagation.KeyFactory<K> keyFactory) {
        return delegate.create(keyFactory);
    }

    @Override
    public Propagation<String> get() {
        return delegate.get();
    }

    @Override
    public TraceContext decorate(TraceContext context) {
        TraceContext decoratedContext = delegate.decorate(context);
        String correlationId = CORRELATION_ID_BAGGAGE_FIELD.getValue(decoratedContext);
        if (!StringUtils.hasText(correlationId)) {
            CORRELATION_ID_BAGGAGE_FIELD.updateValue(decoratedContext, UUID.randomUUID().toString());
        }
        return decoratedContext;
    }
}

Pom.xml

<properties>
        <aspectjrt-version>1.9.9.1</aspectjrt-version>
        <brave-propagation-aws-version>0.23.4</brave-propagation-aws-version>
        <brave-version>5.14.0</brave-version>
        <zipkin-version>2.23.18</zipkin-version>
        
    </properties>

    <dependencies>
        <dependency>
            <groupId>io.micrometer</groupId>
            <artifactId>micrometer-tracing-bridge-brave</artifactId>
            
        </dependency>
        <dependency>
            <groupId>org.aspectj</groupId>
            <artifactId>aspectjrt</artifactId>
            <version>${aspectjrt-version}</version>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-actuator</artifactId>
            <scope>test</scope>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-context-slf4j</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-http</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-httpasyncclient</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-httpclient</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-jms</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-kafka-clients</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-kafka-streams</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-messaging</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-mongodb</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-rpc</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.brave</groupId>
            <artifactId>brave-instrumentation-spring-rabbit</artifactId>
            <version>${brave-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.aws</groupId>
            <artifactId>brave-propagation-aws</artifactId>
            <version>${brave-propagation-aws-version}</version>
        </dependency>
        <dependency>
            <groupId>io.zipkin.zipkin2</groupId>
            <artifactId>zipkin</artifactId>
            <version>${zipkin-version}</version>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
    </dependencies>

The Utility:

public class Correlator {

    private Correlator() {
    }


    @Nullable
    public static String getCorrelationId() {
        return CorrelationAutoConfiguration.CORRELATION_ID_BAGGAGE_FIELD.getValue();
    }

    
    public static void setCorrelationId(String correlationId) {
        CorrelationAutoConfiguration.CORRELATION_ID_BAGGAGE_FIELD.updateValue(correlationId);
    }
}

TestCase:

import com.test.tracing.configuration.CorrelationAutoConfiguration;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import io.micrometer.tracing.ScopedSpan;
import io.micrometer.tracing.Tracer;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

@SpringBootTest(classes = CorrelationAutoConfiguration.class)
@EnableAutoConfiguration
class CorrelatorTest {

    @Autowired
    private Tracer tracer;

    private ScopedSpan span;

    @BeforeEach
    public void setUp() {
        this.span = this.tracer.startScopedSpan("my-scoped-span");
    }

    @AfterEach
    void tearDown() {
        this.span.end();
    }

    @Test
    public void correlationId() {
        String correlationId = "test-5";
        Correlator.setCorrelationId(correlationId);
        assertEquals(correlationId, Correlator.getCorrelationId());
    }

}
0

There are 0 answers