Logging with MDC in functional style

Introduction

If you are already familiar with MDC (Mapped Diagnostic Context) you can go to the next section of this article. If not, please take a look at this first to get a better understanding of what MDC is.

Use case

In one of our services we had an endpoint responsible for invoking process related to car auction with a lot of logs for debugging purposes. We already had a request-id to trace it in Kibana, but we wanted to filter logs in a more convenient way, by value that we know even before invoking this endpoint which was the stock number of the car. So we decided to use MDC as it allows us to add diagnostic key-value pairs to every log in given Thread.

MDC code drawbacks

Although MDC is a great tool, it forces programmers to write clunky code due to the fact that every key-value pair that was set should be cleared, otherwise in some edge cases it could be used by another thread.

Clearing can be done by invoking MDC.clear() or by MDC.remove(key). First one clears all MDC keys so we could accidentally remove keys that were set higher in stack. Second one is better then but it has to be invoked for every key. And also in both cases it should be done in finally block to be sure that no matter what happened, context was cleared.

try {
  MDC.put("key1", "value1");
  MDC.put("key1", "value2");
  //some process
  log.info("very important log");
}finally {
  MDC.remove("key1");
  MDC.remove("key2");
}

MDC also has method putCloseable(key, value) which returns AutoClosable object but it allows only one key-value pair and has to be assigned to a variable, which may look better for a single key but it is not perfect. IDE will complain about never used variable

try (MDC.MDCCloseable closeable = MDC.putCloseable("key", "value")) {
  //some process
  log.info("very important log");
}

Solution

To stand against MDC drawbacks, we came up with this functional style solution:

LogContext.withLogKey("key1", "value1")
  .andKey("key2", "value2")
  .andKey("key3", "value3")
  .execute(() -> {
    // some process
    log.info("very important log");
  });

At first glance we can see that there is no Try Catch block anymore and we are only setting key-value pairs and not clearing them, however everything works exactly the same. So how is it done?

Solution is pretty simple. It is based on a container class for storing MDC context keys mixed with Step Builder Pattern. Let's take a look at implementation.

LogContextSteps

@NoArgsConstructor(access = AccessLevel.PRIVATE)
public  final  class  LogContextSteps {

  public  interface  ContextStep {
    ExecuteStep andKey(final String key, final Object value);
  }

  public  interface  ExecuteStep  extends  ContextStep {

    <T> T execute(final Supplier<T> action);

    default  void  execute(final Runnable action) {
      if (action == null) {
        throw  new IllegalArgumentException("Action for execution cannot be null.");
      }

      execute(() -> {
        action.run();
        return  null;
      });
    }
  }
}

LogContextSteps is a wrapping class for two step interfaces:

• ContextStep which contains method for adding new log key-value pair

• ExecuteStep which also extends ContextStep and serves execute method responsible for invoking action that we want to wrap with MDC.

LogContext

public  class  LogContext  implements  ContextStep, ExecuteStep {
  private  static  final MDCLogContextSetter MDC_LOG_CONTEXT_SETTER = new MDCLogContextSetter();

  private  final LogContextSetter contextSetter;
  private  final Set<String> contextKeys = new HashSet<>();

  private  LogContext(final LogContextSetter contextSetter) {
    this.contextSetter = contextSetter;
  }

  public  static ExecuteStep withLogKey(String key, Object value) {
    return withContextSetter(MDC_LOG_CONTEXT_SETTER).andKey(key,value);
  }

  static ExecuteStep withContextSetter(LogContextSetter contextSetter){
    return  new LogContext(contextSetter);
  }

  @Override
  public ExecuteStep andKey(String key, Object value) {
    contextSetter.put(key, value);
    contextKeys.add(key);
    return  this;
  }

  @Override
  public <T> T execute(final Supplier<T> action) {
    if (action == null) {
      throw  new IllegalArgumentException("Action for execution cannot be null.");
    }

    try {
      return action.get();
    } finally {
      contextKeys.forEach(contextSetter::remove);
    }
  }

  interface  LogContextSetter {

    void  remove(final String key);

    void  put(final String key, final Object value);
  }

  private  static  class  MDCLogContextSetter  implements  LogContextSetter {

    @Override
    public  void  remove(String key) {
      MDC.remove(key);
    }

    @Override
    public  void  put(String key, Object value) {
      MDC.put(key, Objects.toString(value));
    }
  }
}

LogContext is our main class here. It:

• implements interfaces from LogContextSteps

• stores MDC keys in Set

• executes action that we want to wrap with MDC context

For testing purposes there are some package protected classes and methods like LogContextSetter interface and its implementation MDCLogContextSetter, as it would be hard to test static methods of MDC class.

So, how does this flow work?

We begin with the static factory method LogContext.withLogKey("key1", "value1") which sets up our MDCLogContextSetter and invokes the andKey("key2", "value2") method for the first time. Then we can chain next invocations of andKey where every invocation causes two things:

• MDC key is stored in Set<String> contextKeys as only keys are needed to clear context

• LogContextSetter::put method is invoked where implementation is: MDC.put(key, Objects.toString(value)) and at that moment key-value is stored in MDC

Finally we can invoke execute(() -> action()) which is classic try-finally block where try invokes our action and finally clears MDC iterating every through contextKeys and invoking LogContextSetter::remove where implementation id MDC.remove(key)

Practical usage (filtering in Kibana)

Assuming that we have an ELK stack for monitoring and proper log configuration we can go to Kibana to find interesting logs via our MDC key-value pairs. Depending on the log configuration, key can be directly accessible or under the flat field (which was in our case)

For key-value pair: stock-number - SN12356 Kibana filter may look like this: