密集計算場景下的 JNI 實戰

vivo網際網路技術發表於2022-08-23

作者：vivo 網際網路伺服器團隊-  Wei Qianzi、Li Haoxuan

在 Java 發展歷程中，JNI 一直都是一個不可或缺的角色，但是在實際的專案開發中，JNI 這項技術應用的很少。在筆者經過艱難的踩坑之後，終於將 JNI 運用到了專案實戰，本文筆者將簡單介紹 JNI 技術，並介紹簡單的原理和效能分析。透過分享我們的實踐過程，帶各位讀者體驗 JNI 技術的應用。

一、 背景

計算密集型場景中，Java 語言需要花費較多時間最佳化 GC 帶來的額外開銷。並且在一些底層指令最佳化方面，C++ 這種“親核性”的語言有著較好的優勢和大量的業界實踐經驗。那麼作為一個多年的 Java 程式設計師，能否在 Java 服務上面執行 C++ 程式碼呢？答案是肯定的。

JNI (Java Native Interface) 技術正是應對該場景而提出的解決方案。雖然 JNI 技術讓我們能夠進行深度的效能最佳化，其較為繁瑣的開發方式也不免讓新人感到頭疼。本文透過 step by step 的方式介紹如何完成 JNI 的開發，以及我們最佳化的效果和思考。

開始正文前我們可以思考三個問題：

  1. 為什麼選擇使用 JNI 技術？

  2. 如何在 Maven 專案中應用 JNI 技術？

  3. JNI 真的好用嗎？

二、關於 JNI：為什麼會選擇它？

2.1 JNI 基本概念

JNI 的全稱叫做 Java Native Interface ，翻譯過來就是 Java 本地介面。愛看 JDK 原始碼的小夥伴會發現，JDK 中有些方法宣告是帶有 native 修飾符的，並且找不到具體實現，其實是在非 Java 語言上，這就是 JNI 技術的體現。

早在 JDK1.0 版本就已經有了 JNI，官方給 JNI 的定義是：

Java Native Interface (JNI) is a standard programming interface for writing Java native methods and embedding the Java virtual machine into native applications. The primary goal is binary compatibility of native method libraries across all Java virtual machine implementations on a given platform.

JNI 是一種標準的程式介面，用於編寫 Java 本地方法，並且將 JVM 嵌入 Native 應用程式中。是為了給跨平臺上的 JVM 實現本地方法庫進行二進位制相容。

JNI 最初是為了保證跨平臺的相容性，而設計出來的一套介面協議。並且由於 Java 誕生很早，所以 JNI 技術絕大部分情況下呼叫的是 C/C++ 和系統的 lib 庫，對其他語言的支援比較侷限。隨著時間的發展，JNI 也逐漸被開發者所關注，比如 Android 的 NDK，Google 的 JNA，都是對 JNI 的擴充套件，讓這項技術能夠更加輕鬆的被開發者所使用。

我們可以看一下在 JVM 中 JNI 相關的模組，如圖 1：

圖片

圖1 - JVM 記憶體和引擎執行關係

在 JVM 的記憶體區域，Native Interface 是一個重要的環節，連線著執行引擎和執行時資料區。本地介面 (JNI) 的方法在本地方法棧中管理 native 方法，在 Execution Engine 執行時載入本地方法庫。

JNI 就像是打破了 JVM 的束縛，擁有著和 JVM 同樣的能力，可以直接使用處理器中的暫存器，不僅可以直接使用處理器中的暫存器，還可以直接找作業系統申請任意大小的記憶體，甚至能夠訪問到 JVM 虛擬機器執行時的資料，比如搞點堆記憶體溢位什麼的：）

2.2 JNI 的功能

JNI 擁有著強大的功能，那它能做哪些事呢？官方文件給出了參考答案。

  1. 標準 Java 類庫不支援應用程式所需的平臺相關特性。

  2. 您已經有一個用另一種語言編寫的庫，並希望透過 JNI 使其可供 Java 程式碼訪問。

  3. 您想用較低階別的語言（例如彙編）實現一小部分耗時短的程式碼。

當然還有一些擴充，比如：

  1. 不希望所寫的 Java 程式碼被反編譯；

  2. 需要使用系統或已有的 lib 庫；

  3. 期望使用更快速的語言去處理大量的計算；

  4. 對影像或本地檔案操作頻繁；

  5. 呼叫系統驅動的介面。

或許還有別的場景，可以使用到 JNI，可以看到 JNI 技術有著非常好的應用潛力。

三、JNI 實戰：探究踩坑的全過程

我們的業務中存在一個計算密集型場景，需要從本地載入資料檔案進行模型推理。專案組在 Java 版本進行了幾輪最佳化後發現沒有什麼大的進展，主要表現為推理耗時較長，並且載入模型時存在效能抖動。經過調研，如果想進一步提高計算和載入檔案的速度，可以使用 JNI 技術去編寫一個 C++ 的 lib 庫，由 Java native 方法進行呼叫，預計會有一定的提升。

然而專案組目前也沒有 JNI 的實踐經驗，最終效能是否能有提升，還是要打個問號。本著初生牛犢不怕虎的精神，我鼓起勇氣主動認領了這個最佳化任務。下面就分享一下我實踐 JNI 的過程和遇到的問題，給大家拋磚引玉。

3.1 場景準備

實戰就不從 Hello world 開始了，我們直接敲定場景，思考該讓 C++ 實現哪部分邏輯。

場景如下：

圖片

圖2 實戰場景

在計算服務中，我們將離線計算資料轉換成 map 結構，輸入一組 key 在 map 中查詢並對 value 應用演算法公式求值。透過分析 JVM 堆疊資訊和火焰圖 (flame graph)，發現效能瓶頸主要在大量的邏輯迴歸運算和 GC 上面，由於快取了量級很大的 Map 結構，導致佔用 heap 記憶體很大，因此 GC Mark-and-Sweep 耗時很長，所以我們決定將載入檔案和邏輯迴歸運算兩個方法改造為 native 方法。

程式碼如下：

/**
 * 載入檔案
 * @param path 檔案本地路徑
 * @return C++ 建立的類物件的指標地址
 */
public static native long loadModel(String path);
 
/**
 * 釋放 C++ 相關類物件
 * @param ptr  C++ 建立的類物件的指標地址
 */
public static native void close(long ptr);
 
/**
 * 執行計算
 * @param ptr C++ 建立的類物件的指標地址
 * @param keys 輸入的列表
 * @return 輸出的計算結果
 */
public static native float compute(long ptr, long[] keys);

那麼，我們為什麼要傳遞指標呢，並且設計了一個 close 方法呢？

  1. 便於相容現有實現的考慮：雖然整個計算過程都在 C++ 執行時中進行，但物件的生命週期管理是在 Java 中實現的，所以我們選擇回傳載入並初始化後的模型物件指標，之後每次求值時僅傳遞該指標即可；

  2. 記憶體正確釋放的考慮：利用 Java 自身的 GC 和模型管理器程式碼機制，在模型解除安裝時顯式呼叫 close 方法釋放 C++ 執行時管理的記憶體，防止出現記憶體洩漏。

當然，這個建議只適用於需要 lib 執行時將部分資料快取在記憶體中的場景，只使用 native 方法進行計算，無需考慮這種情況。

3.2 環境搭建

下面簡單介紹一下我們所使用的環境和專案結構，這部分介紹的不是很多，如果有疑問可以參考文末的參考資料或者在網上進行查閱。

我們使用的是簡單的 maven 專案，使用 Docker 的 ubuntu-20.04 容器進行編譯和部署，需要在容器中安裝 GCC，Bazel，Maven，openJDK-8 等。如果是在 Windows 下進行開發，也可以安裝相應的工具並編譯成 .dll 檔案，效果是一樣的。

