[JNI] nativeからJavaのStreamの消費とパフォーマンス

現在従事しているプロジェクトのプロダクトはJava製ですが、一部の実装にnative (C++)実装を含んでおり、JNIでJava-nativeコミュニケーションを実装しています。
今回は少し前にnativeに計算させるデータの渡し方を修正したので、改めてその実装のパフォーマンスを計ってみようという記事です。

## array実装版
とりあえず修正前の実装と同様のサンプルコードです。
(計測に使った最終的なコードは[こちら](https://github.com/kikuchi-m/ac2021-jni-sample)。)

nativeに渡すデータは、複数の`long`のセットのコレクションです。array版ではすべてをいったん`long[][]`にして一度に渡します。

```java:Application.java
package jp.co.atware.ac2021.jni.app;

...

public class Application {

    public static void main(String... args) {
        System.out.println("JNI performance by array");
        var nativeApi = new NativeApi();

        // prepare
        var size = Integer.parseInt(args[0]);

        // perform
        var start = LocalTime.now();

        var dataStream = Data.generateDataStream(size);
        var data = new long[size][];
        var dataList = dataStream.collect(toList());
        for (var i = 0; i < size; i++) {
            data[i] = dataList.get(i).toPrimitiveArray();
        }
        var res = nativeApi.calcByArray(data);

        var end = LocalTime.now();

        // result
        var dur = Duration.between(start, end).toMillis() / 1000.0;
        System.out.printf("%sS (start: %s / end: %s)\n", dur, start, end);
        System.out.printf("calculated value: %d\n", res);
    }
}
```

```java:Data.java
package jp.co.atware.ac2021.jni.lib;

...

public class Data {

    private final List<Long> values;

    private Data(List<Long> values) {
        this.values = values;
    }

    public long[] toPrimitiveArray() {
        var raw = new long[values.size()];
        IntStream.range(0, values.size()).forEach(i -> raw[i] = values.get(i));
        return raw;
    }

    public static Stream<Data> generateDataStream(int size) {
        return IntStream.range(0, size)
            .map(i -> i % 10 + 1)
            .mapToObj(i -> IntStream.range(0, i).mapToLong(v -> i).boxed().collect(toList()))
            .map(Data::new);
    }
}
```

```java:NativeApi.java
package jp.co.atware.ac2021.jni.lib;

public class NativeApi {

    public NativeApi() {
        System.loadLibrary("jni-performance-sample");
    }

    public native long calcByArray(long[][] data);
}
```

native側は単純にすべてのvalueを足して返すだけにします。(headerは省略。)

```cpp
#include "jp_co_atware_ac2021_jni_lib_NativeApi.h"

#include <jni.h>

jlong Java_jp_co_atware_ac2021_jni_lib_NativeApi_calcByArray(JNIEnv* env, jobject, jobjectArray data) {

    jlong result = 0;
    auto dataSize = env->GetArrayLength(data);

    for (jsize i = 0; i < dataSize; ++i) {
        auto d = static_cast<jlongArray>(env->GetObjectArrayElement(data, i));
        auto ds = env->GetArrayLength(d);
        auto dPtr = env->GetLongArrayElements(d, JNI_FALSE);
        for (jsize di = 0; di < ds; ++di) {
            result += dPtr[di];
        }
        env->ReleaseLongArrayElements(d, dPtr, JNI_ABORT);
    }
    return result;
}
```

これで計測します。5回実行の平均は 0.4064 secs でした。

* Data (collection) size = 1,000,000
* JVM heap = 256MB

まぁ、こんなもんでしょうか。

## メモリ……が……足りない……
1M(1,000,000)のcollectionはなんとかなりましたが、実際のプロダクトではこの数倍~数十倍のデータを扱います。実際data sizeが2Mになるとlist -> arrayの部分でout of memoryで落ちます。

```java:Application.java
for (var i = 0; i < size; i++) {
    data[i] = dataList.get(i).toPrimitiveArray();
}
```

十数年前のPCでもないしメモリはあるので、とりあえずheapを増やしておきます。(デプロイして使っている環境では1,024MBで運用しています。)

* Data size = 2,000,000
* JVM heap = 1,024MB

同じように5回実行し、平均は 0.6266 secs。だんだんdata sizeを増やしていきます。

data size | avg. |
--- | --- |
3M  | 0.9039 |
5M  | 1.4679 |
10M | NaN |

10MでまたOOMになりました。このくらいのdata sizeもまだ扱うレベルなので困ります。
(実際の問題はちょっと違っていて、OOMで落ちず、list -> arrayの箇所で処理が止まったままになり、native callはされず……という状態でした。)

## Streamのままnativeで消費する
arrayを使うのはメモリ効率が悪いのと、また一度`List`にする必要も実際にはなかったので、`Stream`のままnativeでaccumlateするように実装を変更しました。`Stream`のままと言いましたが、単純に`Iterator`にして回すだけですが、途中で長大な`List`を作らずに回すタイミングで遅延評価(= create `Data` instance)できるのでメモリ効率はよくなります。Stream APIができてこういうことがやりやすくなりましたね。

では実装を追加、変更します。(抜粋)

```java:Application.java
    public static void main(String... args) {
        ...
        var dataStream = Data.generateDataStream(size);
        var res = nativeApi.calcByStream(StreamIterator.of(dataStream));
        ...
    }
```

```java:StreamIterator.java
package jp.co.atware.ac2021.jni.lib;

...

public class StreamIterator {

    private final Iterator<Data> iter;

    private StreamIterator(Stream<Data> dataStream) {
        this.iter = dataStream.iterator();
    }

    public static StreamIterator of(Stream<Data> dataStream) {
        return new StreamIterator(dataStream);
    }

    public boolean hasNext() {
        return iter.hasNext();
    }

    public long[] next() {
        return iter.next().toPrimitiveArray();
    }
}
```

```java:NativeApi.java
    public native long calcByStream(StreamIterator iter);
```

```cpp
jlong Java_jp_co_atware_ac2021_jni_lib_NativeApi_calcByStream(JNIEnv* env, jobject, jobject iter) {
    auto streamIterator = env->FindClass("jp/co/atware/ac2021/jni/lib/StreamIterator");
    auto hasNext = env->GetMethodID(streamIterator, "hasNext", "()Z");
    auto next = env->GetMethodID(streamIterator, "next", "()[J");
    jlong result = 0;

    while (env->CallBooleanMethod(iter, hasNext)) {
        auto d = static_cast<jlongArray>(env->CallObjectMethod(iter, next));
        auto ds = env->GetArrayLength(d);
        auto dPtr = env->GetLongArrayElements(d, JNI_FALSE);
        for (jsize di = 0; di < ds; ++di) {
            result += dPtr[di];
        }
        env->ReleaseLongArrayElements(d, dPtr, JNI_ABORT);
        env->DeleteLocalRef(d);
    }
    return result;
}
```

オーバヘッドが大きそうですがこれで計測します。先の結果を併せて比較します。

* JVM heap = 256MB

Data Size | avg. - array | avg. - stream |
--- | --- | --- |
1M | 0.4064 | 0.4923 |
2M |    NaN | 0.8762 |

* JVM heap = 1,024MB

Data Size | avg. - array | avg. - stream |
--- | --- | --- |
1M  | 0.3475 | 0.5843 |
2M  | 0.6266 | 0.9702 |
3M  | 0.9039 | 1.3964 |
5M  | 1.4679 | 2.1834 |
10M |    NaN | 4.1892 |

遅いですね。プロダクトでは比較計測はせず(できず)、この実装で処理が進んでちゃんと結果を得られるようになったでここまでにしました。
今回記事を書くモチベーションとして、Stream版の実装のパフォーマンスを見たかったので、その結果を得られました。やはりJava-nativeのコミュニケーションが多いほどパフォーマンスが落ちますね。まぁ順当な結果と言えるでしょう。

(数百万回native -> JVMのmethod callをするので、正直劇的に遅くなるかな？と思っていましたが、そこまでじゃないという印象。)

## Java-nativeのコミュニケーションを減らす
`Stream`のままnativeとのコミュニケーションを減らすために、nativeで必要になる値をchunkにしつつ`Stream`のまま回します。

```java:Application.java
    public static void main(String... args) {
        ...
        // prepare
        var size = Integer.parseInt(args[0]);
        var chunkSize = Integer.parseInt(args[1]);

        // perform
        var start = LocalTime.now();

        var dataStream = Data.generateDataStream(size);
        var res = nativeApi.calcByStreamChunk(StreamIterator.of(dataStream, chunkSize));
        ...
    }
```

```java:StreamIterator.java
    private final Iterator<Data> iter;
    private final int chunkSize;

    private StreamIterator(Stream<Data> dataStream, int chunkSize) {
        this.iter = dataStream.iterator();
        this.chunkSize = chunkSize;
    }
    public static StreamIterator of(Stream<Data> dataStream, int chunkSize) {
        return new StreamIterator(dataStream, chunkSize);
    }

    public long[][] nextChunk() {
        var chunk = new ArrayList<long[]>(chunkSize);
        while (iter.hasNext() && chunk.size() < chunkSize) {
            chunk.add(iter.next().toPrimitiveArray());
        }
        return chunk.stream().toArray(long[][]::new);
    }
```

```cpp
jlong Java_jp_co_atware_ac2021_jni_lib_NativeApi_calcByStreamChunk(JNIEnv* env, jobject, jobject iter) {
    auto streamIterator = env->FindClass("jp/co/atware/ac2021/jni/lib/StreamIterator");
    auto hasNext = env->GetMethodID(streamIterator, "hasNext", "()Z");
    auto nextChunk = env->GetMethodID(streamIterator, "nextChunk", "()[[J");
    jlong result = 0;

    while (env->CallBooleanMethod(iter, hasNext)) {
        auto dc = static_cast<jobjectArray>(env->CallObjectMethod(iter, nextChunk));
        auto dcs = env->GetArrayLength(dc);
        for (jsize dci = 0; dci < dcs; ++dci) {
            auto d = static_cast<jlongArray>(env->GetObjectArrayElement(dc, dci));
            auto ds = env->GetArrayLength(d);
            auto dPtr = env->GetLongArrayElements(d, JNI_FALSE);
            for (jsize di = 0; di < ds; ++di) {
                result += dPtr[di];
            }
            env->ReleaseLongArrayElements(d, dPtr, JNI_ABORT);
            env->DeleteLocalRef(d);
        }
        env->DeleteLocalRef(dc);
    }
    return result;
}
```

chunk sizeを変えながら計測した結果はこちら。これも同様に5回実行の平均。

* JVM heap = 1,024MB

data size | avg. - chunk 1K | avg. - chunk 10K | avg. - chunk 50K | chunk 100K |
---  | --- | --- | --- | --- |
1M  | 0.3942 | 0.4145 | 0.4402 | 0.4402 |
2M  | 0.6294 | 0.6898 | 0.7276 | 0.7352 |
3M  | 0.8426 | 0.9282 | 1.0214 | 1.0116 |
5M  | 1.2692 | 1.4232 | 1.5586 | 1.5588 |
10M | 2.3632 | 2.7088 | 2.9113 | 2.9528 | 

chunkが小さい方がパフォーマンスがいい感じですね。ちょっと意外。
これまでの結果と併せて3つの実装の比較をしてみます。

Data Size | avg. - array | avg. - stream | avg. - chunk 1K |
--- | --- | --- | -- |
1M  | 0.3475 | 0.5843 | 0.3942 |
2M  | 0.6266 | 0.9702 | 0.6294 |
3M  | 0.9039 | 1.3964 | 0.8426 |
5M  | 1.4679 | 2.1834 | 1.2692 |
10M |    NaN | 4.1892 | 2.3632 |

data size 1Mではarray ver.がいいですが、2Mではほぼ同等、3M~ではchunk ver.に軍配が上がりました。

## おわりに
今回はnative側の実装を`Stream`を回す部分だけを実装し、native -> JVMのAPI callの部分ではexception handlingはしていません。(`env->ExceptionCheck()`、`env->ExceptionOccurred()`など。) この辺りをちゃんと実装するとまた少し結果は変わるでしょう。
この結果から今のプロダクトでの実装がまだパフォーマンスを改善できそうなことがわかったので、あとで修正チケットをあげようと思います。

---

### Note
* environment
    * Ubuntu 20.04
    * Intel Core i9-9900K 3.6GHz
    * memory 64GB
    * OpenJDK 11.0.2
* source https://github.com/kikuchi-m/ac2021-jni-sample
* [Java Native Interface Specification Contents (Java 11)](https://docs.oracle.com/en/java/javase/11/docs/specs/jni/index.html)