blog.anqou.net
rss
author
tags

Jsonnet コンパイラを書いて Kubernetes のマニフェスト生成を速くできる？

#tech

この記事は、CYBOZU SUMMER BLOG FES ’24 (クラウド基盤本部 Stage) DAY 6 の記事です。

こんにちは。サイボウズのクラウド基盤本部の伴野です。この記事では、自分が趣味で作った Jsonnet コンパイラである JITsonnet を紹介したいと思います。私が所属している Datastore CSA チームでは Jsonnet というプログラミング言語を使って Kubernetes のマニフェスト生成を行っています。Jsonnet の公式の処理系は、入力する Jsonnet プログラムによっては実行にとても時間がかかる場合があり、それが業務中に問題になることがありました。そこで Jsonnet を実行可能バイナリにコンパイルしてから実行するような Jsonnet コンパイラを試しに作成してみたところ、意外とうまく動きました。この記事では、そもそもなぜ Jsonnet コンパイラを作ろうと思ったのか、どのような設計で作ったのか、また作った結果わかったことなどを皆さんに共有できればと思います。

過度な期待を避けるために（？）先に述べておくと、今回作成した処理系は全く production ready ではありませんし、実際の業務で使っているということもありません[1]。ただ、この記事を読んで、ある種の機能に特化したような言語処理系を作る面白さを感じて貰えると良いなと思っています。

#Jsonnet とは

Jsonnet は JSON を拡張したようなプログラミング言語で、通常の JSON のような記法を受け付けるほか、ローカル変数や関数といった仕組みが入っています。例えば次のようなプログラムを書くことができます： 

// test.jsonnet
{
    local a = 10, // ローカル変数の定義
    f(x):: x + a, // 関数 f(x) の定義

    foo: self.f(20),    // 関数 f を引数 20 で呼ぶ
                        // self はこのオブジェクト自体を指す
}

このプログラムを jsonnet コマンドに入力すると、次のような、"foo" に 30 が紐づいたオブジェクトの JSON が出力されます[2]： 

$ jsonnet test.jsonnet
{
  "foo": 30
}

Jsonnet の言語としての特徴の一つは、Jsonnet ではほとんどの関数が純粋[3]で、かつ式の評価が遅延評価で行われる点です。例えば次のような Jsonnet プログラムは、正常に終了して true と出力します： 

local foo = error "hello";
    // error "..." はエラーを出力して実行を異常終了させるが、
    // このプログラム中で変数 foo が使用されることは無いため、
    // error "..." が実行されることもない。
true

Jsonnet のもう一つの特徴は、標準ライブラリに含まれる関数の豊富さです（公式リファレンス）。標準ライブラリの関数は全て std というオブジェクトに含まれていて、 std.length のように呼び出します。例えば、Jsonnet で Kubernetes のマニフェストを生成する場合、Jsonnet の値を YAML のフォーマットに書き換えてくれる標準ライブラリの関数である std.manifestYamlDoc を使用して YAML フォーマットのマニフェストを生成します： 

std.manifestYamlDoc({ a: 1, b: 2 }, quote_keys=false)
// 以下のような YAML 形式の文字列が出力される
// （実際には JSON としてクオートされる）：
//
//   a: 1
//   b: 2

#Jsonnet 処理系とその問題点

Jsonnet の公式の処理系には google/jsonnet と google/go-jsonnet の 2 つがあります。元々の実装は google/jsonnet で、C++ で実装されています。google/go-jsonnet は後継の実装で、その名の通り Go で実装されています。どちらもメインの開発陣は同じですが、現在は後継の google/go-jsonnet の方が推奨の実装ということになっています（参考）。

以下では「Jsonnet」と J を大文字で書いた場合はプログラミング言語としての Jsonnet の方を、「jsonnet」・「go-jsonnet」と小文字で書いた場合には、それぞれ C++・Go で実装された処理系の方を指すことにします。

さて jsonnet と go-jsonnet は共にインタプリタとして実装されており、入力する Jsonnet のコードによってはとても動作が遅いことが知られています[4]。例えば次のような std.manifestYamlDoc を大量に呼び出す Jsonnet コードを jsonnet と go-jsonnet に入力して実行時間を測ると[5]、どちらの処理系でもかなり時間がかかることが分かります。先に見たように std.manifestYamlDoc は Kubernetes のマニフェスト生成によく使うため、この現象はマニフェスト生成が遅くなることに直結します： 

