Distributed Tracing in Asynchronous Java Applications (2.x)

Distributed Tracing in Asynchronous Java Applications

This article provides guidance on propagating context in asynchronous java applications. It covers two major scenario:

1. Context propagation in Async request.
2. Context propagation in scheduled events.

Note: This article explains above mentioned scenarios in SpringBoot framework. Users, can follow similar/equivalent guidance for other frameworks.

Prerequisite:

Java SpringBoot application instrumented with latest Application Insights SpringBoot Starter (1.2.0-BETA) or higher and Application Insights Java Agent (2.4.0-BETA) or higher. Please follow this article on instrumenting Application with Application Insights Java SDK. Take a look on this document on enabling Java Agent.

context propagation in Async Requests

Step 1: Ensure that prerequisites are satisfied and application has enabled asynchronous processing. Please look at official spring documentation for more details.

Step 2: Create a rest controller with asynchronous method. Here is an example of rest controller which makes asynchronous calls to GitHub lookup service:

@RestController
public class LookupController {

    private static final Logger logger = LoggerFactory.getLogger(LookupController.class);

    private final GitHubLookupService gitHubLookupService;

    public LookupController(GitHubLookupService gitHubLookupService) {
        this.gitHubLookupService = gitHubLookupService;
    }

    @Async
    @GetMapping("/fetchUsers")
    public CompletableFuture<List<User>> fetchUsers() throws InterruptedException, ExecutionException {

        // Kick of multiple, asynchronous lookups
        CompletableFuture<User> page1 = gitHubLookupService.findUser("PivotalSoftware");
        CompletableFuture<User> page2 = gitHubLookupService.findUser("CloudFoundry");
        CompletableFuture<User> page3 = gitHubLookupService.findUser("Spring-Projects");

        // Wait until they are all done
        CompletableFuture.allOf(page1,page2,page3).join();
        //return CompletableFuture.allOf(page1, page2, page3);

        // Print results, including elapsed time
        logger.info("--> " + page1.get());
        logger.info("--> " + page2.get());
        logger.info("--> " + page3.get());
        List<User> result = new ArrayList<>();
        result.add(page1.get());
        result.add(page2.get());
        result.add(page3.get());
        return CompletableFuture.completedFuture(result);
    }
}

Step 3: Override the springboot TaskExecutor default implementation with instrumented task executor. This helps in propagating context. Here is how you can do it.

@Configuration
public class configuration {

    @Bean
    public Executor taskExecutor() {

        // Use a custom ThreadPool Executor
        ThreadPoolTaskExecutor executor = new MyThreadPoolTaskExecutor();
        executor.setCorePoolSize(2);
        executor.setMaxPoolSize(2);
        executor.setQueueCapacity(500);
        executor.setThreadNamePrefix("GithubLookup-");
        executor.initialize();
        return executor;
    }

    /**
     * Custom ThreadPoolExecutor for passing a wrapped runnable
     */
    private final class MyThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

        @Override
        public void execute(Runnable command) {
            super.execute(new Wrapped(command, ThreadContext.getRequestTelemetryContext()));
        }
    }

    /**
     * A wrapper class that holds the instance of runnable and the associated context
     */
    private final class Wrapped implements Runnable {
        private final Runnable task;
        private final RequestTelemetryContext rtc;

        Wrapped(Runnable task, RequestTelemetryContext rtc) {
            this.task = task;
            this.rtc = rtc;
        }

        @Override
        public void run() {
            if (ThreadContext.getRequestTelemetryContext() != null) {
                ThreadContext.remove();
            }

            // Set the context explicitly
            ThreadContext.setRequestTelemetryContext(rtc);
            task.run();
        }
    }
}

Step 4: Build and run the application.

Context propagation in scheduled events.

Step 1: Ensure that prerequisites are satisfied. Enable scheduling in Springboot apps. Please follow this basic guidelines from official spring docs on scheduling

Step 2: Create a scheduled method which makes outbound calls (i.e HTTP, SQL etc.). Here is an example about it.

@Component
public class SampleBean {

    @Autowired
    private RestTemplate template;

    @Scheduled(initialDelay=5000, fixedDelay=10000)
    public void init() {
       // make rest call from scheduled method.
       String output = template.exchange("http://localhost:8082/hello", HttpMethod.GET, null, String.class).getBody();
       System.out.println("From RestTemplate Call:"+output);
    }
}

Step 3: Step 3: Override the springboot TaskExecutor default implementation with instrumented task executor. This helps in propagating context. Here is how it can be done.

@Configuration
public class AppConfig implements SchedulingConfigurer {

    @Override
    public void configureTasks(ScheduledTaskRegistrar scheduledTaskRegistrar) {
        // use custom ThreadPoolTaskScheduler

        ThreadPoolTaskScheduler taskScheduler = new MyThreadPoolTaskScheduler();
        taskScheduler.setThreadNamePrefix("my-scheduled-task-pool-");
        taskScheduler.initialize();

        scheduledTaskRegistrar.setTaskScheduler(taskScheduler);
    }

    /**
     * Custom ThreadPoolTaskScheduler to pass a wrapper over runnable
     */
    private static class MyThreadPoolTaskScheduler extends ThreadPoolTaskScheduler {
        @Override
        public void execute(Runnable command) {
            super.execute(new Wrapped(command, new RequestTelemetryContext(new Date().getTime(), null)));
        }

        @Override
        public ScheduledFuture<?> schedule(Runnable task, Trigger trigger) {
            return super.schedule(new Wrapped(task,
                    new RequestTelemetryContext(new Date().getTime(), null)), trigger);
        }

        @Override
        public ScheduledFuture<?> schedule(Runnable task, Date startTime) {
            return super.schedule(new Wrapped(task,
                    new RequestTelemetryContext(new Date().getTime(), null)), startTime);
        }

        @Override
        public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, Date startTime, long period) {
            return super.scheduleAtFixedRate(new Wrapped(task,
                    new RequestTelemetryContext(new Date().getTime(), null)), startTime, period);
        }

        @Override
        public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, long period) {
            return super.scheduleAtFixedRate(new Wrapped(task,
                    new RequestTelemetryContext(new Date().getTime(), null)), period);
        }

        @Override
        public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, Date startTime, long delay) {
            return super.scheduleWithFixedDelay(new Wrapped(task,
                    new RequestTelemetryContext(new Date().getTime(), null)), startTime, delay);
        }

        @Override
        public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, long delay) {
            return super.scheduleWithFixedDelay(new Wrapped(task,
                    new RequestTelemetryContext(new Date().getTime(), null)), delay);
        }
    }

    private static final class Wrapped implements Runnable {
        private final Runnable task;
        private RequestTelemetryContext rtc;

        Wrapped(Runnable task, RequestTelemetryContext rtc) {
            this.task = task;
            this.rtc = rtc;
        }

        @Override
        public void run() {
            if (ThreadContext.getRequestTelemetryContext() != null) {
                ThreadContext.remove();

                // Since this runnable is ran on schedule, update the context on every run
                rtc = new RequestTelemetryContext(new Date().getTime());
            }
            ThreadContext.setRequestTelemetryContext(rtc);
            task.run();
        }
    }
}

Step 4: Build and run the application.