密集计算场景下的 JNI 实战

1,326 阅读16分钟

作者: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