Supercharge your logging MDC with Kotlin: Part 2. Spring Controllers.

Photo by Andy Hermawan on Unsplash

In the previous article, we improved Kotlin MDC to make it more convenient and safe. However, there is an issue with it — it isn’t compatible with Spring Controllers. Why is that so and what can we do about it? Read further to find out!

The problem with error handling

Let’s take our MDC helper that we implemented in the previous article and use it in Spring Controller:

@RestController
@RequestMapping("/dogs")
class DogController(service: DogService) {
  @PostMapping("/{id}/pet")
  fun pet(@PathVariable("id") id: DogId) = withMdc("dogId" to id) {
    service.petDogById(id)
  }
}

What if our method throws an exception? What response will be returned to our clients?
In many cases, this will not be all the code required. Very often it’s necessary to add error handling in order to wrap internal errors and return clear HTTP codes to the client. 
A typical way to handle controller errors is by using @ControllerAdvice.

@ControllerAdvice
class DogExceptionHandler {
  @ExceptionHandler
  fun handleNotFound(ex: DogNotFoundException): ResponseEntity<String> {
    logger.error("Unable to find dog", ex)
    return ResponseEntity(HttpStatus.NOT_FOUND)
  }

  @ExceptionHandler
  fun handleRuntime(ex: RuntimeException): ResponseEntity<String> {
    logger.error("Unexpected error", ex)
    return ResponseEntity(HttpStatus.INTERNAL_ERROR)
  }
}

But if we check the logs carefully, we will notice that error logs don’t contain the MDC values.
In other words, if we call POST /dogs/123/pet and an exception is thrown, we will not see dogId in the logs.

// This is what we expect to see
202402–28T12:00:00 [main] [dogId=123] Unable to find dog
EntityNotFoundException:
 <stacktrace>
// But we actually will see this (note empty MDC)
202402–28T12:00:00 [main] [] Unable to find dog
EntityNotFoundException:
 <stacktrace>

Let’s check the details of the MDC scope and Spring Controller Advice flow to understand the reason for such behavior.

In general, MDC is defined on the thread level and is shared by any code executed by a particular thread. This causes the shuffle of the log context in the code executed by different endpoints when one thread processes first one endpoint and then switches to another (see Part 1). In order to solve this, we decided to implement a wrapper for all endpoints to clean up and set up MDC for each Controller method, i.e., to limit the MDC scope. So, now, the MDC scope is located inside the Controller, and everything works very well during successful execution.

But if the process catches an exception, it first cleans MDC (and consequently closes its context), then exits the Controller, and only after that does Spring pass the exception from Controller to Controller Advice (see Diagram 1). Unfortunately, the error handling code is outside the MDC scope.

Of course, if we add error handling to Controller’s methods, then this problem wouldn’t occur. But that is a bad practice because it would lead to code duplication in each endpoint.

Attempt 1: pass MDC within exceptions

OK, so we don’t have access to the MDC scope in ControllerAdvice. But what do we have? We have access to the actual exceptions. Can we use the exceptions themselves to pass the context?

Yes, we can do this by adding extra fields to the exceptions. Let’s try to do it.

class DogNotFoundException(
  message: String,
  val dogId: String,
) : RuntimeException(message)

@ControllerAdvice
class DogExceptionHandler {
  @ExceptionHandler
  fun handleNotFound(ex: DogNotFoundException): ResponseEntity<String> = withMdc("dogId" to ex.dogId) {
    logger.error("Unable to find dog", ex)
    ResponseEntity(HttpStatus.NOT_FOUND)
  }
  // what to do with this one?
  @ExceptionHandler
  fun handleRuntime(ex: RuntimeException): ResponseEntity<String> {
    logger.error("Unexpected error", ex)
    return ResponseEntity(HttpStatus.INTERNAL_ERROR)
  }
}

This is a draft attempt, but you already can see an inherent issue with this solution: you’ll need custom exceptions. In other words, you will have to wrap (by catch and rethrow) every non-custom exception from every library you have. Even IllegalArgumentException or ValidationException! This is definitely not a pleasant solution.

Attempt 2: intercept requests and wrap MDC in Spring Filter

In order to find alternatives, let’s take a step back. Our problem is that we don’t have access to the MDC scope in exception handling because the scope is already closed. We tried to save and pass the scope via exceptions, but it turns not to be a viable option.

But what if we try to extend the scope instead? We can use Spring Filter to intercept request processing flow and wrap entire processing with our MDC scope! Because Filter is higher on the stack than Controller and ControllerAdvice, we will be able first to initialize the log context, then move to Controller and execute the logic, after that handle errors in Controller Advice, and in the end return to Filter and clean up the context. Both logic and error handling will occur inside our log context (see Diagram 2). Awesome!

