SpringBoot: Interceptor to read particular field from request and set it in the response

Question

All requests and responses handled by our Spring Rest Controller has a Common section which has certain values:

{
    "common": {
        "requestId": "foo-bar-123",
        "otherKey1": "value1",
        "otherKey2": "value2",
        "otherKey3": "value3"
    },
    ...
}

Currently all my controller functions are reading the common and copying it into the response manually. I would like to move it into an interceptor of some sort.

I tried to do this using ControllerAdvice and ThreadLocal:

@ControllerAdvice
public class RequestResponseAdvice extends RequestBodyAdviceAdapter
    implements ResponseBodyAdvice<MyGenericPojo> {

  private ThreadLocal<Common> commonThreadLocal = new ThreadLocal<>();

  /* Request */

  @Override
  public boolean supports(
      MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> aClass) {
    return MyGenericPojo.class.isAssignableFrom(methodParameter.getParameterType());
  }

  @Override
  public Object afterBodyRead(
      Object body,
      HttpInputMessage inputMessage,
      MethodParameter parameter,
      Type targetType,
      Class<? extends HttpMessageConverter<?>> converterType) {
    var common = (MyGenericPojo)body.getCommon();
    if (common.getRequestId() == null) {
       common.setRequestId(generateNewRequestId()); 
    }
    commonThreadLocal(common);
    return body;
  }

  /* Response */

  @Override
  public boolean supports(
      MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
    return MyGenericPojo.class.isAssignableFrom(returnType.getParameterType());
  }

  @Override
  public MyGenericPojo beforeBodyWrite(
      MyGenericPojo body,
      MethodParameter returnType,
      MediaType selectedContentType,
      Class<? extends HttpMessageConverter<?>> selectedConverterType,
      ServerHttpRequest request,
      ServerHttpResponse response) {
    body.setCommon(commonThreadLocal.get());
    commonThreadLocal.remove();
    return body;
  }
}

This works when I test sending one request at a time. But, is it guaranteed that afterBodyRead and beforeBodyWrite is called in the same thread, when multiple requests are coming?

If not, or even otherwise, what is the best way of doing this?

Answer 1

I think that there is no need of your own ThreadLocal you can use request attributes.

@Override
public Object afterBodyRead(
        Object body,
        HttpInputMessage inputMessage,
        MethodParameter parameter,
        Type targetType,
        Class<? extends HttpMessageConverter<?>> converterType) {

    var common = ((MyGenericPojo) body).getCommon();
    if (common.getRequestId() == null) {
        common.setRequestId(generateNewRequestId());
    }

    Optional.ofNullable((ServletRequestAttributes) RequestContextHolder.getRequestAttributes())
            .map(ServletRequestAttributes::getRequest)
            .ifPresent(request -> {request.setAttribute(Common.class.getName(), common);});

    return body;
}


@Override
public MyGenericPojo beforeBodyWrite(
        MyGenericPojo body,
        MethodParameter returnType,
        MediaType selectedContentType,
        Class<? extends HttpMessageConverter<?>> selectedConverterType,
        ServerHttpRequest request,
        ServerHttpResponse response) {

    Optional.ofNullable(RequestContextHolder.getRequestAttributes())
            .map(rc -> rc.getAttribute(Common.class.getName(), RequestAttributes.SCOPE_REQUEST))
            .ifPresent(o -> {
                Common common = (Common) o;
                body.setCommon(common);
            });

    return body;
}

EDIT

Optionals can be replaced with

RequestContextHolder.getRequestAttributes().setAttribute(Common.class.getName(),common,RequestAttributes.SCOPE_REQUEST);

RequestContextHolder.getRequestAttributes().getAttribute(Common.class.getName(),RequestAttributes.SCOPE_REQUEST);

EDIT 2

About thread safety

1) standard servlet-based Spring web application we have thread-per-request scenario. Request is processed by one of the worker threads through all the filters and routines. The processing chain will be executed by the very same thread from start to end . So afterBodyRead and beforeBodyWrite guaranteed to be executed by the very same thread for a given request.

2) Your RequestResponseAdvice by itself is stateless. We used RequestContextHolder.getRequestAttributes() which is ThreadLocal and declared as

private static final ThreadLocal<RequestAttributes> requestAttributesHolder =
        new NamedThreadLocal<>("Request attributes");

And ThreadLocal javadoc states:

his class provides thread-local variables. These variables differ from their normal counterparts in that each thread that accesses one (via its get or set method) has its own, independently initialized copy of the variable.

So I don't see any thread-safety issues into this sulotion.

Answer 2

Quick answer: RequestBodyAdvice and ResponseBodyAdvice are invoked within the same thread for one request.

You can debug the implementation at: ServletInvocableHandlerMethod#invokeAndHandle

The way you're doing it is not safe though:

• ThreadLocal should be defined as static final, otherwise it's similar to any other class property
• Exception thrown in body will skip invocation of ResponseBodyAdvice (hence the threadlocal data is not removed)

"More safe way": Make the request body supports any class (not just MyGenericPojo), in the afterBodyRead method:

• First call ThreadLocal#remove
• Check if type is MyGenericPojo then set the common data to threadlocal