// bench01.jsonnet
[
  // std.manifestYamlDoc の呼び出しを 20,000 回繰り返す。
  std.manifestYamlDoc({ x: { y: { z: x } } }, quote_keys=false)
  for x in std.range(0, 20000)
]
ベンチマーク＼実装 jsonnet (C++ 実装) go-jsonnet (Go 実装)
bench01 3056 (sec.) 184 (sec.)

これほど実行が遅くなってしまう一つの原因は、Jsonnet にある標準ライブラリ関数の多くが Jsonnet のコードとして実装されている点です[6]。例に漏れず std.manifestYamlDoc も Jsonnet で実装されています。そのため、標準ライブラリ関数を呼び出すだけの上記のようなコードでも、大量の Jsonnet コードを逐次解釈する必要があり、結果として遅くなってしまいます。

この現象への対策として、go-jsonnet ではいくつかの標準ライブラリの関数が Go で実装され、処理系に組み込まれています[7]。例えば std.lstripChars という標準ライブラリ関数は組み込み関数になっていて、以下のようなコードを jsonnet と go-jsonnet で走らせると、 go-jsonnet が桁違いに速く実行できることが分かります： 

// bench02.jsonnet
std.lstripChars(std.repeat("hoge", 500), std.repeat("hoge", 10))
ベンチマーク＼実装 jsonnet (C++ 実装) go-jsonnet (Go 実装)
bench02 112 (sec.) 0.021 (sec.)

go-jsonnet ではいくつかの標準ライブラリ関数しか組み込み関数になっていません。一部のサードパーティ実装では、より多くの標準ライブラリを組み込み関数にすることで、更なる高速化を図っています。ただ、これらの実装では標準ライブラリを自前で実装し直しているという都合上、オリジナルである jsonnet や go-jsonnet の標準ライブラリの振る舞いと必ずしも一致しません。

#解決策のアイデア：コンパイルする

ある時、上で述べたような問題を解決しつつ Jsonnet コードを高速に実行するための方法として、Jsonnet コードを実行可能バイナリへコンパイルして実行するということを考えつきました。図にすると以下のような形です：

通常 jsonnet や go-jsonnet を使って Jsonnet コードを実行する場合、（＊）の矢印をたどって Jsonnet コードから直接 JSON を得ます。今回のアイデアではそうではなく、まず Jsonnet コードを実行可能バイナリへとコンパイルし（①）、その後実行する（②）ことで JSON を得ます。

これであれば、標準ライブラリを含めて全ての Jsonnet コードをそのまま用いるため、オリジナルと全く同じ振る舞いを担保できます[8]。その上、Jsonnet コードの代わりにネイティブコードを実行することによって動作が高速化することが見込めます。ただ、この構成では Jsonnet コードを実行可能バイナリへコンパイルするというオーバーヘッドが発生するため、実際にどのくらい速くなるかは作ってみないとわかりません。

#Jsonnet コンパイラ JITsonnet

こういう面白いアイデアを思いつくと、気になって夜しか眠れません。というわけで、完全に業務とは関係ない趣味の時間を使って自由研究で作ってみたのが拙作の JITsonnet です。発音はｼﾞｯﾂｫﾈｯﾄです。コードは GitHub で公開しています。

名前は JITsonnet ですが、実際には JIT（Just-In-Time）っぽいことはしておらず、 AOT（Ahead-Of-Time）コンパイラになっています。大まかな構造は以下の図のような形です：

まず、Jsonnet コードを Haskell コードへコンパイル（トランスパイル）します（①）。その後、コンパイルされた Haskell コードを GHC（Haskell のコンパイラ）を使ってコンパイルして実行バイナリを得て（②）、それを実行して最終的な結果を得ます（③）。① 以外はすでに存在するものをそのまま使うだけなので、JITsonnet の本体は ① の部分です。

もともとのアイデアは Jsonnet コードをそのまま実行バイナリへコンパイルするというものだったのですが、JITsonnet では一度 Haskell にコンパイルしています。 Jsonnet の言語仕様では、ほとんどの関数が純粋である点や評価が遅延評価である点など、Haskell に似た部分が多くあります。そのため、この方針によってコンパイラが書きやすくなります。また、GHC に搭載されている豊富な最適化の仕組みを活用することもできます。一方で、GHC の実行時間がオーバーヘッドとして載ってしまうというデメリットもあります。どのくらいのパフォーマンスが発揮できるかは後述します。