我們建立好 maven 專案的目錄，如下：

/src # 主目錄
-/main
--/cpp  # c++ 倉庫目錄
---export_jni.h  # java 匯出的檔案
---computer.cc  # 具體的 C++ 程式碼
---/third_party  # 三方庫
---WORKSPACE  # bazel 根目錄
---BUILD  # bazel 構建檔案
--/java  # java 倉庫目錄
---/com
----/vivo
-----/demo
------/model
-------ModelComputer.java  # java 程式碼
--/resources  # 存放 lib 的資源目錄
-/test
--/java
----ModelComputerTest.java  # 測試類
pom.xml  # maven pom

3.3 實戰過程

都已經準備好了，那麼就直入正題：

package com.vivo.demo.model;
import java.io.*;
 
public class ModelComputer implements Closeable {
    static {
        // 載入 lib 庫
        loadPath("export_jni_lib");
    }
 
    /**
     * C++ 類物件地址
     */
    private Long ptr;
 
    public ModelComputer(String path) {
        // 建構函式，呼叫 C++ 的載入
        ptr = loadModel(path);
    }
 
    /**
     * 載入 lib 檔案
     *
     * @param name lib名
     */
    public static void loadPath(String name) {
        String path = System.getProperty("user.dir") + "\\src\\main\\resources\\";
        path += name;
        String osName = System.getProperty("os.name").toLowerCase();
        if (osName.contains("linux")) {
            path += ".so";
        } else if (osName.contains("windows")) {
            path += ".dll";
        }
        // 如果存在本檔案，直接載入，並返回
        File file = new File(path);
        if (file.exists() && file.isFile()) {
            System.load(path);
            return;
        }
        String fileName = path.substring(path.lastIndexOf('/') + 1);
        String prefix = fileName.substring(0, fileName.lastIndexOf(".") - 1);
        String suffix = fileName.substring(fileName.lastIndexOf("."));
 
        // 建立臨時檔案，注意刪除
        try {
            File tmp = File.createTempFile(prefix, suffix);
            tmp.deleteOnExit();
 
            byte[] buff = new byte[1024];
            int len;
            // 從jar中讀取檔案流
            try (InputStream in = ModelComputer.class.getResourceAsStream(path);
                    OutputStream out = new FileOutputStream(tmp)) {
                while ((len = in.read(buff)) != -1) {
                    out.write(buff, 0, len);
                }
            }
            // 載入庫檔案
            System.load(tmp.getAbsolutePath());
        } catch (Exception e) {
            throw new RuntimeException();
        }
    }
 
    // native 方法
    public static native long loadModel(String path);
    public static native void close(long ptr);
    public static native float compute(long ptr, long[] keys);
 
    @Override
    public void close() {
        Long tmp = ptr;
        ptr = null;
        // 關閉 C++ 物件
        close(tmp);
    }
 
    /**
     * 計算
     * @param keys 輸入的列表
     * @return 輸出的結果
     */
    public float compute(long[] keys) {
        return compute(ptr, keys);
    }
}

  • 踩坑1：啟動時報 java.lang.UnsatisfiedLinkError 異常

這是因為 lib 檔案在壓縮包中，而載入 lib 的函式尋找的是系統路徑下的檔案，透過 InputStream 和 File 操作從壓縮包中讀取該檔案到臨時資料夾，獲取其路徑，再進行載入就可以了。上文中 getPath 方法作為解決辦法的示例可以參考：System.load() 函式輸入的路徑必須是全路徑下的檔名，也可以使用 System.loadLibrary() 載入 java.library.path 下的lib庫，不需要 lib 檔案的字尾。

儲存上文的 Java 程式碼，透過 Javah 指令可以生成對應的 C++ 標頭檔案，前文目錄結構中的 export_jni.h 就是透過該指令生成的。

