Java: AutoCloseable で MDC を管理する

Java 7 では、入出力ストリームやDB接続のクローズを簡単・確実に行えるようにするために、ARM (Automatic Resource Management) が追加された。元ネタは C# の using ブロックと IDisposable インタフェースで、Java 独自の新発明というわけではないんだけど、これで今後は同じような finally ブロックを大量生産しなくてよくなった。
これを、Log4j の MDC を管理するために使ってみたらどうなるか。

今までどうしていたか

Log4j では各スレッドが Hashtable を 1 つ持っていて、任意のキー (String) に対して値をセットすることができる。設定ファイルの側で %X{キー} という記述を含めておくと、自動的に Hashtable からその値を取ってきて出力してくれる。
つまり、log4j.xml に

     <layout class="org.apache.log4j.PatternLayout">
        <param name="ConversionPattern" value="%d %5p [%X{id}] - %m%n" />
     </layout>

とか書いておいて、Java 側で

	MDC.put("id", "12345");
	LOG.info("処理中だよー! ほんとだよー!!");

とすると、

2011-10-10 03:47:37,934  INFO [12345] - 処理中だよー! ほんとだよー!!

みたいな出力が得られていた。
ここまではいいんだけど、問題は次のような場合。

スレッドプールを使っているので、1 回分の処理が終わったらスレッドに紐付いている値をクリアしたい
値を MDC に設定してから別の処理を呼び出すが、その中でだけ MDC に別の値を設定し、終わったら元に戻したい

素直に対処すると、次のようになる。

	public void run() {
		String id = createID();
		
		// 変更前の値を保存してから、新しい値を設定
		Object oldValue = MDC.get("id");
		try {
			MDC.put("id", id);
			
			LOG.info("処理中だよー! ほんとだよー!!");
			
		} finally {
			// 処理が終わったら、元の値に戻す
			if (oldValue == null) {
				MDC.remove("id");
			} else {
				MDC.put("id", oldValue);
			}
		}
	}

ユーティリティクラスを作ればもう数行は短くできるけど、いずれにしても、確実に値を元に戻すには try/finally ブロックを正確に書かなくてはならなかった。

ARM を使ってみる

まず、AutoCloseable インタフェースを実装したクラスを作り、MDC の設定と値の復元をその中に閉じ込める。

import org.apache.log4j.MDC;

public class MDCResource implements AutoCloseable {
	private String key;
	private Object oldValue;
	
	private MDCResource(String key, Object value) {
		this.key = key;
		this.oldValue = MDC.get(key);
		
		MDC.put(key, value);
	}
	
	public static MDCResource put(String key, Object value) {
		return new MDCResource(key, value);
	}
	
	public void close() {
		if (this.oldValue == null) {
			MDC.remove(this.key);
		} else {
			MDC.put(this.key, this.oldValue);
		}
	}
}

MDC を設定したい箇所では、上記のクラスを使って次のように書ける。

	public void run() {
		String id = createID();
		try (MDCResource mdc = MDCResource.put("id", id)) {
			LOG.info("処理中だよー! ほんとだよー!!");
		}
	}

これでだいぶスッキリした。try ブロックを抜けたら、MDC はブロックに入る前の状態に確実に戻っている。

似たようなことへの応用

今回は Log4j の MDC を例にしたけど、SLF4J みたいな他のロギングライブラリにも似たような仕組みがあって、同じ手法で対処できる。また、もう少し広げて「スレッドに紐付いている値をどうこうする」ということで言えば、java.lang.ThreadLocal にも応用可能だし、何かの処理の開始・終了で何かをしたい場合に、流用することもできるかもしれない。例えば「デバッグフラグが有効の場合のみ、処理にかかった時間を計測して出力する」とか。
まぁ、あんまり乱用すると意味不明になるかもしれないけど。

ちなみに……

上記の最後のコードで、try ブロックのところが

		try (MDCResource mdc = MDCResource.put("id", id)) {

となっている。わざわざローカル変数を宣言しているが、今回の場合、try ブロック内でこの変数を使うことはない。なので、

		try (MDCResource.put("id", id)) {

みたいに書ける……とさらに短くてよかったんだけど、どうやら構文上ローカル変数の宣言は必須らしい。したがって、やっぱり

		try (MDCResource mdc = MDCResource.put("id", id)) {

のように書く必要がある。無念。