How do I pass MDC Context while using virtual threads in DGS and Spring Boot?

[govi20]

I'm getting a 'request_id' in request header and I want to put it in the MDC context.

I've set up an AOP Aspect using DGS annotations like @DgsData, @DgsQuery, etc to capture request-id header from ‘dfe’ and put them into MDC. But, I'm having trouble with async resolvers such as below:

Resolver:

  @DgsData(parentType = "Department", field = "employees")
  public CompletionStage<Employees> getDepartmentEmployees(
    final DataFetchingEnvironment dfe,
  ) {
    // ----> resolver is using common forkjoin pool thread. 
    Department department = dfe.getSource();
    DataLoader<Department, Employees> batchDataLoader =
      dfe.getDataLoader("employeesDataLoader");

    return batchDataLoader.load(department);
  }

Dataloader:

  @Override
  public CompletionStage<List<Employees>> load(List<Department> departments) {
    // ----> MDC is null here because it's running on DGS virtual thread
    return CompletableFuture.supplyAsync(() -> departmentService.getEmployees(departments), customThreadPool);
  }

I'm facing an issue where the MDC context isn't flowing correctly because the resolver and the 'load' method are running on different threads.

How can I set MDC in data loader? Is there a way to ensure that the load method runs on a caller method thread or custom thread pool.

The other option I'm thinking about making another Aspect for the 'load' method, but I'm not sure if it's a foolproof solution

[govi20]

@paulbakker @srinivasankavitha any suggestions?

[paulbakker]

So the short answer is that you'll need to propagate the context that's on thread locals some way. The best option is probably to leverage MicroMeter Context, which Spring Boot is now standardizing on. To integrate that with DGS you'll have to provide an ExecutorService as a bean, where that bean is configured to be context propagation enabled.
We're doing exactly that at Netflix, since our Spring Boot extensions take care of that.
Doing something like that manually will unblock this for you.

I'm also starting to realize that we should probably do something better out of the box. I'll think about that some more, but it will technically be about the same thing.
Also, I'm wondering if you happened to have tried this with the just released Spring GraphQL integration we have in DGS. I haven't tested this specifically, but it might just work there, because the new webmvc handler does context propagation (using MicroMeter Context) already.

Let me know what you find, that will be helpful input to decide on what to do out of the box.

[paulbakker]

@govi20 I experimented some more and was able to fix this nicely. See PR here: #1875

[govi20]

Thanks @paulbakker 🚀

I tested the behaviour with DGS 5.3.0 8.5.3 (webmvc) earlier today, it doesn’t seem to be working currently.

[paulbakker]

My fix will be in the next release, but heads up that Virtual Thread support in general is only available in the Spring Boot 3.x releases of DGS. 5.3.0 is Spring Boot 2, so there's no virtual thread support there at all.

[govi20]

My bad. What I meant was, it's not working in the recent version of DGS 8.5.3 out of the box.

But, following your advice, I've used the Micrometer Context Propagation Library for a custom setup by overriding dgsAsyncTaskExecutor Bean. it seems to be functioning well.

[paulbakker]

@govi20 Released in 8.5.5