javah -jni -encoding utf-8 -classpath com.vivo.demo.model.ModelComputer -o ../cpp/extern_jni.h
# -classpath 表示所在的package
# -d 表示輸出的檔名

開啟可以看到生成出來的檔案如下：

#include <jni.h>  // 引入的標頭檔案， 該標頭檔案在 $JAVA_HOME/include 下，隨Java版本變化而改變
#ifndef _Included_com_vivo_demo_model_ModelComputer // 宏定義 格式 _Included_包名_類名
#define _Included_com_vivo_demo_model_ModelComputer
#ifdef __cplusplus
extern "C" {  // 保證函式、變數、列舉等在所有的原始檔中保持一致，這裡應用於匯出的函式名稱不被改變
#endif
// 生成的loadModel函式，可以看到JNI的修飾和jlong返回值，函式名稱格式為 Java_包名_類名_函式名
// 函式的前兩個引數是 JNIEnv 表示當前執行緒的 JVM 環境引數，jclass 表示呼叫的 class 物件，可以透過這兩個引數去操作 Java 物件。
JNIEXPORT jlong JNICALL Java_com_vivo_demo_model_ModelComputer_loadModel
  (JNIEnv *, jclass, jstring);
 
JNIEXPORT void JNICALL Java_com_vivo_demo_model_ModelComputer_close
  (JNIEnv *, jclass, jlong);
 
JNIEXPORT jfloat JNICALL Java_com_vivo_demo_model_ModelComputer_compute
  (JNIEnv *, jclass, jlong, jlongArray);
 
#ifdef __cplusplus
}
#endif
#endif

  • 踩坑2：Javah 執行失敗

如果生成失敗，可以參考上面 JNI 格式的 “.h” 檔案手寫一個出來，只要格式無誤，效果是一樣的。其中 jni.h 是 JDK 路徑下的一個檔案，裡面定義了一些 JNI 的型別，返回值， 異常， JavaVM 結構體以及一些方法（型別轉化，欄位獲取，JVM 資訊獲取等）。jni.h 還依賴了一個 jni_md.h 檔案，其中定義了 jbyte，jint 和 jlong，這三個型別在不同的機器下的定義是有差異的。

我們可以看下 JNI 常用資料型別與 Java 的對應關係：

圖片

圖3 JNI常用資料型別

如圖3，JNI 定義了一些基本資料型別和引用資料型別，可以完成 Java 和 C++ 的資料轉化。JNIEnv 是一個指向本地執行緒資料的介面指標，通俗的來講，我們透過 JNIEnv 中的方法，可以完成 Java 和 C++ 的資料轉化，透過它，可以使 C++ 訪問 Java 的堆記憶體。

對於基本的資料型別，透過值傳遞，可以進行強制轉化，可以理解為只是定義的名稱發生改變，和 java 基本資料型別差異不大。

而引用資料型別，JNI 定義了 Object 型別的引用，那麼就意味著，java 可以透過引用傳遞任意物件到 C++ 中。對於像基礎型別的陣列和 string 型別，如果透過引用傳遞，那麼 C++ 就要訪問 Java 的堆記憶體，透過 JNIEnv 中的方法來訪問 Java 物件，雖然不需要我們關心具體邏輯，但是其效能消耗要高於 C++ 指標操作物件的。所以 JNI 將陣列和 string 複製到本地記憶體（緩衝區）中，這樣不但提高了訪問速度，還減輕了 GC 的壓力，缺點就是需要使用 JNI 提供的方法進行建立和釋放。

