JDK 22 外部函数与内存 API (FFM) 实战

JNI 那套胶水代码，我写了三年，写烦了

我们有个自研的向量检索库 libhnswwrap.so，C++ 写的，Java 侧通过 JNI 调用。三年来每次 C++ 侧改接口，我都要同步改三处：C 的胶水层、Java 的 native 方法声明、还有两边的结构体序列化代码。

上次升级的改动是增加了一个 float* scores 出参。我改了 JNI 胶水层，忘了改 Java 侧声明的数组长度常量，上线后跑了一周才在某个特定查询下触发堆外内存越界，JVM 直接 core dump，hs_err 文件里只有一行：

# Problematic frame:
# C  [libhnswwrap.so+0x3a21]  hnsw_search_batch+0x111

没有 Java 堆栈，没有业务上下文。查了两天。这事儿之后我下定决心把 FFM 搞明白。

FFM 是什么，JDK 22 里到什么程度了

Foreign Function & Memory API 由 JEP 454 在 JDK 22 正式转正（2024 年 3 月），之前经历了十几轮预览。它替代的是两样东西：JNI（调用本地函数）和 sun.misc.Unsafe / ByteBuffer（操作堆外内存）。

核心类型就四个：

• MemorySegment：一段内存（堆内或堆外）的视图，带边界检查和生命周期
• Arena：内存生命周期的作用域，关掉之后整段内存一起释放
• Linker：把 Java 方法调用桥接到本地函数
• SymbolLookup：在本地库里找符号

我们的目标函数签名：

// hnsw_wrap.h
typedef struct {
    int64_t  index_id;
    uint32_t dim;
    void*    data;
} IndexHandle;

int  hnsw_load(const char* path, int dim, IndexHandle* out);
int  hnsw_search(IndexHandle* h, const float* query,
                 int topk, int64_t* out_ids, float* out_scores);
void hnsw_free(IndexHandle* h);

手写一个绑定

先找到库里的符号，拿到方法句柄：

public final class HnswNative {
    private static final Arena LIB = Arena.global();
    private static final SymbolLookup LOOKUP;
    private static final Linker LINKER = Linker.nativeLinker();

    static {
        System.loadLibrary("hnswwrap");
        LOOKUP = SymbolLookup.loaderLookup();
    }

    // 缓存 MethodHandle，查找一次复用，千万别每次调用都查
    private static final MethodHandle HNSW_SEARCH = LINKER.downcallHandle(
            LOOKUP.find("hnsw_search").orElseThrow(),
            FunctionDescriptor.of(
                ValueLayout.JAVA_INT,                       // 返回 int
                ValueLayout.ADDRESS,                        // IndexHandle*
                ValueLayout.ADDRESS,                        // const float*
                ValueLayout.JAVA_INT,                       // topk
                ValueLayout.ADDRESS,                        // int64_t*
                ValueLayout.ADDRESS),                       // float*
            Linker.Option.critical(false));

    private static final MethodHandle HNSW_LOAD = LINKER.downcallHandle(
            LOOKUP.find("hnsw_load").orElseThrow(),
            FunctionDescriptor.of(ValueLayout.JAVA_INT,
                ValueLayout.ADDRESS, ValueLayout.JAVA_INT, ValueLayout.ADDRESS));
}

然后封装成普通的 Java 方法：

public SearchResult search(float[] query, int topK) {
    // try-with-resources 保证这段堆外内存在方法结束时一定释放
    try (Arena arena = Arena.ofConfined()) {
        MemorySegment querySeg = arena.allocate(ValueLayout.JAVA_FLOAT, query.length);
        querySeg.copyFrom(MemorySegment.ofArray(query));

        MemorySegment ids   = arena.allocate(ValueLayout.JAVA_LONG, topK);
        MemorySegment score = arena.allocate(ValueLayout.JAVA_FLOAT, topK);

        int rc = (int) HNSW_SEARCH.invokeExact(
                (MemorySegment) handle,
                querySeg, topK, ids, score);

        if (rc != 0) throw new IllegalStateException("hnsw_search failed, rc=" + rc);

        long[] outIds = ids.toArray(ValueLayout.JAVA_LONG);
        float[] outScores = score.toArray(ValueLayout.JAVA_FLOAT);
        return new SearchResult(outIds, outScores);
    } catch (Throwable t) {
        throw new RuntimeException(t);
    }
}

对比一下原来的 JNI 胶水层，那段 C 代码有 130 行，现在全部删掉了。

Arena 的四种，选错就是内存泄漏或使用崩溃

这是 FFM 里最需要想清楚的概念。Arena 决定了它分配出来的 MemorySegment 什么时候失效。

几个必须记住的点：

一、confined arena 里分配的内存，别的线程访问会直接抛异常。 这个检查是运行时的，不是编译期的：