@Component
class MdcFilter : OncePerRequestFilter() {
  override fun doFilterInternal(
  request: HttpServletRequest,
  response: HttpServletResponse,
  filterChain: FilterChain
  ) {
    withMdc(/** what to put here? **/) {
      filterChain.doFilter(request, response)
    }
  }
}

This way, we can access the same context in the Controller and in ControllerAdvice. But… how exactly can we set the MDC? How can we access the values that we need?

Actually, there are 4 ways to pass some information in the request:

  • in headers
  • in query parameters
  • in path variables
  • in request body

The trick with handling requests in Spring Filter is that in this place, we don’t yet have any information as to which specific endpoint this request will be routed. However, query parameters, path variables, and request body differ depending on the endpoint that will process the request. And the same data might be located in different places at various endpoints. For example, record id can be passed as a path variable in PUT request and as part of the body in other endpoints.

Whereas the header usually contains some generic info that is the same for all endpoints. So, we can access header info in the Filter, and it doesn’t matter which endpoint processes the request. This means that in the Filter, we can reliably access only the headers to set up our log context. This is clearly not enough!

Working solution: split MDC cleanup and setup between Spring Filter and Controllers

The truth is that the best place to fill the context values is the Controller itself because we have access to all important parts of the requests. And it’s already pinpointed to a specific URL. Now we have a guarantee that the Controller method is already within the scope.

Wild idea: what if we remove our helper and set MDC directly?

@PostMapping("/{id}/pet")
fun pet(@PathVariable("id") id: DogId) {
  MDC.put("dogId", id)
  service.petDogById(id)
}

But… isn’t this something we want to avoid in the first place? There is a whole article (Part 1) explaining why we should not write the code like this!

Well, the answer is no. Firstly, it works differently here. The main issue with direct MDC manipulations is context cleanup, which we already handle in the Filter in this case. Secondly, our goal is not to avoid direct library calls but to make our day-to-day work more pleasant by making our code more readable and less error-prone.

It turns out the best way to solve the issues with error handling and the intricacies of context scope in Spring is by splitting the MDC wrapping logic between Filter and Controllers. This way, we can capture all requests and do all generic initialization and cleaning inside the Filter but set up all specific details depending on the URL inside Controllers, where we can access all these details. Now, our logging works correctly!

MDC setup refactoring

Ok, we already have a working solution, but if we examine it more closely, we will find that we can refactor it and improve its readability.

Let’s examine our controller again. This time, we’ll use another method with more parameters as an example so it’s easier to recognize a pattern. What’s wrong with this code? Can we improve something?

// defined somewhere else
object LogConstants {
  // There are multiple places where we reference our MDC keys
  // So we use a constant to keep them unified and avoid mistakes
  const val BREED_ID = "dogId"
  const val FAVORITE_FOOD_ID = "favoriteFoodId"
}

typealias BreedId = String
// Let's make things more difficult - one of the ID's is Int
typealias FoodId = Int

@PostMapping("search")
fun search(
  @QueryParameter("breedId") breedId: BreedId,
  @QueryParameter("favoriteFood", required = false) favoriteFood: FoodId?
): List<Dog> {
  MDC.put(BREED_ID to breedId)
  MDC.put(FAVORITE_FOOD_ID to favoriteFood.toString())
  return dogRepository.findAll(prepareCriteria(breedId, favoriteFood))
}

What we would like to achieve here is to minimize verbosity, improve expressiveness, and reduce error probability. In order to do that, let’s try to leave only meaningful parts of the code to visualize our end goal.

@PostMapping("search")
fun search(
  @QueryParameter("breedId") breedId: BreedId,
  @QueryParameter("favoriteFood", required = false) favoriteFood: FoodId?
): List<Dog> {
  breedId(breedId) // just a method that accepts BreedId. As short as we can make it
  favoriteFood(favoriteFood) // this method accepts FoodId, no conversion needed
  return dogRepository.findAll(prepareCriteria(breedId, favoriteFood))
}

But now it looks like we call some random methods. It’s not clear what these methods do and why we call them.
We can use static calls, like with ‘’’MDC.put’’’, or add something related to logs to the methods names. However, there is a more interesting option.

@PostMapping("search")
fun search(
  @QueryParameter("breedId") breedId: BreedId,
  @QueryParameter("favoriteFood", required = false) favoriteFood: FoodId?
): List<Dog> {
  // we don't need to pass or store the context, therefore we can apply it on place instead of creating an instance
  LogContext {
    // now it's clear what these methods do - they set a log context!
    breedId(breedId)
    favoriteFood(favoriteFood)
  }
  return dogRepository.findAll(prepareCriteria(breedId, favoriteFood))
}