// 可以使用下列三組函式，其中 tpye 為基本資料型別,後兩組有 Get 和 Release 方法，Release 方法的作用是提醒 JVM 釋放記憶體
// 資料量小的時候使用此方法，原理是將資料複製到C緩衝區，分配在 C 堆疊上，因此只適用於少量的元素，Set 操作是對快取區進行修改
Get<type>ArrayRegion
Set<type>ArrayRegion
// 將陣列的內容複製到本地記憶體中，供 C++ 使用
Get<type>ArrayElement
Release<type>ArrayElement
// 有可能直接返回 JVM 中的指標，否則的話也會複製一個陣列出來，和 GetArrayElement 功能相同
GetPrimitiveArrayCritical
ReleasePrimitiveArrayCritical

透過這三組方法的介紹，也就大致瞭解了 JNI 的資料型別轉化，如果沒有 C++ 建立修改 Java Object 的操作的話，那編寫 C++ 程式碼和正常的 C++ 開發無異，下面給出了 “export_jni.h” 程式碼示例。

#include "jni.h" // 這裡改為相對引用，是因為把 jni.h 和 jni_md.h 複製到專案中，方便編譯
#include "computer.cc"
#ifndef _Included_com_vivo_demo_model_ModelComputer
#define _Included_com_vivo_demo_model_ModelComputer
#ifdef __cplusplus
extern "C" {
#endif
    JNIEXPORT jlong JNICALL Java_com_vivo_demo_model_ModelComputer_loadModel
    (JNIEnv* env, jclass clazz, jstring path) {
        vivo::Computer* ptr = new vivo::Computer();
        const char* cpath = env->GetStringUTFChars(path, 0); // 將 String 轉為 char*
        ptr->init_model(cpath);
        env->ReleaseStringUTFChars(path, cpath); // 釋放String
        return (long)ptr;
    };
 
    JNIEXPORT void JNICALL Java_com_vivo_demo_model_ModelComputer_close
    (JNIEnv* env, jclass clazz, jlong ptr) {
        vivo::Computer* computer = (vivo::Computer*)ptr; // 獲取到物件
        delete computer; // 刪除物件
    };
 
    JNIEXPORT jfloat JNICALL Java_com_vivo_demo_model_ModelComputer_compute
    (JNIEnv* env, jclass clazz, jlong ptr, jlongArray array) {
        jlong* idx_ptr = env->GetLongArrayElements(array, NULL); // 將 array 轉為 jlong*
        vivo::Computer* computer = (vivo::Computer*)ptr; // 獲取到 C++ 物件
        float result = computer->compute((long *)idx_ptr); // 執行 C++ 方法
        env->ReleaseLongArrayElements(array, idx_ptr, 0); // 釋放 array
        return result; // 返回結果
    };
 
#ifdef __cplusplus
}
#endif
#endif

C++ 程式碼編譯完成後，把 lib 檔案放到 resource 目錄指定位置，如果為了方便，可以寫個 shell 指令碼一鍵執行。

  • 踩坑3：伺服器啟動時報java.lang.UnsatisfiedLinkError 異常

又是這個異常，前文已經介紹了一種解決方案，但在實際應用中仍然頻繁出現，比如：

  1. 執行環境有問題（比如在 linux 下編譯在 windows 上執行，這是不可以的）；

  2. JVM 位數和 lib 的位數不一致 （比如一個是 32 位，一個是 64 位）；

  3. C++ 函式名寫錯；

  4. 生成的 lib 檔案中並沒有相對應的方法。

對於這些問題，只要認真分析異常日誌，便可以逐一解決，也有工具可以協助我們解決問題。

使用 dumpbin/objdump 分析 lib，更快速地解決 UnsatisfiedLinkError。

對於 lib 庫中的函式檢查，不同作業系統也提供了不同的工具。

在 windows 下，可以使用 dumpbin 工具或者 Dependency Walker 工具分析 lib 中是否存在所編寫的 C++ 方法。dumpbin 指令如下：

dumpbin /EXPORTS xxx.dll

圖片

圖4 dumpbin 檢視 dll 檔案

而 Dependency Walker 只需要開啟 dll 檔案就可以看到相關資訊了。

圖片