#Jsonnet の式を正しくコンパイルする

JITsonnet の本体の部分（上図の ① の部分）は OCaml で実装しています。中の構造は極めて教科書的です；まず与えられた Jsonnet コードに対して字句解析と構文解析を行い AST（抽象構文木）を得ます。その AST に対して必要な脱糖（desugar）や静的チェックなどを施したあと、ノードを再帰的に辿っていって Haskell の式に変換します。 Jsonnet の言語仕様は公式サイトにまとめられています。特に、文法が BNF として、意味論が big-step の操作的意味論としてまとめられているので、これを参考に実装しました。

実際のコンパイルの手順をより詳細に見てみます。例えば、JITsonnet の実装で足し算（e1 + e2）をコンパイルするコードは次のようになっています（コードは説明のため若干改変しています）。足し算の左辺 e1 と右辺 e2 について再帰的にコンパイルした後、その結果を足し合わせるような Haskell の式を生成します： 

(* compile_expr は Jsonnet の AST を受け取り、Haskell の AST を返す関数 *)
let rec compile_expr env (e0 : Syntax.Core.expr) : Haskell.expr =
  match e0.v with
  (* ... *)
  | Binary (e1, `Add, e2) -> (* e1 + e2 のコンパイル *)
      (* 予め用意した binaryAdd という Haskell の関数を
         e1 と e2 をコンパイルした結果で呼び出すように、つまり
            binaryAdd (e1 のコンパイル結果) (e2 コンパイル結果)
         となるようにコンパイルする。 *)
      make_call (Symbol "binaryAdd") [
        (* ... *)
        compile_expr env e1; (* e1 を再帰的にコンパイル *)
        compile_expr env e2  (* e2 を再帰的にコンパイル *)
      ]
  (* ... *)

基本的には上記のように部分式に対して再帰的にコンパイルすればそれで済むのですが、一部の Jsonnet の仕様は特別な考慮が必要です。その最たる例が self です。self は特別なキーワードで、基本的には self がスコープに入っている一番近いオブジェクト自身を指します。例えば次のように動きます： 

{
    a: 10,
    b: self.a,  // 上の a を参照し 10 になる。
}   // <-- self はこのオブジェクトを指す

単純そうに見える self の挙動ですが、ここにオブジェクトの加算が絡むと話が突然ややこしくなります。例えば次のような Jsonnet のコードを考えてみます： 

// 引用：https://github.com/google/go-jsonnet/blob/2b4d7535f540f128e38830492e509a550eb86d57/testdata/obj_local_right_level2.jsonnet
{
    local foo = self.bar,
    bar: "wrong",   // ①
    answer: foo
} + //（＊）
{
    bar: "right"    // ②
}

このコードに現れる加算（＊）では、左辺で定義された bar フィールド（①）が、右辺の同名のフィールド（②）によって上書きされます。したがって、足し算の結果得られるオブジェクトの answer の値は "right" になります。この self の挙動によって Jsonnet ではオブジェクトを使ったオブジェクト指向プログラミング（OOP）を可能にしていて、 Jsonnet の特徴の一つになっています。チュートリアルにもその使い方が書かれています。

しかし、このコードを単に足し算の左辺と右辺に分けて再帰的にコンパイルすると self が含まれている左辺のオブジェクトをコンパイルしている最中には右辺のオブジェクトについて知らないため、self を左辺のオブジェクトに結びつけるほかありません。その結果 "right" の代わりに "wrong" が answer に紐づくような誤ったオブジェクトを与えるコンパイルになってしまいます。この現象は、もとをたどると、self という名前が指す先がその式のコンパイル時には（つまり静的には）決まらないということに起因しています[9]。

JITsonnet ではこの self の挙動を正しくコンパイルするために、オブジェクトは単なる値ではなく、self を外から受け取る関数[10]だと思ってコンパイルしています。 self の値はそのオブジェクトの使用時には（つまり動的には）定まるので、実行時にその値を引数として渡し self として用いることで、正しい振る舞いが可能になります。ただ、オブジェクト使用時に毎回この関数を実行するのはオーバーヘッドがかかりすぎるため、self が単にその self が含まれるオブジェクトを指すような単純なケース（つまり足し算が前後にないようなケース）については値をキャッシュし、毎回関数が走らずに済むようにしています[11]。

#パフォーマンス

この記事を執筆している時点での最新の JITsonnet（コミット：acf8f64）で冒頭のベンチマーク（bench01.jsonnet, bench02.jsonnet）を実行すると次のような結果になりました。なお実行時間には、上図の ①・②・③ の全ての時間を含んでいます：

ベンチマーク＼実装 jsonnet (C++ 実装) go-jsonnet (Go 実装) JITsonnet
bench01 3056 (sec.) 184 (sec.) 6.84 (sec.)
bench02 112 (sec.) 0.021 (sec.) 1.43 (sec.)

jsonnet と比較すると bench01 では約 447 倍、bench02 では約 26.9 倍高速に実行することができました。また、公式推奨の実装である go-jsonnet と比べても、bench01 では約 16.4 倍高速に実行できました。一方で bench02 では約 68.1 倍遅くなっています。前に述べたように、このベンチマークでは go-jsonnet は組み込みの（つまり Go で実装された）標準ライブラリ関数を使用しています。Go でハンドメイドされた関数と比べると JITsonnet のコンパイル結果はまだまだ遅いようです。

#真面目に使う場合の課題

冒頭にも書いたとおり、 JITsonnet を書き始めた大本のモチベーションは、業務で使用している Jsonnet コードのマニフェスト生成スピードが遅いのをなんとかしたいというものでした。では JITsonnet でこれが解決したかというと、実のところ JITsonnet は PoC という立ち位置で、そのままプロダクション環境に使うことは全く想定していません。事実、社内でも使用していませんし、今後使用する予定も今のところありません。

現状の JITsonnet を真面目に使おうとした際に一番問題になるのは、 jsonnet や go-jsonnet で高速に実行できるコードでは、どうしても JITsonnet の実行時間のほうが長くなってしまうという点です。 JITsonnet では、たとえどんなに短い Jsonnet プログラムでも、1 秒程度のコンパイルのオーバーヘッドがかかってしまいます。そのため、短い Jsonnet プログラムを大量に動かすようなユースケースでは jsonnet・go-jsonnet のほうが JITsonnet よりも遥かに速く実行を終えることができるでしょう。

その他の問題点として、 jsonnet や go-jsonnet が備えているコールスタックの機能が JITsonnet では不十分にしか備わっていない点があります。例えば次のようなコードを考えてみます： 

local f() = error "msg"; // エラーを送出する
local g() = f();
local h = g();
h

このコードを jsonnet で実行すると次のようにコールスタックが表示されます： 

RUNTIME ERROR: msg
	hoge.jsonnet:1:13-24	function <f>
	hoge.jsonnet:2:13-16	function <g>
	hoge.jsonnet:3:11-14	thunk <h>
	hoge.jsonnet:4:1

一方 JITsonnet では次のように、コールスタックの一部しか表示されません。より詳細に言うと、エラーが発生した箇所から連続する関数呼び出し（function <...> の行）のみが表示され、サンク（thunk <...> の行）の実行が挟まるとそれ以降は表示されません： 

RUNTIME ERROR: msg
	hoge.jsonnet:1:13-24	function <f>
	hoge.jsonnet:2:13-16	function <g>
	hoge.jsonnet:3:11-14

これは、現在のコンパイル方法による本質的な制約だと考えています。現在の JITsonnet の構成では Jsonnet コードを一度 Haskell コードにコンパイルしてから実行可能バイナリを得ます。その際、Jsonnet の遅延評価の仕組みは Haskell の遅延評価の機構をそのまま流用してコンパイルされます。その結果、Haskell がいつある式を実行するか（サンクを実行するか）は Haskell （のランタイム）のみが知っていて、JITsonnet 側では知り得ません。そのため、jsonnet や go-jsonnet が出力するような、どのサンクの実行時にエラーが発生したかという情報は出力できないのです[12]。

#まとめ

Jsonnet 公式の処理系は、入力する Jsonnet コードによっては、動作がとても遅くなってしまいます。この挙動は、Jsonnet で Kubernetes のマニフェスト生成を行う際に問題となることがありました。そこで、Jsonnet コードを Haskell を経由して実行可能バイナリにコンパイルしてから実行することで、公式処理系と同様の振る舞いを担保しつつ高速に実行を行う JITsonnet を書いてみました。結果、公式の処理系と比較して（ベンチマークによっては）大幅な高速化を見ることができました。

インタプリタを JIT コンパイラに書き換えるのは簡単な高速化手法のように思えますが、様々な言語で見られるように、全ての入力で高速な処理を実現するためには、それぞれの言語に合わせた工夫が必要なことがよく知られています。今回の JITsonnet はそれほど手間をかけてチューニングしたというわけでもなく、実際公式処理系よりも遅くなるケースも多々あります。それでも、もともとの motivating example で十分高速化できたのは幸運でした。この記事に書くには小さすぎるコンパイルのネタも含め、色々と興味深い自由研究になりました。

最後にもう一度 JITsonnet の GitHub レポジトリへのリンクを貼っておきます：

https://github.com/ushitora-anqou/jitsonnet

興味が湧いた方はぜひ README を参考に動かしてみてください。

今日のお相手は伴野でした。明日の CYBOZU SUMMER BLOG FES ’24 クラウド基盤本部 Stage の担当は池添さんの予定です。

#注釈

  1. 業務では、自分とは別のメンバーが、まず Jsonnet で処理した結果を一度 JSON として出力し（これはそれなりに高速）、その JSON を Python スクリプトで YAML に整形するという方式で問題を解決してくれました。 ↩

  2. このように、コマンドライン引数を何も渡さない場合、Jsonnet のコードを実行すると JSON が出力されます。 Jsonnet で Kubernetes のマニフェスト生成を行う場合は JSON を出力されても困るので一工夫必要です。例えば私のチームでは jsonnet コマンドの -m や -S オプションなどを利用し、Jsonnet コードを実行した結果得られたマニフェストの文字列を複数ファイルに分けて出力するようにしています。 ↩

  3. 一応、純粋でない関数も存在します。例えば、エラーを出力して実行を異常終了させる error や、デバッグ出力を行うための std.trace などは純粋ではありません。 ↩

  4. 詳細は https://www.databricks.com/blog/2018/10/12/writing-a-faster-jsonnet-compiler.html にて調査されています。 ↩

  5. このベンチマークを含め、この記事に含まれる全てのベンチマークは、以下の環境で計測しています：

    • CPU: Intel Core i7-8700
    • RAM: 32 GiB
    • OS: Ubuntu 24.04 LTS
    • jsonnet: v0.20.0
    • go-jsonnet: v0.20.0
    ↩

  6. 標準ライブラリの実装は jsonnet の std.jsonnet ファイルに存在します。 ↩

  7. 面白い話として、一部の標準ライブラリの関数が組み込み関数として実装されているという go-jsonnet の事情を悪用すると、Jsonnet コード中で、いま jsonnet と go-jsonnet のどちらで実行されているかを判定できます。例えば以下のような式を実行すると、jsonnet では "jsonnet" と表示され、 go-jsonnet では "go-jsonnet" と表示されます。 

    if (std { objectFields(x):: [] }).equals({}, {x: 1})
then "jsonnet"
else "go-jsonnet"

    これは jsonnet では std.equals は組み込み関数ではなく std.objectFields などを呼び出す Jsonnet コードとして実装されているのに対し、go-jsonnet では std.equals が組み込み関数になっているので std.objectFields が呼ばれないということを利用しています。 ↩

  8. もちろん、コンパイラが正しく実装されていることを仮定しています。 ↩

  9. この仕組みは、Jsonnet の通常の変数であれば変数の指す先は静的に決まる（レキシカルスコープである）ことと対照的です。 ↩

  10. より正確には、オブジェクトに含まれる各フィールドが self を外から受け取ります。また super にも同様の問題があるため super も追加で受け取っています。このようなエンコードは OOP ができる言語（C++ など）でよく見られるものですが、 Jsonnet では継承関係が実行時にしか決まらないという違いがあります。そのため、self 経由で呼ばれる関数や変数の値が動的にしか決まらず、コンパイルの最適化が難しくなっているように感じます。 ↩

  11. go-jsonnet でもこのあたりの実装は工夫されていて、オブジェクトの値を表す構造体にそのオブジェクトの中身を表す AST を含め、必要に応じてその AST を評価しなおすことでこの self の挙動を実現しているように見えます（参考）。 ↩

  12. Haskell には詳しくないので、実は抜け穴的なやり方があるかもしれませんが、今のところ見つけられていません。 ↩