E2E: async context propagation on Spring Boot sample

end2end/spring_boot_postgres.py

@@ -6,6 +6,27 @@
     safe_request=Request("/api/pets/create", body={"name": "Bobby"}),
     unsafe_request=Request("/api/pets/create", body={"name": "Malicious Pet', 'Gru from the Minions') -- "})
 )
+
+for endpoint in [
+    "completable-future-single",
+    "submit-callable",
+    "thread-pool-execute",
+    "fork-join-submit",
+    "scheduled-callable",
+    "spring-task-executor",
+    "spring-async-annotation",
+]:
+    spring_boot_postgres_app.add_payload(
+        f"sql async context propagation {endpoint}",
+        safe_request=Request(
+            f"/api/pets/create/async/{endpoint}",
+            body={"name": "Bobby"}
+        ),
+        unsafe_request=Request(
+            f"/api/pets/create/async/{endpoint}",
+            body={"name": "Malicious Pet', 'Gru from the Minions') -- "}
+        )
+    )
 spring_boot_postgres_app.add_payload("command injection",
     safe_request=Request("/api/commands/execute/Johnny", method='GET'),
     unsafe_request=Request("/api/commands/execute/%27%3B%20sleep%202%3B%20%23%20", method='GET'),

sample-apps/SpringBootPostgres/src/main/java/com/example/demo/AsyncContextPropagationConfig.java

@@ -0,0 +1,23 @@
+package com.example.demo;
+
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.scheduling.annotation.EnableAsync;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
+
+import java.util.concurrent.Executor;
+
+@Configuration
+@EnableAsync
+public class AsyncContextPropagationConfig {
+    @Bean(name = "asyncContextPropagationExecutor")
+    public Executor asyncContextPropagationExecutor() {
+        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
+        executor.setThreadNamePrefix("async-context-");
+        executor.setCorePoolSize(2);
+        executor.setMaxPoolSize(2);
+        executor.setQueueCapacity(10);
+        executor.initialize();
+        return executor;
+    }
+}

sample-apps/SpringBootPostgres/src/main/java/com/example/demo/AsyncContextPropagationController.java

@@ -0,0 +1,138 @@
+package com.example.demo;
+
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.http.MediaType;
+import org.springframework.web.bind.annotation.*;
+
+import java.util.concurrent.*;
+
+@RestController
+@RequestMapping("/api/pets/create/async")
+public class AsyncContextPropagationController {
+    private final Executor springExecutor;
+    private final AsyncContextPropagationService asyncContextPropagationService;
+
+    public AsyncContextPropagationController(
+        @Qualifier("asyncContextPropagationExecutor") Executor springExecutor,
+        AsyncContextPropagationService asyncContextPropagationService
+    ) {
+        this.springExecutor = springExecutor;
+        this.asyncContextPropagationService = asyncContextPropagationService;
+    }
+
+    private record PetCreate(String name) {}
+
+    @PostMapping(
+        path = "/completable-future-single",
+        consumes = MediaType.APPLICATION_JSON_VALUE,
+        produces = MediaType.APPLICATION_JSON_VALUE
+    )
+    public PetsController.Rows completableFutureSingle(@RequestBody PetCreate pet) throws Exception {
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        try {
+            return CompletableFuture
+                .supplyAsync(() -> createPet(pet.name()), executor)
+                .get();
+        } finally {
+            executor.shutdown();
+        }
+    }
+
+    @PostMapping(
+        path = "/submit-callable",
+        consumes = MediaType.APPLICATION_JSON_VALUE,
+        produces = MediaType.APPLICATION_JSON_VALUE
+    )
+    public PetsController.Rows submitCallable(@RequestBody PetCreate pet) throws Exception {
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        try {
+            return executor.submit(() -> createPet(pet.name())).get();
+        } finally {
+            executor.shutdown();
+        }
+    }
+
+    @PostMapping(
+        path = "/thread-pool-execute",
+        consumes = MediaType.APPLICATION_JSON_VALUE,
+        produces = MediaType.APPLICATION_JSON_VALUE
+    )
+    public PetsController.Rows threadPoolExecute(@RequestBody PetCreate pet) throws Exception {
+        ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1);
+        try {
+            CompletableFuture<PetsController.Rows> future = new CompletableFuture<>();
+            executor.execute(() -> {
+                try {
+                    future.complete(createPet(pet.name()));
+                } catch (Throwable throwable) {
+                    future.completeExceptionally(throwable);
+                }
+            });
+            return future.get();
+        } finally {
+            executor.shutdown();
+        }
+    }
+
+    @PostMapping(
+        path = "/fork-join-submit",
+        consumes = MediaType.APPLICATION_JSON_VALUE,
+        produces = MediaType.APPLICATION_JSON_VALUE
+    )
+    public PetsController.Rows forkJoinSubmit(@RequestBody PetCreate pet) throws Exception {
+        return ForkJoinPool.commonPool()
+            .submit(() -> createPet(pet.name()))
+            .get();
+    }
+
+    @PostMapping(
+        path = "/scheduled-callable",
+        consumes = MediaType.APPLICATION_JSON_VALUE,
+        produces = MediaType.APPLICATION_JSON_VALUE
+    )
+    public PetsController.Rows scheduledCallable(@RequestBody PetCreate pet) throws Exception {
+        ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
+        try {
+            return executor.schedule(
+                () -> createPet(pet.name()),
+                1,
+                TimeUnit.MILLISECONDS
+            ).get();
+        } finally {
+            executor.shutdown();
+        }
+    }
+
+    @PostMapping(
+        path = "/spring-task-executor",
+        consumes = MediaType.APPLICATION_JSON_VALUE,
+        produces = MediaType.APPLICATION_JSON_VALUE
+    )
+    public PetsController.Rows springTaskExecutor(@RequestBody PetCreate pet) throws Exception {
+        CompletableFuture<PetsController.Rows> future = new CompletableFuture<>();
+        springExecutor.execute(() -> {
+            try {
+                future.complete(createPet(pet.name()));
+            } catch (Throwable throwable) {
+                future.completeExceptionally(throwable);
+            }
+        });
+        return future.get();
+    }
+
+    @PostMapping(
+        path = "/spring-async-annotation",
+        consumes = MediaType.APPLICATION_JSON_VALUE,
+        produces = MediaType.APPLICATION_JSON_VALUE
+    )
+    public PetsController.Rows springAsyncAnnotation(@RequestBody PetCreate pet) throws Exception {
+        return asyncContextPropagationService
+            .createPetWithAsyncAnnotation(pet.name())
+            .get();
+    }
+
+    private PetsController.Rows createPet(String name) {
+        Integer rowsCreated = DatabaseHelper.createPetByName(name);
+        return new PetsController.Rows(rowsCreated);
+    }
+}

sample-apps/SpringBootPostgres/src/main/java/com/example/demo/AsyncContextPropagationService.java

@@ -0,0 +1,15 @@
+package com.example.demo;
+
+import org.springframework.scheduling.annotation.Async;
+import org.springframework.stereotype.Service;
+
+import java.util.concurrent.CompletableFuture;
+
+@Service
+public class AsyncContextPropagationService {
+    @Async("asyncContextPropagationExecutor")
+    public CompletableFuture<PetsController.Rows> createPetWithAsyncAnnotation(String name) {
+        Integer rowsCreated = DatabaseHelper.createPetByName(name);
+        return CompletableFuture.completedFuture(new PetsController.Rows(rowsCreated));
+    }
+}