圖5 Dependency Walker 檢視 dll 檔案

在 Linux 下，可以使用 objdump 工具分析 so 檔案中的資訊。

objdump 指令如下：

objdump -t xxx.so

圖片

圖6 objdump 檢視 so 檔案

3.4 效能分析

根據之前的調研，我們注意到 Java 對 native 方法的呼叫本身也存在額外效能開銷，針對此我們用 JMH 進行了簡單測試。圖 7 展示的是 JNI 空方法呼叫和 Java 的對比：

圖片

圖7 - 空函式呼叫對比 （資料來源自個人機器JMH測試，僅供參考）

其中 JmhTest.code 為呼叫 native 空方法， JmhTest.jcode 為呼叫 java 空方法，從中可以看出，直接呼叫 java 的方法要比呼叫 native 方法快十倍還要多。我們對堆疊呼叫進行了簡單分析，發現呼叫 native 的過程比直接呼叫 java 方法要繁瑣一些，進入了 ClassLoad 的  findNative 方法。

// Invoked in the VM class linking code.
// loader 為類載入器， name 為C++方法的 name，eg: Java_com_vivo_demo_model_ModelComputer_compute
static long findNative(ClassLoader loader, String name) {
    // 選擇 nativeLibary   
    Vector<NativeLibrary> libs =
        loader != null ? loader.nativeLibraries : systemNativeLibraries;
    synchronized (libs) {
        int size = libs.size();
        for (int i = 0; i < size; i++) {
            NativeLibrary lib = libs.elementAt(i);
            // 找到 name 持有的 handel
            long entry = lib.find(name); 
            if (entry != 0)
                // 返回 handel
                return entry;
        }
    }
    return 0;
}

堆疊資訊如下：

圖片

圖8 呼叫 native 堆疊資訊

find 方法是一個 native 方法，堆疊上也列印不出相關資訊，但不難得出，透過 find 方法去呼叫 lib 庫中的方法，還要再經過至少一輪的對映才能找到對應的 C++ 函式執行，然後將結果返回。瞬間回想起圖一，這種呼叫鏈路，透過 Native Interface 來串起本地方法棧，虛擬機器棧，nativeLibrary 和執行引擎之間的關係，邏輯勢必會複雜一些，相對的呼叫耗時也會增加。

做了這麼多工作，差點忘了我們的目標：提高我們的計算和載入速度。經過上文的最佳化後，我們在壓測環境進行了全鏈路壓測，發現即使 native 的呼叫存在額外開銷，全鏈路的效能仍然有了較為明顯的提升。

我們的服務在模型推理的核心計算上耗時降低了 80%，載入和解析模型檔案耗時也降低了 60%（分鐘級到秒級），GC 的平均耗時也降低了 30%，整體的收益非常明顯。

圖片

圖9 young GC 耗時對比

四、思考和總結：JNI 帶來的收益

JNI 在一些特定場景下的成功應用開啟了我們的最佳化思路，尤其是在 Java 上進行了較多最佳化嘗試後並沒有進展時，JNI 確實值得一試。

又回到了最初的問題：JNI 真的好用嗎？我的答案是：它並不是很好用。如果是一名很少接觸 C++ 程式設計的工程師，那麼在第一步的環境搭建和編譯上，就要耗費大量的時間，再到後續的程式碼維護，C++ 調優等等，是一個非常頭疼的事情。但我還是非常推薦去了解這項技術和這項技術的應用，去思考這項技術能夠給自己的伺服器效能帶來提升。

或許有一天，JNI 能為你所用！

參考資料：

  1. Oracle JNI Guide: Java Native Interface

  2. bazel 概述

  3. docker hub

  4. JMH GitHub

  5. Dumpbin refrence

來自 “ ITPUB部落格 ” ，連結：http://blog.itpub.net/69912579/viewspace-2910714/，如需轉載，請註明出處，否則將追究法律責任。

相關文章