どうも、趣味でOpenJDKのコミッタをやってます。JVMの専門家ではないです。
今回はJITコンパイルによる暖気が十分に行われてから処理を受けられるようにする方法を紹介します。
今回の実装は Oracle の有償機能から OpenJDK へ寄贈され OpenJDK 11 に追加された JDK Flight Recorder（JFR）と OpenJDK 14 に追加された JFR Event Streaming を使用します。
JITコンパイルイベント
OpenJDK では、JITコンパイラがコンパイルした情報を内部的にイベントとして記録しています。今回はこのイベントを使用していきます。このイベントのデータ形式は以下になります。
@Name("jdk.Compilation") @Category({"Java Virtual Machine", "Compiler"}) @Label("Compilation") class Compilation extends jdk.jfr.Event { @Label("Start Time") @Timestamp("TICKS") long startTime; @Label("Duration") @Timespan("TICKS") long duration; @Label("Event Thread") @Description("Thread in which event was committed in") Thread eventThread; @Unsigned @CompileId @Label("Compilation Identifier") int compileId; @Label("Compiler") String compiler; @Label("Method") Method method; @Unsigned @Label("Compilation Level") short compileLevel; @Label("Succeeded") boolean succeded; @Label("On Stack Replacement") boolean isOsr; @Unsigned @DataAmount("BYTES") @Label("Compiled Code Size") long codeSize; @Unsigned @DataAmount("BYTES") @Label("Inlined Code Size") long inlinedBytes; }
イベントの名前はjdk.Compilationです。このイベントは、コンパイルされた時間、コンパイルにかかった期間、コンパイルされたメソッド、コンパイルのレベルなどを記録します。
このイベントは実際に以下のような内容が記録されます。
jdk.Compilation { startTime = 14:52:33.762 duration = 1.029 ms compileId = 4710 compiler = "c1" method = sun.rmi.transport.ConnectionInputStream.done(Connection) compileLevel = 3 succeded = true isOsr = false codeSize = 12.1 kB inlinedBytes = 266 bytes eventThread = "C1 CompilerThread0" (javaThreadId = 10) }
JITのウォーミングアップの監視の実装
JIT コンパイルは Java プロセスの起動時から開始されます。そのため、JIT コンパイルをより正確に監視するには Java アプリケーションが起動する前から監視する必要があります。監視は Java の Pre Main で開始して、ウォーミングアップの状態は MBean で持ちます。これらの状態を SpringBoot の中から参照していきます。
注：最初は SpringBoot の中で全てやっていましたが、誤解が広まると嫌なので書き換えました。
JIT コンパイルによるウォーミングアップの状態は以下の MXBean に持たせます。
public interface JitWarmUpMXBean { boolean isWarmedUp(); }
実装は以下になります。
public class JitWarmUpMXBeanImpl implements JitWarmUpMXBean{ private volatile boolean isWarmedUp = false; public void warmedUp() { this.isWarmedUp = true; } @Override public boolean isWarmedUp() { return this.isWarmedUp; } }
状態は isWarmedUp フィールドに持ちます。値がfalseならウォーミングアップが十分ではなく、trueなら期待するレベルでウォーミングアップが完了しています。
次に、暖気による Readiness Probe を確認できるようにしましょう。まずは、暖気が終わってなければDOWNを、暖気が終わっていればUPを返す実装を作ります。これはSpringのお作法に則り、HealthIndicatorを実装します。HealthIndicatorは以下のとおりです。
import org.springframework.boot.actuate.health.Health; import org.springframework.boot.actuate.health.HealthIndicator; import org.springframework.stereotype.Component; import javax.management.*; import java.lang.management.ManagementFactory; @Component public class JitWarmUpHealthIndicator implements HealthIndicator { private JitWarmUpMXBean jitWarmUpMXBean; public JitWarmUpHealthIndicator() { try { MBeanServer mbs = ManagementFactory.getPlatformMBeanServer(); ObjectName mxbeanName = new ObjectName("com.example:type=JitWarmUp"); this.jitWarmUpMXBean = JMX.newMBeanProxy(mbs, mxbeanName, JitWarmUpMXBean.class); } catch (MalformedObjectNameException e) { e.printStackTrace(); } } @Override public Health health() { return jitWarmUpMXBean != null && jitWarmUpMXBean.isWarmedUp() ? Health.up().build() : Health.down().build(); } }
このクラスは暖気が終わったかどうかを保持するJitWarmUpMXBeanへのプロキシを持ちます。このMBeanはisWarmedUpメソッドを持ち、このメソッドは起動直後は暖気が終わっていないのでfalseを返しますが、暖気が終わるとtrueを返します。health()メソッドは、JitWarmUpMXBeanクラスのisWarmedUpメソッドを見て、暖気が終わっていればUPを、終わっていなければDOWNを返します。
次は、JITコンパイルされたメソッドの数を数える処理です。ここが本質だと思います。今回は、JITコンパイラに一定個数のメソッドがサーバコンパイル（Level 4）されたら暖気完了とします。
注意：先ほどのコンパイルイベントのデータ構造にメソッド名などもありました。それを使えば特定のメソッドがコンパイルされたら暖気を完了するという実装もできます。ですが、サーバコンパイルされるかどうかはJITの気分次第であり、インライン化されたりサーバコンパイルされないと言った可能性もあるのでお奨めしません。実際、何度か検証しましたが、狙ったメソッドがサーバコンパイルされないことが多々ありました。
JITコンパイラの監視は、Pre Main で実行します。JITコンパイルの情報は JDK Flight Recorder に記録され、イベントの数の集計処理はその記録から JFR Event Streaming で処理します。
public class JitWarmUpAgent { private static final int threshold = 2000; private static final int expectCompileLevel = 4; private static volatile boolean isClosed = false; private final static void start() { try { MBeanServer mbs = ManagementFactory.getPlatformMBeanServer(); ObjectName mxbeanName = new ObjectName("com.example:type=JitWarmUp"); JitWarmUpMXBeanImpl mxbean = new JitWarmUpMXBeanImpl(); mbs.registerMBean(mxbean, mxbeanName); RecordingStream es = new RecordingStream(); es.enable("jdk.Compilation").withThreshold(Duration.ZERO); AtomicInteger compiledMethodNum = new AtomicInteger(0); Consumer<RecordedEvent> compilation = e -> { if (e.getShort("compileLevel") >= expectCompileLevel) { int methodNum = compiledMethodNum.incrementAndGet(); if (methodNum >= threshold) { if (!isClosed) { synchronized (es) { if (!isClosed) { System.out.println("閾値を超える数の " + methodNum + " メソッドがコンパイルされました"); mxbean.warmedUp(); es.close(); isClosed = true; } } } } } }; es.onEvent("jdk.Compilation", compilation); es.startAsync(); Runtime.getRuntime().addShutdownHook(new Thread("JitWarmUpShutdownHook") { public void run() { if (!isClosed) { es.close(); } } }); } catch (MalformedObjectNameException e) { e.printStackTrace(); } catch (NotCompliantMBeanException e) { e.printStackTrace(); } catch (InstanceAlreadyExistsException e) { e.printStackTrace(); } catch (MBeanRegistrationException e) { e.printStackTrace(); } } // Used when loading agent from command line public static void premain(String agentArgs, Instrumentation inst) { start(); } // Used when loading agent during runtime. public static void agentmain(String agentArgs, Instrumentation inst) { start(); } }
このクラスにはパラメータとして、何個のメソッドをコンパイルすれば暖気終了と見なすのかを表す閾値であるthresholdフィールドと、期待するコンパイルレベルを示すexpectCompileLevelを持ちます。
監視開始の処理はstart()メソッドで、監視の終了はウォーミングアップが済んだときとシャットダウンフックで行います。
起動時の処理ではまず、JFR Event Streaming の記録を開始します。この時、コンパイルのイベントを有効にします。通常、コンパイルのイベントはコンパイルに 100ms 以上かかったものだけが記録されます。しかし、それではコンパイル時間が長いイベントだけが記録されるだけであり、コンパイルされた全てのメソッド数は記録できません。そのため、全てのイベントを記録できるように0 msを指定します。これらの設定を指定して、記録を開始します。
コンパイルされたメソッド数を保持する変数としてAtomicIntegerクラスのcompiledMethodNumを作成します。
次に、コンパイルのイベントで処理されるコンシューマを作成します。この中では、来たコンパイルのイベントが期待するコンパイルレベルかを確認します。期待するコンパイルレベルだった場合には、コンパイルされたメソッドとして記録しています。同じメソッドが何度もコンパイルされる可能性もありますが、今回の実装では無視しています。メソッドの名前などを使えばより正しく記録できるようになるでしょう。その後、コンパイルされたメソッド数が期待する閾値を超えると標準出力へログを出力して、HealthIndicatorを暖気済に変更して、処理を停止させます。
最後にこの処理をコンパイルのイベントで処理されるように設定して、非同期で処理を開始します。
これらの処理によって、起動後に暖気用のリクエストを受けてJITコンパイルされたメソッドが増えると、いつしかその閾値を超えてJitWarmUpMXBean の状態が true になります。そうすると、HealthIndicatorのステータスがUPを返すようになり、そのアプリケーションが本番のリクエストを受けられるようになります。
実際に確認してみましょう。HealthIndicatorの様子が分かるようにapplication.propertiesに以下を追加しています。
management.endpoint.health.show-details=always
まずは起動直後です。jitWarmUpのステータスがDOWNなため、全体のステータスもDOWNになっています。
次に、暖気処理後です。jitWarmUpのステータスがUPになり、全体のステータスもUPになっています。
いかがでしたでしょうか。
この方法を使えば、JITコンパイルが十分に行われてからJavaプロセスが処理を受けられるようになりました。もし、性能に関する要件が厳しく、JITコンパイルが十分に行わないといけない場合にはこの様な方法を使用してみて下さい。そういった要件が無い場合は、使用する必要はありません。
