SpringWebFluxでコンテキストを併用するとMDCをアプリケーション間で伝播できるらしいじゃん

まえがき

SpringWebFluxでMDCを使用するにはコンテキストを使うんだってさって記事でふれた、MDCをアプリケーション間で伝播させる話について書きます。
内容的に上記記事の知見が必要です。

先駆者の話

• Line公式アカウントのチャットシステムにおけるSpringおよびWebFluxの活用事例
  
  • MDCにリクエスト情報乗っけるには？という話題。WebFilter,Hook,Contextを使って解決するパターン
• Passing Context with Spring WebFlux
  • HttpHeaderを利用して自動的にコンテキストに情報を乗っけちゃおうぜ、という試み。本記事で紹介する。

本題

昨今流行りのマイクロサービスには色んな複雑さがつきまとうが、その一つがログである。
Aシステム->Bシステム->Cシステムと流れてくる通信を、あとから一本のログとしてみたいという需要は当然としてある話で、それを解決するには色々な方法が取られてきた。(Spring Sleuthとかね)

それと同様に、色んなデバッグ情報をともに引き継がせちゃおうぜ、という需要もあり、我々のシステムもまさにソレに巻き込まれた。

今回、Passing Context with Spring WebFluxで紹介されている方法の一つを取り上げてみる。

具体的なコード

端的に今から紹介するWebFilterのやっていることを書くと

• リクエストについて、X-MDC-というプレフィックスがついたHttpHeaderをすべて取得し、プレフィックスを削ってコンテキストにぶち込んでいる。コンテキストのKey値は定数CONTEXT_MAPである。
• レスポンスについて、定数CONTEXT_MAPから取り出した値を、プレフィックスX-MDC-をつけて返している。

MdcHeaderFilter.java

import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;

import java.util.Map;

import static java.util.stream.Collectors.toMap;
import static jp.co.example.helper.LogHelper.*;

@Component
@Order(Ordered.HIGHEST_PRECEDENCE)
public class MdcHeaderFilter implements WebFilter {
  private static final String MDC_HEADER_PREFIX = "X-MDC-";

  @Override
  public Mono<Void> filter(
    ServerWebExchange ex,
    WebFilterChain chain) {
    ex.getResponse().beforeCommit(
      () -> addContextToHttpResponseHeaders(ex.getResponse())
    );

    return chain.filter(ex)
      .subscriberContext(
        ctx -> addRequestHeadersToContext(ex.getRequest(), ctx)
      );
  }

  private Context addRequestHeadersToContext(
    final ServerHttpRequest request,
    final Context context) {

    final Map<String, String> contextMap = request
      .getHeaders().toSingleValueMap().entrySet()
      .stream()
      .filter(x -> x.getKey().startsWith(MDC_HEADER_PREFIX))
      .collect(
        toMap(v -> v.getKey().substring(MDC_HEADER_PREFIX.length()),
          Map.Entry::getValue
        )
      );

    // 例えばCookieをいれたければこうすればよい
    String cookie = request.getCookies().containsKey("EXAMPLE") ?
      request.getCookies().getFirst("EXAMPLE").getValue() : "none";
    contextMap.put("example-cookie", cookie);

    return context.put(CONTEXT_MAP, contextMap);
  }

  private Mono<Void> addContextToHttpResponseHeaders(
    final ServerHttpResponse res) {

    return Mono.subscriberContext().doOnNext(ctx -> {
      if (!ctx.hasKey(CONTEXT_MAP)) return;

      final HttpHeaders headers = res.getHeaders();
      ctx.<Map<String, String>>get(CONTEXT_MAP).forEach(
        (key, value) -> headers.add(MDC_HEADER_PREFIX + key, value)
      );
    }).then();
  }
}

やっていることは非常にシンプルだが、汎用性が高い。
我々のF●ckシステムのように、デバッグのためのユーザ情報としてあるCookieを利用せざるを得ない場合などにも使える。

おわり

お客さんに返すレスポンスに乗っけないようにね！
この記事は自ブログの書きなぐりを加筆修正、部分抜粋して掲載しました。