Supercharge your logging MDC with Kotlin: Part 1

Supercharge your MDC: Flying toy rocket

If you’re using any JVM-based language, then you are likely familiar with logging MDC (Mapped Diagnostic Context) – a default solution for passing the logging context. However, while being widely used and well known, with good support from libraries and services (such as ELK stack), MDC still has some issues and rough edges. Let’s use Kotlin to make it safer and more convenient!

This article covers only blocking code. Reactive code and coroutines have their own issues that should be targeted differently.

What all the fuss is about?

But what exactly wrong with MDC? Well, to explain it we need to take a step back.

If you didn’t have experience with logging MDC before, here’s what it looks like:

fun petTheDogById(id: DogId) {
  MDC.put("dogId", id.toString())
  val dog = repository.getById(id)
  petTheDog(dog)
}

fun petTheDog(dog: Dog) {
  // no additional data passed
  logger.info("Dog '{}' has been pet", dog.name)
  // in log we will see something like this
  // 2024-02-28T12:00:00 [main] [dogId=123] Dog 'Hugo' has been pet
}

Looks good, right? You just set a context, and it is automatically passed to all subsequent logger calls. This works especially well for libraries.

But the issue comes from the same thing as the benefits. MDC is based on ThreadLocal, so the values are stored per thread and aren’t cleaned automatically. What does it mean? In almost all cases, you want to reuse threads by utilizing thread pools, so you will use the same thread over and over again.

For example, let’s look at the Spring HTTP endpoint implementation:

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

@RestController
@RequestMapping("/food")
class FoodController(service: FoodService) {
  @PostMapping("/{id}/resupply")
  fun resupply(@PathVariable("id") id: FoodId) {
     MDC.put("foodId", id)
     service.resupplyDogFood(id)
  }
}

Here we have two endpoints with totally independent logging context. However, the same thread may handle both endpoints and MDC values will ‘leak’ from one context to another.

// first we call 'pet' endpoint
2024-02-28T12:05:00 [main] [dogId=123] Dog 'Hugo' has been pet
// then we call 'resupply' endpoint
// it is possible that the same thread will process the call
2024-02-28T12:10:00 [main] [dogId=123, foodId=456] The food {id=456, name=Meat} has beed resupplied

But we don’t want to see food logs when we search by DogId. It’s confusing and not convenient at all. To avoid such leaks we have to clean MDC at some point. And it’s much clearer and more reliable to remove MDC values when we no longer need them than to leave our application in a dirty state. In other words – clean after the actual logic, not before.

Classic solution

You may ask: “If we need to clear MDC, then why don’t we?”
Well, below is the code we get if we implement it straightforwardly.

fun petTheDogById(id: DogId) {
  val oldContext = MDC.copyOfContextMap()
  try{
    MDC.put("dogId", id)
    //--- only this part contains application logic 
    val dog = repository.getById(id)
    petTheDog(dog)
    //---  
  } finally {
    MDC.setContextMap(oldContext)
  }
}

The code no longer looks so good — it has become verbose and cumbersome. We have to use try-finally because any function can throw an exception. And we have to restore the old context instead of cleaning it because it can contain other values we would like to preserve (like requestId). Also, such code is error-prone, as someone can easily forget to restore the old context.

Better way

Let’s try to write something better.

What we need is to limit the scope of the new context. Our intention is to run some code with enriched context and then revert it to the old one. In other words, we want to wrap our code with some MDC logic. And the easiest way to wrap the code is to use lambda syntax.

fun withMdc(key: String, value: String, wrappedCall: ()->Unit) {
  val oldContext = MDC.copyOfContextMap()
  try{
    MDC.put(key, value)
    wrappedCall() //<- application logic was here
  } finally {
    MDC.setContextMap(oldContext)
  }
}

// utilize two Kotlin features:
// 1. simplified function definition with '='
// 2. trailing lambdas can be moved outside of parenthesis
fun petTheDogById(id: DogId) = withMdc("dogId", id) {
    //--- only this part contains application logic 
    val dog = repository.getById(id)
    petTheDog(dog)
    //---  
}

Looks much better already!

Multiple MDC values

But what if we want to set more than one key? It will lead to ugly nested calls.

fun feedTheDog(dogId: DogId, food: DogFood) = withMdc("dogId" to dogId) {
    withMDC("foodId" to food.id) {
        val dog = dogRepository.getById(dogId)
        feeder.feedTheDog(dog, food)
    }
}

Let’s adjust our helper to handle multiple values.

The easiest way to pass multiple values is to use a Map since MDC is very similar to it.

fun withMdc(values: Map<String, String>, wrappedCall: ()->Unit) {
    ...
}

OK, but in this case we will have to create a Map every time and write the same code that doesn’t express any logic.

fun feedTheDog(dogId: DogId, food: DogFood) = withMdc(mapOf(
    "dogId" to dogId,
    "foodId" to food.id
)) {
    val dog = dogRepository.getById(dogId)
    feeder.feedTheDog(dog, food)
}

Yet we can pass multiple values to the mapOf function. Let’s do the same for our helper!

fun withMdc(vararg pairs: Pair<String, String>, wrappedCall: ()->Unit) {
  val oldContext = MDC.getCopyOfContextMap()
  try {
    pairs.forEach{ MDC.put(it.key, it.value) }
    wrappedCall()
  }
  finally {
    MDC.setContextMap(oldContext)
  }
}

fun feedTheDog(dogId: DogId, food: DogFood) = withMdc(
  "dogId" to dogId.toString(), // `to` is an extension function to create `Pair`
  "foodId" to food.id.toString()
){
  val dog = dogRepository.getById(dogId)
  feeder.feedTheDog(dog, food)
}

Good, this is something we can work with.

Return the result of the wrapped call

What other common cases can we face?
Obviously, sometimes we want to return the results of our code. However, the current implementation does not allow this.

// This will produce a compilation error
fun goodBoyCheck(dogId: DogId): GoodBoyStatus = withMdc("dogId", dogId) {
    service.extractGoodBoyStatus(dogId)
}

Luckily, returning the result is easy. Such a feature will make our helper much more versatile.

// use a template here because we don't know the return type in advance
fun <T> withMdc(vararg pairs: <String, String>, wrappedCall: ()->T): T {
  val oldContext = MDC.getCopyOfContextMap()
  try {
    pairs.forEach{ MDC.put(it.key, it.value) }
    return wrappedCall() // and add a `return` here
  }
  finally {
    MDC.setContextMap(oldContext)
  }
}

// No error with new implementation!
fun goodBoyCheck(dogId: DogId): GoodBoyStatus = withMdc("dogId", dogId) {
    service.extractGoodBoyStatus(dogId)
}

Handy indeed!

Handle nullable values

Another typical scenario is a nullable value. Kotlin has great compile-time support for nullability, but our implementation isn’t compatible with nullable values. Since we don’t want to see nulls in our MDC values, we’ll have to filter them out.

fun search(breedId: BreedId, favoriteFood: FoodId?): List<Dog> {
    val logContext = listOf(
        "breedId" to breedId,
        "favoriteFoodId" to favoriteFood
    )
        .filter { it.second != null }
        // to avoid casting to List<Pair<String, String>>
        .map { it.first to it.second!! }
    return withMdc(*logContext.toTypedArray()) {
        dogRepository.findAll(prepareCriteria(breedId, favoriteFood))
    }
}

Ugh! So much useless and verbose code!
Let’s add a nullability check to make it clean and pretty again.

fun <T> withMdc(vararg pairs: <String, String?>, wrappedCall: ()->T): T {
  val oldContext = MDC.getCopyOfContextMap()
  try {
    pairs
      // with this check we don't have to handle nullability in our main logic  
      .filter{ it.value != null }
      .forEach{ MDC.put(it.key, it.value) }
    return wrappedCall()
  }
  finally {
    MDC.setContextMap(oldContext)
  }
}

fun search(breedId: BreedId, favoriteFood: FoodId?): List<Dog> = withMdc(
    "breedId" to breedId,
    "favoriteFoodId" to favoriteFood
) {
   dogRepository.findAll(prepareCriteria(breedId, favoriteFood))
}

Nice! Clear, short and reliable!

Conclusion

Here, we took a well-known logging MDC solution and tried to address its limitations. We applied some Kotlin features to improve it even more by making it less error-prone and verbose. And finally, it’s pleasant to read and easy to use!

Libraries can be hindered by their release cycle, backward compatibility requirements, or sheer lack of support, but this doesn’t mean you have to abide by the existing API. There is no ‘one size fits all’. So, don’t be afraid to create your own helpers. You can always figure out something that fits this specific project or your own preferences. After all, you are the developer.

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