From 56dd13bfc176729c5b23a36dee48995a0da424f0 Mon Sep 17 00:00:00 2001
From: Mateusz Rzeszutek <mrzeszutek@splunk.com>
Date: Tue, 12 Jan 2021 06:32:21 +0100
Subject: [PATCH] Spring Batch instrumentation tests: splits (#2017)

---
 .../spring-batch-3.0-javaagent.gradle         |  2 +
 .../test/groovy/SpringBatchApplication.groovy | 39 +++++++++++++++++++
 .../src/test/groovy/SpringBatchTest.groovy    | 35 +++++++++++++++++
 .../META-INF/batch-jobs/splitJob.xml          | 15 +++++++
 .../src/test/resources/baseContext.xml        |  2 +
 .../src/test/resources/spring-batch.xml       | 15 +++++++
 6 files changed, 108 insertions(+)
 create mode 100644 instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/META-INF/batch-jobs/splitJob.xml

diff --git a/instrumentation/spring/spring-batch-3.0/javaagent/spring-batch-3.0-javaagent.gradle b/instrumentation/spring/spring-batch-3.0/javaagent/spring-batch-3.0-javaagent.gradle
index 0720fd5e43..bc261a593e 100644
--- a/instrumentation/spring/spring-batch-3.0/javaagent/spring-batch-3.0-javaagent.gradle
+++ b/instrumentation/spring/spring-batch-3.0/javaagent/spring-batch-3.0-javaagent.gradle
@@ -13,6 +13,8 @@ dependencies {
   library group: 'org.springframework.batch', name: 'spring-batch-core', version: '3.0.0.RELEASE'
 
   testImplementation group: 'javax.inject', name: 'javax.inject', version: '1'
+  // SimpleAsyncTaskExecutor context propagation
+  testInstrumentation project(':instrumentation:spring:spring-core-2.0:javaagent')
 }
 
 tasks.withType(Test) {
diff --git a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchApplication.groovy b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchApplication.groovy
index b059477611..8f1459f705 100644
--- a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchApplication.groovy
+++ b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchApplication.groovy
@@ -17,6 +17,8 @@ import org.springframework.batch.item.ItemWriter
 import org.springframework.beans.factory.annotation.Autowired
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
+import org.springframework.core.task.AsyncTaskExecutor
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor
 import springbatch.TestItemProcessor
 import springbatch.TestItemReader
 import springbatch.TestItemWriter
@@ -118,4 +120,41 @@ class SpringBatchApplication {
       .tasklet(new TestTasklet())
       .build()
   }
+
+  // split job
+  @Bean
+  Job splitJob() {
+    jobs.get("splitJob")
+      .start(splitFlowStep1())
+      .split(asyncTaskExecutor())
+      .add(splitFlow2())
+      .build()
+      .build()
+  }
+
+  @Bean
+  Step splitFlowStep1() {
+    steps.get("splitFlowStep1")
+      .tasklet(new TestTasklet())
+      .build()
+  }
+
+  @Bean
+  Flow splitFlow2() {
+    new FlowBuilder<SimpleFlow>("splitFlow2")
+      .start(splitFlowStep2())
+      .build()
+  }
+
+  @Bean
+  Step splitFlowStep2() {
+    steps.get("splitFlowStep2")
+      .tasklet(new TestTasklet())
+      .build()
+  }
+
+  @Bean
+  AsyncTaskExecutor asyncTaskExecutor() {
+    new ThreadPoolTaskExecutor()
+  }
 }
diff --git a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchTest.groovy b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchTest.groovy
index fb31a7074b..454d154ed5 100644
--- a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchTest.groovy
+++ b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchTest.groovy
@@ -147,6 +147,41 @@ abstract class SpringBatchTest extends AgentTestRunner {
       }
     }
   }
+
+  def "should trace split flow job"() {
+    when:
+    runJob("splitJob")
+
+    then:
+    assertTraces(1) {
+      trace(0, 5) {
+        span(0) {
+          name "BatchJob splitJob"
+          kind INTERNAL
+        }
+        span(1) {
+          name ~/BatchJob splitJob\.splitFlowStep[12]/
+          kind INTERNAL
+          childOf span(0)
+        }
+        span(2) {
+          name ~/BatchJob splitJob\.splitFlowStep[12]\.Chunk/
+          kind INTERNAL
+          childOf span(1)
+        }
+        span(3) {
+          name ~/BatchJob splitJob\.splitFlowStep[12]/
+          kind INTERNAL
+          childOf span(0)
+        }
+        span(4) {
+          name ~/BatchJob splitJob\.splitFlowStep[12]\.Chunk/
+          kind INTERNAL
+          childOf span(3)
+        }
+      }
+    }
+  }
 }
 
 class JavaConfigBatchJobTest extends SpringBatchTest implements ApplicationConfigTrait {
diff --git a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/META-INF/batch-jobs/splitJob.xml b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/META-INF/batch-jobs/splitJob.xml
new file mode 100644
index 0000000000..a9f8356f46
--- /dev/null
+++ b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/META-INF/batch-jobs/splitJob.xml
@@ -0,0 +1,15 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<job id="splitJob" xmlns="http://xmlns.jcp.org/xml/ns/javaee" version="1.0">
+  <split id="split">
+    <flow id="splitFlow1">
+      <step id="splitFlowStep1">
+        <batchlet ref="jsr.TestBatchlet"/>
+      </step>
+    </flow>
+    <flow id="splitFlow2">
+      <step id="splitFlowStep2">
+        <batchlet ref="jsr.TestBatchlet"/>
+      </step>
+    </flow>
+  </split>
+</job>
diff --git a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/baseContext.xml b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/baseContext.xml
index ab5a675da0..eb1ce9bdc6 100644
--- a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/baseContext.xml
+++ b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/baseContext.xml
@@ -3,6 +3,8 @@
   xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
   xsi:schemaLocation="http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans.xsd">
 
+  <bean id="asyncTaskExecutor" class="org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor"/>
+
   <bean id="jobRepository"
     class="org.springframework.batch.core.repository.support.MapJobRepositoryFactoryBean"/>
   <bean id="jobLauncher" class="org.springframework.batch.core.launch.support.SimpleJobLauncher">
diff --git a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/spring-batch.xml b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/spring-batch.xml
index 9c96e74144..539c668212 100644
--- a/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/spring-batch.xml
+++ b/instrumentation/spring/spring-batch-3.0/javaagent/src/test/resources/spring-batch.xml
@@ -37,6 +37,21 @@
     </b:step>
   </b:flow>
 
+  <b:job id="splitJob">
+    <b:split id="split" task-executor="asyncTaskExecutor">
+      <b:flow>
+        <b:step id="splitFlowStep1">
+          <b:tasklet ref="testTasklet"/>
+        </b:step>
+      </b:flow>
+      <b:flow>
+        <b:step id="splitFlowStep2">
+          <b:tasklet ref="testTasklet"/>
+        </b:step>
+      </b:flow>
+    </b:split>
+  </b:job>
+
   <bean id="itemReader" class="springbatch.TestItemReader"/>
   <bean id="itemProcessor" class="springbatch.TestItemProcessor"/>
   <bean id="itemWriter" class="springbatch.TestItemWriter"/>