java.lang.WrongThreadException: Attempted access outside of the thread
    owning the confined arena

我们第一版把 querySeg 缓存起来复用，然后在虚拟线程里访问，就撞上了。虚拟线程是换线程跑的，一个虚拟线程的不同执行段可能在不同载体线程上，所以缓存 MemorySegment 跨虚拟线程复用，必须用 ofShared()。

二、Arena 关闭后访问内存，抛的是 IllegalStateException 而不是段错误。 这是 FFM 相对 JNI 最大的安全提升：

Arena a = Arena.ofConfined();
MemorySegment s = a.allocate(1024);
a.close();
s.get(ValueLayout.JAVA_BYTE, 0);
// java.lang.IllegalStateException: Already closed

JVM 不会崩，有完整的 Java 堆栈。光这一点就值回票价了。

三、ofAuto() 看着方便，但释放时机不可控。我们的索引句柄（2.4 GB 的堆外内存）一开始用了 auto，结果 Full GC 频繁时才释放，堆外内存水位长期在 90%。改成 global() 加显式释放之后才正常。大内存必须显式管理生命周期，别指望 GC。

结构体的内存布局

本地函数返回结构体时，布局要手动对齐。用 MemoryLayout.structLayout：

static final GroupLayout INDEX_HANDLE = MemoryLayout.structLayout(
        ValueLayout.JAVA_LONG.withName("index_id"),
        ValueLayout.JAVA_INT.withName("dim"),
        MemoryLayout.paddingLayout(4),                 // 4 字节填充，对齐到 8
        ValueLayout.ADDRESS.withName("data")
);

// 用 VarHandle 访问字段，比算偏移量安全
static final VarHandle VH_DIM = INDEX_HANDLE.varHandle(
        MemoryLayout.PathElement.groupElement("dim"));

MemorySegment handle = arena.allocate(INDEX_HANDLE);
VH_DIM.set(handle, 768L);
int dim = (int) VH_DIM.get(handle);

那个 paddingLayout(4) 是我踩过坑的地方。C 的结构体在 64 位系统上，int64 后面跟 uint32，再跟指针，编译器会自动插入 4 字节填充让指针 8 字节对齐。第一次我忘了写，结果 data 字段读到的是垃圾地址，程序随机崩溃——这种 bug 在 JNI 时代就是段错误，在 FFM 里至少能拿到异常。

验证布局对不对，可以用 jextract 生成的头文件解析结果对照，或者直接打印：

System.out.println(INDEX_HANDLE);
// [[i8](struct{}index_id)[i4](struct{}dim)]... byteSize=16, byteAlignment=8
//                                                    ↑ 确认是 16 而不是 12

性能：比 JNI 快，但没有宣传的那么夸张

测试场景：单次查询调 1 万次 hnsw_search（模拟批量检索），JMH 跑 5 轮取平均，机器 16C32G，JDK 22.0.2。

几个结论：

• MethodHandle 必须缓存成 static final。downcallHandle 每次调用都会做符号查找和调用点链接，开销是 3.2 微秒，比 JNA 还慢。这一条如果没做对，FFM 会比 JNI 慢 50 倍。
• FFM 比 JNI 快 30% 左右，主要省在参数编组上。JNI 要走 GetPrimitiveArrayCritical 之类的 JNIEnv 调用，FFM 直接把 MemorySegment 地址传过去。
• Linker.Option.critical(true) 能再快 14%，但它要求被调函数不能调用任何 JNI 函数、不能长时间阻塞。我们的 hnsw_search 是纯计算，可以用；hnsw_load 要读文件，不能用。

端到端看，我们的检索接口 P99 从 41 毫秒降到 38 毫秒。这点提升说实话不是迁移的主要理由，真正值钱的是删掉了 130 行 C 代码和 480 行 Java native 声明，以及内存越界不再导致进程崩溃。

踩过的坑

坑一：忘了加 --enable-native-access

java.lang.UnsupportedOperationException: Native access is restricted

JDK 22 里受限方法调用需要显式授权（JEP 472）。按模块授权而不是无脑开 ALL-UNNAMED：

java --enable-native-access=com.example.search -jar app.jar
# 或者运行时用 Native 模块的方式

坑二：字符串转换的编码

C 侧要的是 const char*，Java 的 String 是 UTF-16。FFM 提供了 allocateFrom，但它按平台默认编码转换，如果路径里有中文就会炸：

// 有隐患：依赖 file.encoding
MemorySegment path = arena.allocateFrom("/data/index/model.bin");

// 明确指定编码
byte[] bytes = pathStr.getBytes(StandardCharsets.UTF_8);
MemorySegment path = arena.allocate(ValueLayout.JAVA_BYTE, bytes.length + 1);
path.copyFrom(MemorySegment.ofArray(bytes));
path.set(ValueLayout.JAVA_BYTE, bytes.length, (byte) 0);   // 别忘了 \0

