Preserving Context Across Threads

When building a large production-ready stateless microservices architecture, we always come across a common challenge of preserving request context across services and threads, including context propagation to the child threads.

What Is Context Propagation?

Context propagation means passing contextual information or states across different components or services in a distributed system where applications are often composed of multiple services running on different machines or containers. These services need to communicate and collaborate to fulfill a user request or perform a business process.

Context propagation becomes crucial in such distributed systems to ensure that relevant information about a particular transaction or operation is carried along as it traverses different services. This context may include data such as:

• User authentication details
• Request identifiers
• Distributed Tracing information
• Other metadata (that helps in understanding the state and origin of a request)

Key aspects of context propagation include:

• Request Context: When a user initiates a request, it often triggers a chain of interactions across multiple services. The context of the initial request, including relevant information like user identity, request timestamp, and unique identifiers, needs to be propagated to ensure consistent behavior and tracking.
• Distributed Tracing and Logging: Context propagation is closely tied to distributed tracing and logging mechanisms. By propagating context information, it becomes easier to trace the flow of a request through various services, aiding in debugging, performance analysis, and monitoring.
• Consistency: Maintaining a consistent context across services is essential for ensuring that each service involved in handling a request has the necessary information to perform its tasks correctly. This helps avoid inconsistencies and ensures coherent behavior across the distributed system.
• Middleware and Framework Support: Many middleware and frameworks provide built-in support for context propagation. For example, in microservices architectures, frameworks like Spring Cloud, Istio, or Zipkin offer tools for managing and propagating context seamlessly.
• Statelessness: Context propagation is especially important in stateless architectures where each service should operate independently without relying on a shared state. The context helps in providing the necessary information for a service to process a request without needing to store a persistent state.

Effective context propagation contributes to the overall reliability, observability, and maintainability of distributed systems by providing a unified view of the state of a transaction as it moves through different services. It also helps in reducing the code.

The Usecase

Let's say you are building a Springboot Webflux-based Microservices/applications, and you need to ensure that the state of the user (Session Identifier, Request Identifier, LoggedIn Status, etc. ) and client ( Device Type, Client IP,  etc.)  passed in the originating request should be passed between the services.

The Challenges

• Service-to-service call: For internal service-to-service calls, the context propagation does not happen automatically.
• Propagating context within classes: To refer to the context within service and/or helper classes, you need to explicitly pass it via the method arguments. This can be handled by creating a class with a static method that stores the context in the ThreadLocal object.
• Java Stream Operations: Since Java stream functions run in separate executor threads, the Context propagation via ThreadLocal to child threads needs to be done explicitly.
  
• Webflux: Similar to Java Stream functions, Context propagation in Webflux needs to be handled via reactor Hooks.

The Idea here is how to ensure that context propagation happens automatically in the child threads and to the internal called service using a reactive web client. A similar pattern can be implemented for Non reactive code also. 

Solution

Core Java provides two classes, ThreadLocal and  InheritableThreadLocal, to store thread-scoped values. 

• ThreadLocal allows the creation of variables that are local to a thread, ensuring each thread has its own copy of the variable.
• A limitation of  ThreadLocal is that if a new thread is spawned within the scope of another thread, the child thread does not inherit the values of ThreadLocal variables from its parent.

public class ExampleThreadLocal {

    private static ThreadLocal<String> threadLocal = new ThreadLocal<>();



    public static void main(String[] args) {

        threadLocal.set("Main Thread Value");



        new Thread(() -> {

            System.out.println("Child Thread: " + threadLocal.get()); // Outputs: Child Thread: null

        }).start();



        System.out.println("Main Thread: " + threadLocal.get()); // Outputs: Main Thread: Main Thread Value

    }

}

On the other hand; 

• InheritableThreadLocal extends ThreadLocal and provides the ability for child threads to inherit values from their parent threads.

public class ExampleInheritableThreadLocal {

    private static InheritableThreadLocal<String> inheritableThreadLocal = new InheritableThreadLocal<>();



    public static void main(String[] args) {

        inheritableThreadLocal.set("Main Thread Value");



        new Thread(() -> {

            System.out.println("Child Thread: " + inheritableThreadLocal.get()); // Outputs: Child Thread: Main Thread Value

        }).start();



        System.out.println("Main Thread: " + inheritableThreadLocal.get()); // Outputs: Main Thread: Main Thread Value

    }

}

 Hence, in the scenarios where we need to ensure that context must be propagated between parent and child threads, we can use application-scoped static InheritableThreadLocal variables to hold the context and fetch it wherever needed.

@Getter
@ToString
@Builder
public class RequestContext {

  private String sessionId;
  private String correlationId;
  private String userStatus;
  private String channel;
}

public class ContextAdapter {

  final ThreadLocal<RequestContext> threadLocal = new InheritableThreadLocal<>();

  public RequestContext getCurrentContext() {
    return threadLocal.get();
  }

  public void setContext(tRequestContext requestContext) {
    threadLocal.set(requestContext);
  }

  public void clear() {
    threadLocal.remove();
  }
}

public final class Context {
  static ContextAdapter contextAdapter;

  private Context() {}

  static {
    contextAdapter = new ContextAdapter();
  }

  public static void clear() {
    if (contextAdapter == null) {
      throw new IllegalStateException();
    }
    contextAdapter.clear();
  }

  public static RequestContext getContext() {
    if (contextAdapter == null) {
      throw new IllegalStateException();
    }
    return contextAdapter.getCurrentContext();
  }

  public static void setContext(RequestContext requestContext) {
    if (cContextAdapter == null) {
      throw new IllegalStateException();
    }
    contextAdapter.setContext(requestContext);
  }

  public static ContextAdapter getContextAdapter() {
    return contextAdapter;
  }
}

We can then refer to the context by calling the static method wherever required in the code.

Context.getContext()

This solves for: 

• Propagating context within classes.
• Java Stream Operations
  
• Webflux

In order to ensure that context is propagated to external calls via webclient, automatically, we can create a custom ExchangeFilterFunctionto read the context from Context.getContext() and then add the context to the header or query params as required.

public class HeaderExchange implements ExchangeFilterFunction {

  @Override
  public Mono<ClientResponse> filter(
      ClientRequest clientRequest, ExchangeFunction exchangeFunction) {
      return Mono.deferContextual(Mono::just)
        .flatMap(
            context -> {
              RequestContext  currentContext = Context.getContext();
              ClientRequest newRequest = ClientRequest.from(clientRequest)
                        .headers(httpHeaders ->{
                          httpHeaders.add("context-session-id",currentContext.getSessionId() );
                          httpHeaders.add("context-correlation-id",currentContext.getCorrelationId() );
                        }).build();

              return exchangeFunction.exchange(newRequest);
            });
  }
}

Initializing the Context as part of WebFilter.

@Slf4j
@Component
public class RequestContextFilter implements WebFilter {


  @Override
  public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {

        String sessionId =  exchange.getRequest().getHeaders().getFirst("context-session-id");
        String correlationId =  exchange.getRequest().getHeaders().getFirst("context-correlation-id");


        RequestContext requestContext = RequestContext.builder().sessionId(sessionId).correlationId(correlationId).build()

        Context.setContext(requestContext);


        return chain.filter(exchange);
  }
}