diff --git a/docs/src/main/asciidoc/integrations.adoc b/docs/src/main/asciidoc/integrations.adoc index 322c2ff23..9a67a7807 100644 --- a/docs/src/main/asciidoc/integrations.adoc +++ b/docs/src/main/asciidoc/integrations.adoc @@ -584,11 +584,18 @@ This feature is available for all tracer implementations. If you have Spring Cloud Deployer running on the classpath, we wrap the `AppDeployer` in a trace representation. We are polling the application for its status at a default interval. You can change that default by setting the `spring.sleuth.deployer.status-poll-delay` property. In order to disable this instrumentation set `spring.sleuth.deployer.enabled` to `false`. - -[[sleuth-deployer-integration]] +[[sleuth-rsocket-integration]] == Spring RSocket This feature is available for all tracer implementations. If you have Spring RSocket running on the classpath, we wrap the inbound and outbound communication to propagate the tracing context via the metadata. In order to disable this instrumentation set `spring.sleuth.rsocket.enabled` to `false`. + +[[sleuth-batch-integration]] +== Spring Batch + +This feature is available for all tracer implementations. + +If you have Spring Batch running on the classpath, we wrap the `StepBuilderFactory` and the `JobBuilderFactory` to propagate the tracing context. +In order to disable this instrumentation set `spring.sleuth.batch.enabled` to `false`. diff --git a/spring-cloud-sleuth-autoconfigure/pom.xml b/spring-cloud-sleuth-autoconfigure/pom.xml index 525554657..4c4f0075d 100644 --- a/spring-cloud-sleuth-autoconfigure/pom.xml +++ b/spring-cloud-sleuth-autoconfigure/pom.xml @@ -93,6 +93,11 @@ spring-boot-starter-websocket true + + org.springframework.boot + spring-boot-starter-batch + true + org.springframework.cloud spring-cloud-stream @@ -451,6 +456,11 @@ archunit-junit5 test + + com.h2database + h2 + test + diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceBatchAutoConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceBatchAutoConfiguration.java new file mode 100644 index 000000000..5ede393d1 --- /dev/null +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceBatchAutoConfiguration.java @@ -0,0 +1,54 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.autoconfig.instrument.batch; + +import org.springframework.batch.core.configuration.annotation.JobBuilderFactory; +import org.springframework.beans.factory.BeanFactory; +import org.springframework.boot.autoconfigure.AutoConfigureAfter; +import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.cloud.sleuth.Tracer; +import org.springframework.cloud.sleuth.autoconfig.brave.BraveAutoConfiguration; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +/** + * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration + * Auto-configuration} that registers instrumentation for Spring Batch. + * + * @author Marcin Grzejszczak + * @since 3.1.0 + */ +@Configuration(proxyBeanMethods = false) +@ConditionalOnClass(JobBuilderFactory.class) +@ConditionalOnBean(Tracer.class) +@ConditionalOnProperty(value = "spring.sleuth.batch.enabled", matchIfMissing = true) +@AutoConfigureAfter(BraveAutoConfiguration.class) +public class TraceBatchAutoConfiguration { + + @Bean + static TraceJobBuilderFactoryBeanPostProcessor traceJobBuilderFactoryBeanPostProcessor(BeanFactory beanFactory) { + return new TraceJobBuilderFactoryBeanPostProcessor(beanFactory); + } + + @Bean + static TraceStepBuilderFactoryBeanPostProcessor traceStepBuilderFactoryBeanPostProcessor(BeanFactory beanFactory) { + return new TraceStepBuilderFactoryBeanPostProcessor(beanFactory); + } + +} diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceJobBuilderFactoryBeanPostProcessor.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceJobBuilderFactoryBeanPostProcessor.java new file mode 100644 index 000000000..14d25dce9 --- /dev/null +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceJobBuilderFactoryBeanPostProcessor.java @@ -0,0 +1,47 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.autoconfig.instrument.batch; + +import org.springframework.batch.core.configuration.annotation.JobBuilderFactory; +import org.springframework.beans.BeansException; +import org.springframework.beans.factory.BeanFactory; +import org.springframework.beans.factory.config.BeanPostProcessor; +import org.springframework.cloud.sleuth.instrument.batch.TraceJobBuilderFactory; + +/** + * Bean post processor for {@link JobBuilderFactory}. + * + * @author Marcin Grzejszczak + * @since 3.1.0 + */ +public class TraceJobBuilderFactoryBeanPostProcessor implements BeanPostProcessor { + + private final BeanFactory beanFactory; + + public TraceJobBuilderFactoryBeanPostProcessor(BeanFactory beanFactory) { + this.beanFactory = beanFactory; + } + + @Override + public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException { + if (bean instanceof JobBuilderFactory && !(bean instanceof TraceJobBuilderFactory)) { + return new TraceJobBuilderFactory(this.beanFactory, (JobBuilderFactory) bean); + } + return bean; + } + +} diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceStepBuilderFactoryBeanPostProcessor.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceStepBuilderFactoryBeanPostProcessor.java new file mode 100644 index 000000000..cb36096b7 --- /dev/null +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/batch/TraceStepBuilderFactoryBeanPostProcessor.java @@ -0,0 +1,47 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.autoconfig.instrument.batch; + +import org.springframework.batch.core.configuration.annotation.StepBuilderFactory; +import org.springframework.beans.BeansException; +import org.springframework.beans.factory.BeanFactory; +import org.springframework.beans.factory.config.BeanPostProcessor; +import org.springframework.cloud.sleuth.instrument.batch.TraceStepBuilderFactory; + +/** + * Bean post processor for {@link StepBuilderFactory}. + * + * @author Marcin Grzejszczak + * @since 3.1.0 + */ +public class TraceStepBuilderFactoryBeanPostProcessor implements BeanPostProcessor { + + private final BeanFactory beanFactory; + + public TraceStepBuilderFactoryBeanPostProcessor(BeanFactory beanFactory) { + this.beanFactory = beanFactory; + } + + @Override + public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException { + if (bean instanceof StepBuilderFactory && !(bean instanceof TraceStepBuilderFactory)) { + return new TraceStepBuilderFactory(this.beanFactory, (StepBuilderFactory) bean); + } + return bean; + } + +} diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/rsocket/TraceRSocketAutoConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/rsocket/TraceRSocketAutoConfiguration.java index 584f76199..b92d88cc4 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/rsocket/TraceRSocketAutoConfiguration.java +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/rsocket/TraceRSocketAutoConfiguration.java @@ -34,7 +34,6 @@ import org.springframework.boot.context.properties.EnableConfigurationProperties import org.springframework.boot.rsocket.server.RSocketServerCustomizer; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.autoconfig.brave.BraveAutoConfiguration; -import org.springframework.cloud.sleuth.brave.propagation.PropagationType; import org.springframework.cloud.sleuth.instrument.rsocket.TracingRSocketConnectorConfigurer; import org.springframework.cloud.sleuth.instrument.rsocket.TracingRSocketServerCustomizer; import org.springframework.cloud.sleuth.propagation.Propagator; @@ -68,13 +67,13 @@ public class TraceRSocketAutoConfiguration { return builder; } - private boolean containsZipkinPropagationType(List types) { - return types.contains(PropagationType.B3); + private boolean containsZipkinPropagationType(List types) { + return types.stream().anyMatch(s -> s.equalsIgnoreCase("b3")); } @Bean RSocketConnectorConfigurer tracingRSocketConnectorConfigurer(Propagator propagator, Tracer tracer, - @Value("${spring.sleuth.propagation.type:B3}") List types) { + @Value("${spring.sleuth.propagation.type:B3}") List types) { return new TracingRSocketConnectorConfigurer(propagator, tracer, containsZipkinPropagationType(types)); } @@ -82,7 +81,7 @@ public class TraceRSocketAutoConfiguration { // OTel @Bean RSocketServerCustomizer tracingRSocketServerCustomizer(Propagator propagator, Tracer tracer, - @Value("${spring.sleuth.propagation.type:B3}") List types) { + @Value("${spring.sleuth.propagation.type:B3}") List types) { return new TracingRSocketServerCustomizer(propagator, tracer, containsZipkinPropagationType(types)); } diff --git a/spring-cloud-sleuth-autoconfigure/src/main/resources/META-INF/additional-spring-configuration-metadata.json b/spring-cloud-sleuth-autoconfigure/src/main/resources/META-INF/additional-spring-configuration-metadata.json index ef1e667b5..fef960fe4 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/resources/META-INF/additional-spring-configuration-metadata.json +++ b/spring-cloud-sleuth-autoconfigure/src/main/resources/META-INF/additional-spring-configuration-metadata.json @@ -154,6 +154,12 @@ "type": "java.lang.Boolean", "description": "Enable Spring Cloud Config Server instrumentation.", "defaultValue": true + }, + { + "name": "spring.sleuth.batch.enabled", + "type": "java.lang.Boolean", + "description": "Enable Spring Batch instrumentation.", + "defaultValue": true } ] } diff --git a/spring-cloud-sleuth-autoconfigure/src/main/resources/META-INF/spring.factories b/spring-cloud-sleuth-autoconfigure/src/main/resources/META-INF/spring.factories index 89df7248f..4d5007f20 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/resources/META-INF/spring.factories +++ b/spring-cloud-sleuth-autoconfigure/src/main/resources/META-INF/spring.factories @@ -3,6 +3,7 @@ org.springframework.boot.autoconfigure.EnableAutoConfiguration=\ org.springframework.cloud.sleuth.autoconfig.instrument.async.TraceAsyncAutoConfiguration,\ org.springframework.cloud.sleuth.autoconfig.instrument.async.TraceAsyncCustomAutoConfiguration,\ org.springframework.cloud.sleuth.autoconfig.instrument.async.TraceAsyncDefaultAutoConfiguration,\ +org.springframework.cloud.sleuth.autoconfig.instrument.batch.TraceBatchAutoConfiguration,\ org.springframework.cloud.sleuth.autoconfig.instrument.config.TraceSpringCloudConfigAutoConfiguration,\ org.springframework.cloud.sleuth.autoconfig.instrument.circuitbreaker.TraceCircuitBreakerAutoConfiguration,\ org.springframework.cloud.sleuth.autoconfig.instrument.deployer.TraceDeployerAutoConfiguration,\ diff --git a/spring-cloud-sleuth-instrumentation/pom.xml b/spring-cloud-sleuth-instrumentation/pom.xml index 3b2d48a01..7252a853b 100644 --- a/spring-cloud-sleuth-instrumentation/pom.xml +++ b/spring-cloud-sleuth-instrumentation/pom.xml @@ -107,6 +107,11 @@ spring-boot-starter-websocket true + + org.springframework.boot + spring-boot-starter-batch + true + org.springframework.cloud spring-cloud-stream diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceJobBuilderFactory.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceJobBuilderFactory.java new file mode 100644 index 000000000..0aae8cdec --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceJobBuilderFactory.java @@ -0,0 +1,56 @@ +/* + * Copyright 2018-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.batch; + +import org.springframework.batch.core.configuration.annotation.JobBuilderFactory; +import org.springframework.batch.core.job.builder.JobBuilder; +import org.springframework.beans.factory.BeanFactory; +import org.springframework.cloud.sleuth.Tracer; + +/** + * StepBuilderFactory adding {@link TraceJobExecutionListener}. + * + * @author Marcin Grzejszczak + * @since 3.1.0 + */ +public class TraceJobBuilderFactory extends JobBuilderFactory { + + private final BeanFactory beanFactory; + + private final JobBuilderFactory delegate; + + private Tracer tracer; + + public TraceJobBuilderFactory(BeanFactory beanFactory, JobBuilderFactory delegate) { + super(null); + this.beanFactory = beanFactory; + this.delegate = delegate; + } + + @Override + public JobBuilder get(String name) { + return this.delegate.get(name).listener(new TraceJobExecutionListener(tracer())); + } + + private Tracer tracer() { + if (this.tracer == null) { + this.tracer = this.beanFactory.getBean(Tracer.class); + } + return this.tracer; + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceJobExecutionListener.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceJobExecutionListener.java new file mode 100644 index 000000000..c53edc5a9 --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceJobExecutionListener.java @@ -0,0 +1,68 @@ +/* + * Copyright 2018-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.batch; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; + +import org.springframework.batch.core.JobExecution; +import org.springframework.batch.core.JobExecutionListener; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.SpanAndScope; +import org.springframework.cloud.sleuth.Tracer; + +class TraceJobExecutionListener implements JobExecutionListener { + + private final Tracer tracer; + + private static Map SPANS = new ConcurrentHashMap<>(); + + TraceJobExecutionListener(Tracer tracer) { + this.tracer = tracer; + } + + @Override + public void beforeJob(JobExecution jobExecution) { + Span span = this.tracer.nextSpan().name(jobExecution.getJobInstance().getJobName()); + Tracer.SpanInScope spanInScope = this.tracer.withSpan(span.start()); + SPANS.put(jobExecution, new SpanAndScope(span, spanInScope)); + } + + @Override + public void afterJob(JobExecution jobExecution) { + SpanAndScope spanAndScope = SPANS.remove(jobExecution); + List throwables = jobExecution.getFailureExceptions(); + Span span = spanAndScope.getSpan(); + span.tag("batch.job.name", jobExecution.getJobInstance().getJobName()); + span.tag("batch.job.instanceId", String.valueOf(jobExecution.getJobInstance().getInstanceId())); + span.tag("batch.job.executionId", String.valueOf(jobExecution.getId())); + Tracer.SpanInScope scope = spanAndScope.getScope(); + if (!throwables.isEmpty()) { + span.error(mergedThrowables(throwables)); + } + span.end(); + scope.close(); + } + + private IllegalStateException mergedThrowables(List throwables) { + return new IllegalStateException( + throwables.stream().map(Throwable::toString).collect(Collectors.joining("\n"))); + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceStepBuilderFactory.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceStepBuilderFactory.java new file mode 100644 index 000000000..d4ecb9481 --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceStepBuilderFactory.java @@ -0,0 +1,56 @@ +/* + * Copyright 2018-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.batch; + +import org.springframework.batch.core.configuration.annotation.StepBuilderFactory; +import org.springframework.batch.core.step.builder.StepBuilder; +import org.springframework.beans.factory.BeanFactory; +import org.springframework.cloud.sleuth.Tracer; + +/** + * StepBuilderFactory adding {@link TraceStepExecutionListener}. + * + * @author Marcin Grzejszczak + * @since 3.1.0 + */ +public class TraceStepBuilderFactory extends StepBuilderFactory { + + private final BeanFactory beanFactory; + + private final StepBuilderFactory delegate; + + private Tracer tracer; + + public TraceStepBuilderFactory(BeanFactory beanFactory, StepBuilderFactory delegate) { + super(null, null); + this.beanFactory = beanFactory; + this.delegate = delegate; + } + + @Override + public StepBuilder get(String name) { + return this.delegate.get(name).listener(new TraceStepExecutionListener(tracer())); + } + + private Tracer tracer() { + if (this.tracer == null) { + this.tracer = this.beanFactory.getBean(Tracer.class); + } + return this.tracer; + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceStepExecutionListener.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceStepExecutionListener.java new file mode 100644 index 000000000..185e7c553 --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/batch/TraceStepExecutionListener.java @@ -0,0 +1,72 @@ +/* + * Copyright 2018-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.batch; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; + +import org.springframework.batch.core.ExitStatus; +import org.springframework.batch.core.Step; +import org.springframework.batch.core.StepExecution; +import org.springframework.batch.core.StepExecutionListener; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.SpanAndScope; +import org.springframework.cloud.sleuth.Tracer; + +class TraceStepExecutionListener implements StepExecutionListener { + + private final Tracer tracer; + + private static Map SPANS = new ConcurrentHashMap<>(); + + TraceStepExecutionListener(Tracer tracer) { + this.tracer = tracer; + } + + @Override + public void beforeStep(StepExecution stepExecution) { + Span span = this.tracer.nextSpan().name(stepExecution.getStepName()); + Tracer.SpanInScope spanInScope = this.tracer.withSpan(span.start()); + SPANS.put(stepExecution, new SpanAndScope(span, spanInScope)); + } + + @Override + public ExitStatus afterStep(StepExecution stepExecution) { + SpanAndScope spanAndScope = SPANS.remove(stepExecution); + List throwables = stepExecution.getFailureExceptions(); + Span span = spanAndScope.getSpan(); + span.tag("batch.step.name", stepExecution.getStepName()); + span.tag("batch.job.executionId", String.valueOf(stepExecution.getJobExecutionId())); + span.tag("batch.step.executionId", String.valueOf(stepExecution.getId())); + span.tag("batch.step.type", stepExecution.getExecutionContext().getString(Step.STEP_TYPE_KEY)); + Tracer.SpanInScope scope = spanAndScope.getScope(); + if (!throwables.isEmpty()) { + span.error(mergedThrowables(throwables)); + } + span.end(); + scope.close(); + return stepExecution.getExitStatus(); + } + + private IllegalStateException mergedThrowables(List throwables) { + return new IllegalStateException( + throwables.stream().map(Throwable::toString).collect(Collectors.joining("\n"))); + } + +} diff --git a/tests/brave/pom.xml b/tests/brave/pom.xml index b5ea5e1c9..23f1b8daa 100644 --- a/tests/brave/pom.xml +++ b/tests/brave/pom.xml @@ -38,6 +38,7 @@ spring-cloud-sleuth-instrumentation-annotation-tests spring-cloud-sleuth-instrumentation-async-tests spring-cloud-sleuth-instrumentation-baggage-tests + spring-cloud-sleuth-instrumentation-batch-tests spring-cloud-sleuth-instrumentation-config-server-tests spring-cloud-sleuth-instrumentation-circuitbreaker-tests spring-cloud-sleuth-instrumentation-circuitbreaker-reactive-tests diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/pom.xml b/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/pom.xml new file mode 100644 index 000000000..8b066d3ad --- /dev/null +++ b/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/pom.xml @@ -0,0 +1,90 @@ + + + + + 4.0.0 + + spring-cloud-sleuth-instrumentation-batch-tests + jar + Spring Cloud Sleuth Brave Batch Instrumentation Tests + Spring Cloud Sleuth Brave Batch Instrumentation Tests + + + org.springframework.cloud + spring-cloud-sleuth-tests-brave + 3.1.0-SNAPSHOT + .. + + + + true + + + + + + + maven-deploy-plugin + + true + + + + + + + + org.springframework.cloud + spring-cloud-sleuth-tests-common + ${project.version} + + + org.springframework.boot + spring-boot-starter-aop + + + org.springframework.boot + spring-boot-starter-batch + + + com.h2database + h2 + runtime + + + org.springframework.cloud + spring-cloud-starter-sleuth + + + org.springframework.boot + spring-boot-starter-test + + + io.zipkin.brave + brave-tests + + + org.awaitility + awaitility + + + + diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/batch/BatchIntegrationTests.java b/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/batch/BatchIntegrationTests.java new file mode 100644 index 000000000..917e9d8be --- /dev/null +++ b/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/batch/BatchIntegrationTests.java @@ -0,0 +1,54 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.brave.instrument.batch; + +import brave.sampler.Sampler; + +import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.cloud.sleuth.brave.BraveTestSpanHandler; +import org.springframework.cloud.sleuth.test.TestSpanHandler; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.test.context.ContextConfiguration; + +@SpringBootTest +@ContextConfiguration(classes = BatchIntegrationTests.Config.class) +public class BatchIntegrationTests extends org.springframework.cloud.sleuth.instrument.batch.BatchIntegrationTests { + + @Configuration(proxyBeanMethods = false) + @EnableBatchProcessing + static class Config { + + @Bean + TestSpanHandler testSpanHandlerSupplier(brave.test.TestSpanHandler testSpanHandler) { + return new BraveTestSpanHandler(testSpanHandler); + } + + @Bean + Sampler alwaysSampler() { + return Sampler.ALWAYS_SAMPLE; + } + + @Bean + brave.test.TestSpanHandler braveTestSpanHandler() { + return new brave.test.TestSpanHandler(); + } + + } + +} diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/src/test/resources/application.yml b/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/src/test/resources/application.yml new file mode 100644 index 000000000..6d58539e0 --- /dev/null +++ b/tests/brave/spring-cloud-sleuth-instrumentation-batch-tests/src/test/resources/application.yml @@ -0,0 +1,5 @@ +logging.level.org.springframework.cloud: DEBUG +logging.level.com.netflix.discovery.InstanceInfoReplicator: ERROR +logging.level.org.springframework.cloud.sleuth.brave.instrument.web.client.feign: TRACE + +spring.autoconfigure.exclude: org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration, org.springframework.boot.autoconfigure.orm.jpa.HibernateJpaAutoConfiguration, org.springframework.boot.autoconfigure.security.servlet.SecurityAutoConfiguration, org.springframework.boot.autoconfigure.data.web.SpringDataWebAutoConfiguration diff --git a/tests/common/pom.xml b/tests/common/pom.xml index 7c1fb98de..a4f029dc5 100644 --- a/tests/common/pom.xml +++ b/tests/common/pom.xml @@ -64,6 +64,11 @@ spring-boot-starter-websocket true + + org.springframework.boot + spring-boot-starter-batch + true + org.springframework.boot spring-boot-starter-actuator diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/batch/BatchIntegrationTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/batch/BatchIntegrationTests.java new file mode 100644 index 000000000..04b615415 --- /dev/null +++ b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/batch/BatchIntegrationTests.java @@ -0,0 +1,102 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.batch; + +import java.util.List; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Collectors; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import org.springframework.batch.core.ExitStatus; +import org.springframework.batch.core.Job; +import org.springframework.batch.core.JobExecution; +import org.springframework.batch.core.JobParameters; +import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing; +import org.springframework.batch.core.configuration.annotation.JobBuilderFactory; +import org.springframework.batch.core.configuration.annotation.StepBuilderFactory; +import org.springframework.batch.core.launch.JobLauncher; +import org.springframework.batch.repeat.RepeatStatus; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.Tracer; +import org.springframework.cloud.sleuth.exporter.FinishedSpan; +import org.springframework.cloud.sleuth.test.TestSpanHandler; +import org.springframework.context.annotation.Configuration; +import org.springframework.test.context.ContextConfiguration; + +import static org.assertj.core.api.BDDAssertions.then; + +@ContextConfiguration(classes = BatchIntegrationTests.TestConfig.class) +public abstract class BatchIntegrationTests { + + private static final Log log = LogFactory.getLog(BatchIntegrationTests.class); + + @Autowired + TestSpanHandler spans; + + @Autowired + Tracer tracer; + + @Autowired + StepBuilderFactory stepBuilderFactory; + + @Autowired + JobBuilderFactory jobBuilderFactory; + + @Autowired + JobLauncher jobLauncher; + + @BeforeEach + public void setup() { + this.spans.clear(); + } + + @Test + public void should_pass_tracing_information_when_using_batch() throws Exception { + AtomicReference spanFromTasklet = new AtomicReference<>(); + Job job = this.jobBuilderFactory.get("myJob") + .start(this.stepBuilderFactory.get("myTask").tasklet((stepContribution, chunkContext) -> { + log.info("Hello"); + spanFromTasklet.set(this.tracer.currentSpan()); + return RepeatStatus.FINISHED; + }).build()).build(); + + JobExecution jobExecution = this.jobLauncher.run(job, new JobParameters()); + + then(jobExecution.getExitStatus().getExitCode()).isEqualTo(ExitStatus.COMPLETED.getExitCode()); + then(spanFromTasklet.get()).isNotNull(); + List spans = this.spans.reportedSpans(); + then(spans).hasSize(2); + then(spans.stream().map(FinishedSpan::getTraceId).collect(Collectors.toSet())).hasSize(1); + then(spans.get(0).getName()).isEqualTo("myTask"); + then(spans.get(1).getName()).isEqualTo("myJob"); + then(this.tracer.currentSpan()).isNull(); + } + + @Configuration(proxyBeanMethods = false) + @EnableAutoConfiguration + @EnableBatchProcessing + public static class TestConfig { + + } + +} diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/sample/FlatMapTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/sample/FlatMapTests.java index 8288389e5..31bf1e94a 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/sample/FlatMapTests.java +++ b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/sample/FlatMapTests.java @@ -143,7 +143,7 @@ public abstract class FlatMapTests { sender.port = port; spans.clear(); - Awaitility.await().atMost(5, TimeUnit.SECONDS).pollInterval(1, TimeUnit.SECONDS).untilAsserted(() -> { + Awaitility.await().atMost(15, TimeUnit.SECONDS).pollInterval(1, TimeUnit.SECONDS).untilAsserted(() -> { // when LOGGER.info("Start"); spans.clear(); diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/rsocket/TraceRSocketTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/rsocket/TraceRSocketTests.java index 4b7a5180d..736975713 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/rsocket/TraceRSocketTests.java +++ b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/rsocket/TraceRSocketTests.java @@ -1,368 +1,376 @@ -/* - * Copyright 2013-2021 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.springframework.cloud.sleuth.instrument.rsocket; - -import java.net.URI; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.LinkedBlockingDeque; - -import io.rsocket.frame.FrameType; -import org.junit.jupiter.api.Test; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; -import reactor.util.context.ContextView; - -import org.springframework.boot.WebApplicationType; -import org.springframework.boot.autoconfigure.EnableAutoConfiguration; -import org.springframework.boot.builder.SpringApplicationBuilder; -import org.springframework.cloud.sleuth.Span; -import org.springframework.cloud.sleuth.TraceContext; -import org.springframework.cloud.sleuth.Tracer; -import org.springframework.cloud.sleuth.exporter.FinishedSpan; -import org.springframework.cloud.sleuth.test.TestSpanHandler; -import org.springframework.context.ConfigurableApplicationContext; -import org.springframework.context.annotation.Bean; -import org.springframework.context.annotation.Configuration; -import org.springframework.core.env.Environment; -import org.springframework.messaging.handler.annotation.MessageMapping; -import org.springframework.messaging.handler.annotation.Payload; -import org.springframework.messaging.rsocket.RSocketRequester; -import org.springframework.messaging.rsocket.RSocketRequester.Builder; -import org.springframework.messaging.rsocket.RSocketStrategies; -import org.springframework.stereotype.Controller; -import org.springframework.util.MimeType; - -import static org.assertj.core.api.BDDAssertions.then; - -public abstract class TraceRSocketTests { - - public static final String EXPECTED_TRACE_ID = "b919095138aa4c6e"; - - @Test - public void should_instrument_responder() throws Exception { - // setup - ConfigurableApplicationContext context = new SpringApplicationBuilder(MyConfig.class, testConfiguration()) - .web(WebApplicationType.REACTIVE) - .properties("server.port=0", "spring.rsocket.server.transport=websocket", - "spring.rsocket.server.mapping-path=/rsocket", "spring.jmx.enabled=false", - "spring.application.name=TraceRSocketTests", "security.basic.enabled=false", - "management.security.enabled=false") - .run(); - final TestSpanHandler spans = context.getBean(TestSpanHandler.class); - final int port = context.getBean(Environment.class).getProperty("local.server.port", Integer.class); - final TestController controller2 = context.getBean(TestController.class); - final RSocketStrategies strategies = context.getBean(RSocketStrategies.class); - - final Builder rsocketRequesterBuilder = RSocketRequester.builder().rsocketStrategies(strategies); - - final RSocketRequester rSocketRequester = rsocketRequesterBuilder - .websocket(URI.create("ws://localhost:" + port + "/rsocket")); - - // REQUEST FNF - whenRequestFnFIsSent(rSocketRequester, "api.c2.fnf").block(); - - FrameType receivedFrame = controller2.getReceivedFrames().take(); - thenSpanWasReportedWithTags(spans, "api.c2.fnf", receivedFrame); - spans.clear(); - controller2.reset(); - - // REQUEST RESPONSE - whenRequestResponseIsSent(rSocketRequester, "api.c2.rr").block(); - - receivedFrame = controller2.getReceivedFrames().take(); - thenSpanWasReportedWithTags(spans, "api.c2.rr", receivedFrame); - spans.clear(); - controller2.reset(); - - // REQUEST STREAM - whenRequestStreamIsSent(rSocketRequester, "api.c2.rs").blockLast(); - - receivedFrame = controller2.getReceivedFrames().take(); - thenSpanWasReportedWithTags(spans, "api.c2.rs", receivedFrame); - spans.clear(); - controller2.reset(); - - // REQUEST CHANNEL - whenRequestChannelIsSent(rSocketRequester, "api.c2.rc").blockLast(); - - receivedFrame = controller2.getReceivedFrames().take(); - thenSpanWasReportedWithTags(spans, "api.c2.rc", receivedFrame); - spans.clear(); - controller2.reset(); - - // REQUEST FNF - whenNonSampledRequestFnfIsSent(rSocketRequester); - controller2.getReceivedFrames().take(); - // then - thenNoSpanWasReported(spans, controller2, EXPECTED_TRACE_ID); - spans.clear(); - controller2.reset(); - - // REQUEST RESPONSE - whenNonSampledRequestResponseIsSent(rSocketRequester); - controller2.getReceivedFrames().take(); - // then - thenNoSpanWasReported(spans, controller2, EXPECTED_TRACE_ID); - spans.clear(); - controller2.reset(); - - // REQUEST STREAM - whenNonSampledRequestStreamIsSent(rSocketRequester); - controller2.getReceivedFrames().take(); - // then - thenNoSpanWasReported(spans, controller2, EXPECTED_TRACE_ID); - spans.clear(); - controller2.reset(); - - // REQUEST CHANNEL - whenNonSampledRequestChannelIsSent(rSocketRequester); - controller2.getReceivedFrames().take(); - // then - thenNoSpanWasReported(spans, controller2, EXPECTED_TRACE_ID); - spans.clear(); - controller2.reset(); - - // cleanup - context.close(); - } - - @Test - public void should_instrument_requester_and_responder() throws Exception { - // setup - ConfigurableApplicationContext context = new SpringApplicationBuilder(MyConfig.class, testConfiguration()) - .web(WebApplicationType.REACTIVE) - .properties("server.port=0", "spring.rsocket.server.transport=websocket", - "spring.rsocket.server.mapping-path=/rsocket", "spring.jmx.enabled=false", - "spring.application.name=TraceRSocketTests", "security.basic.enabled=false", - "management.security.enabled=false") - .run(); - - final org.springframework.cloud.sleuth.Tracer tracer = context - .getBean(org.springframework.cloud.sleuth.Tracer.class); - final TestSpanHandler spans = context.getBean(TestSpanHandler.class); - final int port = context.getBean(Environment.class).getProperty("local.server.port", Integer.class); - final TestController controller2 = context.getBean(TestController.class); - - final Builder rsocketRequesterBuilder = context.getBean(Builder.class); - - final RSocketRequester rSocketRequester = rsocketRequesterBuilder - .websocket(URI.create("ws://localhost:" + port + "/rsocket")); - - // REQUEST FNF - final org.springframework.cloud.sleuth.Span nextSpanFnf = tracer.nextSpan().start(); - whenRequestFnFIsSent(rSocketRequester, "api.c2.fnf") - .contextWrite(ctx -> ctx.put(TraceContext.class, nextSpanFnf.context())) - .doFinally(signalType -> nextSpanFnf.end()).block(); - controller2.getReceivedFrames().take(); - thenNoSpanWasReported(spans, controller2, nextSpanFnf.context().traceId()); - spans.clear(); - controller2.reset(); - - // REQUEST RESPONSE - final org.springframework.cloud.sleuth.Span nextSpanRR = tracer.nextSpan().start(); - whenRequestResponseIsSent(rSocketRequester, "api.c2.rr") - .contextWrite(ctx -> ctx.put(TraceContext.class, nextSpanRR.context())) - .doFinally(signalType -> nextSpanRR.end()).block(); - - controller2.getReceivedFrames().take(); - thenNoSpanWasReported(spans, controller2, nextSpanRR.context().traceId()); - spans.clear(); - controller2.reset(); - - // REQUEST STREAM - final org.springframework.cloud.sleuth.Span nextSpanRS = tracer.nextSpan().start(); - whenRequestStreamIsSent(rSocketRequester, "api.c2.rs") - .contextWrite(ctx -> ctx.put(TraceContext.class, nextSpanRS.context())) - .doFinally(signalType -> nextSpanRS.end()).blockLast(); - - controller2.getReceivedFrames().take(); - thenNoSpanWasReported(spans, controller2, nextSpanRS.context().traceId()); - spans.clear(); - controller2.reset(); - - // REQUEST CHANNEL - final org.springframework.cloud.sleuth.Span nextSpanRC = tracer.nextSpan().start(); - whenRequestChannelIsSent(rSocketRequester, "api.c2.rc") - .contextWrite(ctx -> ctx.put(TraceContext.class, nextSpanRC.context())) - .doFinally(signalType -> nextSpanRC.end()).blockLast(); - - controller2.getReceivedFrames().take(); - thenNoSpanWasReported(spans, controller2, nextSpanRC.context().traceId()); - spans.clear(); - controller2.reset(); - - // cleanup - context.close(); - } - - protected abstract Class testConfiguration(); - - private void thenSpanWasReportedWithTags(TestSpanHandler spans, String path, FrameType frameType) { - then(spans).hasSize(1); - // TODO: Preferred option would be : [api.c2.{name}] - FinishedSpan span = spans.get(0); - then(span.getName()).isEqualTo(frameType.name() + " " + path); - then(span.getTags()).containsEntry("messaging.controller.class", - "org.springframework.cloud.sleuth.instrument.rsocket.TraceRSocketTests$TestController"); - then(span.getTags()).containsKey("messaging.controller.method"); - } - - private Mono whenRequestFnFIsSent(RSocketRequester requester, String path) { - return requester.route(path).send(); - } - - private Mono whenRequestResponseIsSent(RSocketRequester requester, String path) { - return requester.route(path).retrieveMono(String.class); - } - - private Flux whenRequestStreamIsSent(RSocketRequester requester, String path) { - return requester.route(path).retrieveFlux(String.class); - } - - private Flux whenRequestChannelIsSent(RSocketRequester requester, String path) { - return requester.route(path).data(Flux.fromArray(new String[] { "test1", "test2" })).retrieveFlux(String.class); - } - - private void whenNonSampledRequestFnfIsSent(RSocketRequester requester) { - requester.route("api.c2.fnf").metadata(EXPECTED_TRACE_ID + "-" + EXPECTED_TRACE_ID + "-0", new MimeType("b3") { - @Override - public String toString() { - return "b3"; - } - }).send().block(); - } - - private void whenNonSampledRequestResponseIsSent(RSocketRequester requester) { - requester.route("api.c2.rr").metadata(EXPECTED_TRACE_ID + "-" + EXPECTED_TRACE_ID + "-0", new MimeType("b3") { - @Override - public String toString() { - return "b3"; - } - }).retrieveMono(String.class).block(); - } - - private void whenNonSampledRequestStreamIsSent(RSocketRequester requester) { - requester.route("api.c2.rs").metadata(EXPECTED_TRACE_ID + "-" + EXPECTED_TRACE_ID + "-0", new MimeType("b3") { - @Override - public String toString() { - return "b3"; - } - }).retrieveFlux(String.class).blockLast(); - } - - private void whenNonSampledRequestChannelIsSent(RSocketRequester requester) { - requester.route("api.c2.rc").metadata(EXPECTED_TRACE_ID + "-" + EXPECTED_TRACE_ID + "-0", new MimeType("b3") { - @Override - public String toString() { - return "b3"; - } - }).data(Flux.fromArray(new String[] { "test1", "test2" })).retrieveFlux(String.class).blockLast(); - } - - private void thenNoSpanWasReported(TestSpanHandler spans, TestController controller2, String expectedTraceId) { - // then(spans).isEmpty(); // FIXME: does not work for request case - then(controller2.getSpan()).isNotNull(); - then(controller2.getSpan().context().traceId()).isEqualTo(expectedTraceId); - } - - @Configuration(proxyBeanMethods = false) - @EnableAutoConfiguration - static class MyConfig { - - @Bean - TestController controller(Tracer tracer) { - return new TestController(tracer); - } - - } - - @Controller - @MessageMapping("api.c2") - static class TestController { - - final Tracer tracer; - - Span span; - - ContextView interceptedContext; - - BlockingQueue receivedFrames = new LinkedBlockingDeque<>(); - - TestController(Tracer tracer) { - this.tracer = tracer; - } - - BlockingQueue getReceivedFrames() { - return this.receivedFrames; - } - - Span getSpan() { - return this.span; - } - - void reset() { - this.span = null; - } - - @MessageMapping("fnf") - Mono testFnf() { - - this.span = this.tracer.currentSpan(); - - return Mono.deferContextual(c -> { - interceptedContext = c; - receivedFrames.offer(FrameType.REQUEST_FNF); - return Mono.empty(); - }); - } - - @MessageMapping("rr") - Mono testRR() { - this.span = this.tracer.currentSpan(); - - return Mono.deferContextual(c -> { - interceptedContext = c; - receivedFrames.offer(FrameType.REQUEST_RESPONSE); - return Mono.just("response"); - }); - } - - @MessageMapping("rs") - Flux testRS() { - this.span = this.tracer.currentSpan(); - - return Flux.deferContextual(c -> { - interceptedContext = c; - receivedFrames.offer(FrameType.REQUEST_STREAM); - return Flux.just("stream"); - }); - } - - @MessageMapping("rc") - Flux testRC(@Payload Flux inbound) { - this.span = this.tracer.currentSpan(); - - return Flux.deferContextual(c -> { - interceptedContext = c; - receivedFrames.offer(FrameType.REQUEST_CHANNEL); - return inbound; - }); - } - - } - -} +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.rsocket; + +import java.net.URI; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingDeque; + +import io.rsocket.frame.FrameType; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.util.context.ContextView; + +import org.springframework.boot.WebApplicationType; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.builder.SpringApplicationBuilder; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.TraceContext; +import org.springframework.cloud.sleuth.Tracer; +import org.springframework.cloud.sleuth.exporter.FinishedSpan; +import org.springframework.cloud.sleuth.test.TestSpanHandler; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.core.env.Environment; +import org.springframework.messaging.handler.annotation.MessageMapping; +import org.springframework.messaging.handler.annotation.Payload; +import org.springframework.messaging.rsocket.RSocketRequester; +import org.springframework.messaging.rsocket.RSocketRequester.Builder; +import org.springframework.messaging.rsocket.RSocketStrategies; +import org.springframework.stereotype.Controller; +import org.springframework.util.MimeType; + +import static org.assertj.core.api.BDDAssertions.then; + +public abstract class TraceRSocketTests { + + public static final String EXPECTED_TRACE_ID = "b919095138aa4c6e"; + + @Test + public void should_instrument_responder() throws Exception { + // setup + ConfigurableApplicationContext context = new SpringApplicationBuilder(MyConfig.class, testConfiguration()) + .web(WebApplicationType.REACTIVE) + .properties("server.port=0", "spring.rsocket.server.transport=websocket", + "spring.rsocket.server.mapping-path=/rsocket", "spring.jmx.enabled=false", + "spring.application.name=TraceRSocketTests", "security.basic.enabled=false", + "management.security.enabled=false") + .run(); + final TestSpanHandler spans = context.getBean(TestSpanHandler.class); + final int port = context.getBean(Environment.class).getProperty("local.server.port", Integer.class); + final TestController controller2 = context.getBean(TestController.class); + final RSocketStrategies strategies = context.getBean(RSocketStrategies.class); + + final Builder rsocketRequesterBuilder = RSocketRequester.builder().rsocketStrategies(strategies); + + final RSocketRequester rSocketRequester = rsocketRequesterBuilder + .websocket(URI.create("ws://localhost:" + port + "/rsocket")); + + // REQUEST FNF + whenRequestFnFIsSent(rSocketRequester, "api.c2.fnf").block(); + + FrameType receivedFrame = controller2.getReceivedFrames().take(); + thenSpanWasReportedWithTags(spans, "api.c2.fnf", receivedFrame); + spans.clear(); + controller2.reset(); + + // REQUEST RESPONSE + whenRequestResponseIsSent(rSocketRequester, "api.c2.rr").block(); + + receivedFrame = controller2.getReceivedFrames().take(); + thenSpanWasReportedWithTags(spans, "api.c2.rr", receivedFrame); + spans.clear(); + controller2.reset(); + + // REQUEST STREAM + whenRequestStreamIsSent(rSocketRequester, "api.c2.rs").blockLast(); + + receivedFrame = controller2.getReceivedFrames().take(); + thenSpanWasReportedWithTags(spans, "api.c2.rs", receivedFrame); + spans.clear(); + controller2.reset(); + + // REQUEST CHANNEL + whenRequestChannelIsSent(rSocketRequester, "api.c2.rc").blockLast(); + + receivedFrame = controller2.getReceivedFrames().take(); + thenSpanWasReportedWithTags(spans, "api.c2.rc", receivedFrame); + spans.clear(); + controller2.reset(); + + // REQUEST FNF + whenNonSampledRequestFnfIsSent(rSocketRequester); + controller2.getReceivedFrames().take(); + // then + thenNoSpanWasReported(spans, controller2, expectedTraceId()); + spans.clear(); + controller2.reset(); + + // REQUEST RESPONSE + whenNonSampledRequestResponseIsSent(rSocketRequester); + controller2.getReceivedFrames().take(); + // then + thenNoSpanWasReported(spans, controller2, expectedTraceId()); + spans.clear(); + controller2.reset(); + + // REQUEST STREAM + whenNonSampledRequestStreamIsSent(rSocketRequester); + controller2.getReceivedFrames().take(); + // then + thenNoSpanWasReported(spans, controller2, expectedTraceId()); + spans.clear(); + controller2.reset(); + + // REQUEST CHANNEL + whenNonSampledRequestChannelIsSent(rSocketRequester); + controller2.getReceivedFrames().take(); + // then + thenNoSpanWasReported(spans, controller2, expectedTraceId()); + spans.clear(); + controller2.reset(); + + // cleanup + context.close(); + } + + protected String expectedTraceId() { + return EXPECTED_TRACE_ID; + } + + protected String expectedSpanId() { + return EXPECTED_TRACE_ID; + } + + @Test + public void should_instrument_requester_and_responder() throws Exception { + // setup + ConfigurableApplicationContext context = new SpringApplicationBuilder(MyConfig.class, testConfiguration()) + .web(WebApplicationType.REACTIVE) + .properties("server.port=0", "spring.rsocket.server.transport=websocket", + "spring.rsocket.server.mapping-path=/rsocket", "spring.jmx.enabled=false", + "spring.application.name=TraceRSocketTests", "security.basic.enabled=false", + "management.security.enabled=false") + .run(); + + final org.springframework.cloud.sleuth.Tracer tracer = context + .getBean(org.springframework.cloud.sleuth.Tracer.class); + final TestSpanHandler spans = context.getBean(TestSpanHandler.class); + final int port = context.getBean(Environment.class).getProperty("local.server.port", Integer.class); + final TestController controller2 = context.getBean(TestController.class); + + final Builder rsocketRequesterBuilder = context.getBean(Builder.class); + + final RSocketRequester rSocketRequester = rsocketRequesterBuilder + .websocket(URI.create("ws://localhost:" + port + "/rsocket")); + + // REQUEST FNF + final org.springframework.cloud.sleuth.Span nextSpanFnf = tracer.nextSpan().start(); + whenRequestFnFIsSent(rSocketRequester, "api.c2.fnf") + .contextWrite(ctx -> ctx.put(TraceContext.class, nextSpanFnf.context())) + .doFinally(signalType -> nextSpanFnf.end()).block(); + controller2.getReceivedFrames().take(); + thenNoSpanWasReported(spans, controller2, nextSpanFnf.context().traceId()); + spans.clear(); + controller2.reset(); + + // REQUEST RESPONSE + final org.springframework.cloud.sleuth.Span nextSpanRR = tracer.nextSpan().start(); + whenRequestResponseIsSent(rSocketRequester, "api.c2.rr") + .contextWrite(ctx -> ctx.put(TraceContext.class, nextSpanRR.context())) + .doFinally(signalType -> nextSpanRR.end()).block(); + + controller2.getReceivedFrames().take(); + thenNoSpanWasReported(spans, controller2, nextSpanRR.context().traceId()); + spans.clear(); + controller2.reset(); + + // REQUEST STREAM + final org.springframework.cloud.sleuth.Span nextSpanRS = tracer.nextSpan().start(); + whenRequestStreamIsSent(rSocketRequester, "api.c2.rs") + .contextWrite(ctx -> ctx.put(TraceContext.class, nextSpanRS.context())) + .doFinally(signalType -> nextSpanRS.end()).blockLast(); + + controller2.getReceivedFrames().take(); + thenNoSpanWasReported(spans, controller2, nextSpanRS.context().traceId()); + spans.clear(); + controller2.reset(); + + // REQUEST CHANNEL + final org.springframework.cloud.sleuth.Span nextSpanRC = tracer.nextSpan().start(); + whenRequestChannelIsSent(rSocketRequester, "api.c2.rc") + .contextWrite(ctx -> ctx.put(TraceContext.class, nextSpanRC.context())) + .doFinally(signalType -> nextSpanRC.end()).blockLast(); + + controller2.getReceivedFrames().take(); + thenNoSpanWasReported(spans, controller2, nextSpanRC.context().traceId()); + spans.clear(); + controller2.reset(); + + // cleanup + context.close(); + } + + protected abstract Class testConfiguration(); + + private void thenSpanWasReportedWithTags(TestSpanHandler spans, String path, FrameType frameType) { + then(spans).hasSize(1); + // TODO: Preferred option would be : [api.c2.{name}] + FinishedSpan span = spans.get(0); + then(span.getName()).isEqualTo(frameType.name() + " " + path); + then(span.getTags()).containsEntry("messaging.controller.class", + "org.springframework.cloud.sleuth.instrument.rsocket.TraceRSocketTests$TestController"); + then(span.getTags()).containsKey("messaging.controller.method"); + } + + private Mono whenRequestFnFIsSent(RSocketRequester requester, String path) { + return requester.route(path).send(); + } + + private Mono whenRequestResponseIsSent(RSocketRequester requester, String path) { + return requester.route(path).retrieveMono(String.class); + } + + private Flux whenRequestStreamIsSent(RSocketRequester requester, String path) { + return requester.route(path).retrieveFlux(String.class); + } + + private Flux whenRequestChannelIsSent(RSocketRequester requester, String path) { + return requester.route(path).data(Flux.fromArray(new String[] { "test1", "test2" })).retrieveFlux(String.class); + } + + private void whenNonSampledRequestFnfIsSent(RSocketRequester requester) { + requester.route("api.c2.fnf").metadata(expectedTraceId() + "-" + expectedSpanId() + "-0", new MimeType("b3") { + @Override + public String toString() { + return "b3"; + } + }).send().block(); + } + + private void whenNonSampledRequestResponseIsSent(RSocketRequester requester) { + requester.route("api.c2.rr").metadata(expectedTraceId() + "-" + expectedSpanId() + "-0", new MimeType("b3") { + @Override + public String toString() { + return "b3"; + } + }).retrieveMono(String.class).block(); + } + + private void whenNonSampledRequestStreamIsSent(RSocketRequester requester) { + requester.route("api.c2.rs").metadata(expectedTraceId() + "-" + expectedSpanId() + "-0", new MimeType("b3") { + @Override + public String toString() { + return "b3"; + } + }).retrieveFlux(String.class).blockLast(); + } + + private void whenNonSampledRequestChannelIsSent(RSocketRequester requester) { + requester.route("api.c2.rc").metadata(expectedTraceId() + "-" + expectedSpanId() + "-0", new MimeType("b3") { + @Override + public String toString() { + return "b3"; + } + }).data(Flux.fromArray(new String[] { "test1", "test2" })).retrieveFlux(String.class).blockLast(); + } + + private void thenNoSpanWasReported(TestSpanHandler spans, TestController controller2, String expectedTraceId) { + // then(spans).isEmpty(); // FIXME: does not work for request case + then(controller2.getSpan()).isNotNull(); + then(controller2.getSpan().context().traceId()).isEqualTo(expectedTraceId); + } + + @Configuration(proxyBeanMethods = false) + @EnableAutoConfiguration + static class MyConfig { + + @Bean + TestController controller(Tracer tracer) { + return new TestController(tracer); + } + + } + + @Controller + @MessageMapping("api.c2") + static class TestController { + + final Tracer tracer; + + Span span; + + ContextView interceptedContext; + + BlockingQueue receivedFrames = new LinkedBlockingDeque<>(); + + TestController(Tracer tracer) { + this.tracer = tracer; + } + + BlockingQueue getReceivedFrames() { + return this.receivedFrames; + } + + Span getSpan() { + return this.span; + } + + void reset() { + this.span = null; + } + + @MessageMapping("fnf") + Mono testFnf() { + + this.span = this.tracer.currentSpan(); + + return Mono.deferContextual(c -> { + interceptedContext = c; + receivedFrames.offer(FrameType.REQUEST_FNF); + return Mono.empty(); + }); + } + + @MessageMapping("rr") + Mono testRR() { + this.span = this.tracer.currentSpan(); + + return Mono.deferContextual(c -> { + interceptedContext = c; + receivedFrames.offer(FrameType.REQUEST_RESPONSE); + return Mono.just("response"); + }); + } + + @MessageMapping("rs") + Flux testRS() { + this.span = this.tracer.currentSpan(); + + return Flux.deferContextual(c -> { + interceptedContext = c; + receivedFrames.offer(FrameType.REQUEST_STREAM); + return Flux.just("stream"); + }); + } + + @MessageMapping("rc") + Flux testRC(@Payload Flux inbound) { + this.span = this.tracer.currentSpan(); + + return Flux.deferContextual(c -> { + interceptedContext = c; + receivedFrames.offer(FrameType.REQUEST_CHANNEL); + return inbound; + }); + } + + } + +}