It looks good, doesn’t it? And there is a way to implement such an API — Kotlin DSL. We won’t dive deep into the intricacies of Kotlin DSL because it won’t be possible to fit them into this article. Instead, we will show the implementation straight away.

object LogContext {
  fun invoke(context: LogContextBuilder.() -> Unit) {
    context.invoke(LogContextBuilder)
  }

  const val DOG_ID = "dogId"
  const val BREED_ID = "breedId"
  const val FAVORITE_FOOD_ID = "favoriteFoodId"
}

//domain specific
object LogContextBuilder {
  fun dogId(id: DogId?) = put(DOG_ID, id)
  fun breedId(id: BreedId?) = put(BREED_ID, id)
  fun favoriteFoodId(id: FoodId?) = put(FAVORITE_FOOD_ID, id)

  private fun put(key: String, value: String?) {
    value?.let { MDC.put(key, it) }
  }
}

Here, we have defined our DSL — high-level code to manipulate log context. The key idea is that this code is tightly coupled with our domain or even the specific application. By defining dedicated helper functions (dogId, breedId, etc.), we’ve narrowed the allowed arguments by specific types (DogId, BreedId, etc.). We also have hidden MDC keys, so inserting another key by mistake is impossible. It makes our code less versatile but much more safe and readable.

If you are curious to learn more about Kotlin DSL — you can check the documentation.

Final solution

And this is how our final result looks. As mentioned before, it consists of two parts. The first one is the initial context setup and MDC clean-up inside the helper.

@Component
class MdcFilter : OncePerRequestFilter() {
  override fun doFilterInternal(
    request: HttpServletRequest,
    response: HttpServletResponse,
    filterChain: FilterChain
  ) {
    // Let's add USER_AGENT header to the MDC
    // We don't really need it for our logic, but it is handy to have in logs
    withMdc(LogContext.USER_AGENT to request.getHeader(Headers.USER_AGENT)) {
      filterChain.doFilter(request, response)
    }
  }
}

One interesting remark to add. Although we use our previously developed helper method in Spring Filter, it is unnecessary because there is only one place to clean the context. Consequently, it also makes perfect sense to use try-catch directly here.

And the second essential part is the improved MDC setup using Kotlin DSL.

@PostMapping("search")
fun search(
  @QueryParameter("breedId") breedId: BreedId,
  @QueryParameter("favoriteFood", required = false) favoriteFood: FoodId?
): List<Dog> {
  // Enrich the context in the controller
  LogContext {
    breedId(breedId)
    favoriteFood(favoriteFood)
  }

  return dogRepository.findAll(prepareCriteria(breedId, favoriteFood))
}

Use cases where the helper fits the best

Well, finally, we’ve implemented a solution for Spring Controllers. But does this mean that all our efforts in Part 1 are useless, and should we remove our helper?

Definitely not! There is only one case when we face an issue — in Spring Controllers. Our little helper fits perfectly in all other cases, such as scheduled jobs, event listeners, batch requests, etc.

Here are examples of some of these cases, like the helper used for scheduled jobs.

@Scheduled(FOOD_CLEANUP_INTERVAL)
// This is not a controller, continue to use our helper
fun cleanupFood() = withMdc(JOB_NAME to "cleanup_expired_food") {
  foodService.cleanupExpired()
}

Or inside event listeners.

fun donationEventListener(event: DonationEvent) = withMdc(EVENT_ID to event.id, EVENT_TYPE to "donation") {
  celebrationService.celebrateDonation(event.amount)
}

fun puppiesEventListener(event: NewPuppiesEvent) = withMdc(EVENT_ID to event.id, EVENT_TYPE to "new_puppies") {
  event.puppies.forEach { puppy ->
    // Set specific ID for each celebration call
    // Here we add dogId to the context without overriding existing values 
    withMdc(DOG_ID to puppy.id) {
      celebrationService.celebrateNewPuppie(event.puppy)
    }
  }
}

Conclusion

It’s extremely difficult, and sometimes just impossible, to create a solution that handles all possible cases universally well. It’s much easier, however, to have multiple specialized solutions and combine them when needed. And this is what we practiced in this article.

Here we found out that the solution we implemented before doesn’t support correct error handling in Spring Controllers. And to fix this, we developed a new approach — split setup and clean-up logic between Spring Filters and Spring Controllers. We decided to catch all HTTP requests in Spring Filter and execute general initialization and clean up logic there. The setup code was added to Controllers. We enhanced this code with a specific DSL to improve readability and reduce error chance.

At the same time, the helper function written in Part 1 remains totally correct for all other cases except Spring Controllers. It’s the best option to use in schedulers, event listeners, batch requests, and many other cases.

By subscribing you agree to receive our newsletter emails and agree with our privacy policy You may unsubscribe at any time.