那个 +1 和末尾的零字节也是坑，第一次忘了写，C 侧 strlen 读到了越界的脏数据。

坑三：jextract 生成的代码不好定制

jextract 是 OpenJDK 提供的工具，能从头文件直接生成 Java 绑定，听起来很美：

$ jextract --output src/main/java -t com.example.hnsw hnsw_wrap.h

实际用下来两个问题：一是生成的代码量巨大（我们这个头文件生成了 2400 行），二是它按 C 的原始结构生成，字段名和类型都很原始，业务层还是得再包一层。我们最后的选择是手写核心的 6 个函数绑定，其余用 jextract 生成一次后手工裁剪。

跨线程共享索引句柄

索引句柄（2.4 GB 的堆外数据）是所有查询线程共享的。这里必须用 ofShared()，而且关闭时机要管好。

public final class HnswIndex implements AutoCloseable {
    private final Arena arena = Arena.ofShared();      // 多线程访问，不能是 confined
    private final MemorySegment handle;

    public HnswIndex(String path, int dim) {
        this.handle = arena.allocate(INDEX_HANDLE);
        MemorySegment pathSeg = arena.allocateFrom(path, StandardCharsets.UTF_8);
        int rc = (int) HNSW_LOAD.invokeExact(pathSeg, dim, handle);
        if (rc != 0) {
            arena.close();                              // 加载失败要立刻释放，别泄漏
            throw new IllegalStateException("hnsw_load failed, rc=" + rc);
        }
    }

    public void close() {
        HNSW_FREE.invoke(handle);
        arena.close();                                  // 关闭后所有 segment 全部失效
    }
}

这里有个容易漏的点：构造失败时要手动 close() arena。如果只在 close() 方法里释放，加载失败的路径会泄漏掉已分配的内存（虽然这里只有 16 字节，但如果中间分配了大缓冲就不是小事）。

查询线程这边，每次查询各自开一个 confined arena 放临时缓冲，这些 arena 之间互不干扰：

// 每个查询线程一个临时 arena，方法结束自动释放
try (Arena local = Arena.ofConfined()) {
    MemorySegment querySeg = local.allocate(JAVA_FLOAT, dim);
    MemorySegment ids = local.allocate(JAVA_LONG, topK);
    MemorySegment scores = local.allocate(JAVA_FLOAT, topK);
    // handle 来自 shared arena，不同 arena 的 segment 可以互相传参
    int rc = (int) HNSW_SEARCH.invokeExact(handle, querySeg, topK, ids, scores);
    ...
}

一个真实的踩坑：arena 关闭导致的偶发异常

上线后出现过几次 IllegalStateException: Already closed，频率大概每天两三次。查了很久才发现是优雅停机的问题：应用收到 SIGTERM 后，Spring 先销毁了 HnswIndex bean（关闭了 shared arena），但还有几个在飞的查询线程拿着 handle 去调 hnsw_search。

修复方式是加关闭门闩，等所有在飞的查询结束再关：

public void close() {
    shuttingDown = true;
    // 等所有在飞的查询结束，最多等 10 秒
    long deadline = System.currentTimeMillis() + 10_000;
    while (inflight.get() > 0 && System.currentTimeMillis() < deadline) {
        Thread.sleep(50);
    }
    HNSW_FREE.invoke(handle);
    arena.close();
}

这个 bug 在 JNI 版本里的表现是段错误，进程直接死。在 FFM 版本里是抛异常，查询失败但进程活着，还有完整堆栈。光是这一点，迁移就值了。

什么情况下别用 FFM

• 要调用 C++ 的类和方法。FFM 只支持 C ABI，C++ 的 name mangling、this 指针、异常都没法直接用。我们之所以有个 libhnswwrap 而不是直接调 hnswlib，就是因为要写一层 C 接口包装。
• 回调（upcall）频繁的场景。用 Linker.upcallStub 可以把 Java 方法传给 C，但每次调用的开销比 downcall 大一个量级。我们的进度回调用了一次性 stub 复用，没在热路径上。
• 还在用 JDK 17 或更低。FFM 在 21 里是预览，22 才转正。老项目不值得为了它单独升 JDK。
• 调用频率极低。我们有个每天跑一次的离线任务原来也用 JNI，迁移到 FFM 花了半天，收益为零，纯属给自己找事。

小结

迁移前后的对比：

性能提升是附带的，真正的收益在工程性：不用写 C 代码了，内存错误可诊断了，接口变更成本降了一个数量级。

如果你的项目里也有 JNI 胶水层，我建议先挑一个最简单的函数试着手写绑定，感受一下 Arena 的生命周期模型。这个心智模型建起来之后，剩下的迁移就是体力活。

参考

• 官方文档：java-core
• Baeldung 上的相关实战
• StackOverflow 上的相关讨论